Make the router's routing state authoritative instead of sniffed

lib/cli/ui/stdout_router.rb

@@ -11,11 +11,36 @@
 module CLI
   module UI
     module StdoutRouter
+      WRITE_WITHOUT_CLI_UI = :write_without_cli_ui
+
+      # Defined here rather than on the singleton class, where callers had no
+      # way to name it.
+      NotEnabled = Class.new(StandardError)
+
+      class Route
+        #: Method
+        attr_reader :original_write
+
+        #: bool
+        attr_reader :remove_compatibility_write
+
+        #: (Method original_write, bool remove_compatibility_write) -> void
+        def initialize(original_write, remove_compatibility_write)
+          @original_write = original_write
+          @remove_compatibility_write = remove_compatibility_write
+        end
+      end
+      private_constant :Route
+
       class Writer
-        #: (io_like stream, Symbol name) -> void
-        def initialize(stream, name)
+        # `original_write` must be the stream's pre-routing `write`; `activate`
+        # passes it in. Holding the method itself keeps an in-flight write
+        # working after `deactivate` removes the WRITE_WITHOUT_CLI_UI alias.
+        #: (io_like stream, Symbol name, Method original_write) -> void
+        def initialize(stream, name, original_write)
           @stream = stream
           @name = name
+          @original_write = original_write
         end
 
         #: (*Object args) -> Integer
@@ -38,7 +63,7 @@ def write(*args)
           end
 
           stream_args = prepend_id(@stream, strs) #: as untyped
-          ret = @stream.write_without_cli_ui(*stream_args) #: as Integer
+          ret = @original_write.call(*stream_args) #: as Integer
           if (dup = StdoutRouter.duplicate_output_to)
             begin
               dup_args = prepend_id(dup, strs) #: as untyped
@@ -307,10 +332,6 @@ def synchronize(&block)
       end
 
       class << self
-        WRITE_WITHOUT_CLI_UI = :write_without_cli_ui
-
-        NotEnabled = Class.new(StandardError)
-
         #: io_like?
         attr_accessor :duplicate_output_to
 
@@ -339,62 +360,152 @@ def assert_enabled!
           raise NotEnabled unless enabled?
         end
 
+        # Only unroutes what it routed, so nesting inside an already-enabled
+        # router leaves that router alone.
         #: [T] { -> T } -> T
         def with_enabled(&block)
-          enable
+          activated = activate_current_streams
           yield
         ensure
-          disable
+          activated&.each { |stream| deactivate(stream) if enabled?(stream) }
         end
 
         # TODO: remove this
         #: -> void
         def ensure_activated
-          enable unless enabled?
+          enable
         end
 
+        # Makes the current $stdout/$stderr the only routed streams; returns
+        # whether anything changed. Per-stream so a half-routed process can
+        # recover.
         #: -> bool
         def enable
-          return false if enabled?($stdout) || enabled?($stderr)
-
-          activate($stdout, :stdout)
-          activate($stderr, :stderr)
-          true
+          streams = current_streams
+          activated = activate_streams(streams)
+          deactivated = deactivate_streams_except(streams)
+          deactivated || !activated.empty?
         end
 
         #: (?io_like stream) -> bool
         def enabled?(stream = $stdout)
-          stream.respond_to?(WRITE_WITHOUT_CLI_UI)
+          routed_streams.key?(stream)
         end
 
         #: -> bool
         def disable
-          return false unless enabled?($stdout) && enabled?($stderr)
+          routed = routed_streams.keys
+          routed.each { |stream| deactivate(stream) }
+          !routed.empty?
+        end
+
+        # Writes past the router: no frame inset, capture hooks, or output
+        # duplication. Prefer this to calling WRITE_WITHOUT_CLI_UI directly,
+        # which only exists while the stream is routed.
+        #: (io_like stream, *Object args) -> Integer
+        def write_without_routing(stream, *args)
+          original_write = routed_streams[stream]&.original_write
+          write_args = args #: as untyped
+          return stream.write(*write_args) unless original_write
 
-          deactivate($stdout)
-          deactivate($stderr)
-          true
+          original_write.call(*write_args) #: as Integer
         end
 
         private
 
+        # Routed stream => the state needed to restore it. State lives here
+        # rather than in a method on the stream, so nothing on the stream can
+        # make `enabled?` lie. `deactivate` drops the entry.
+        #: -> Hash[io_like, Route]
+        def routed_streams
+          @routed_streams ||= {}.compare_by_identity
+        end
+
+        #: -> Hash[io_like, Symbol]
+        def current_streams
+          streams = {}.compare_by_identity #: Hash[io_like, Symbol]
+          streams[$stdout] = :stdout
+          streams[$stderr] ||= :stderr
+          streams
+        end
+
+        #: (Hash[io_like, Symbol] streams) -> bool
+        def deactivate_streams_except(streams)
+          stale = routed_streams.keys.reject { |stream| streams.key?(stream) }
+          stale.each { |stream| deactivate(stream) }
+          !stale.empty?
+        end
+
+        #: (Hash[io_like, Symbol] streams) -> Array[io_like]
+        def activate_streams(streams)
+          activated = [] #: Array[io_like]
+          begin
+            streams.each do |stream, streamname|
+              next if enabled?(stream)
+
+              activate(stream, streamname)
+              activated << stream
+            end
+            activated
+          rescue
+            activated.reverse_each { |stream| deactivate(stream) if enabled?(stream) }
+            raise
+          end
+        end
+
+        #: -> Array[io_like]
+        def activate_current_streams
+          activate_streams(current_streams)
+        end
+
         #: (io_like stream) -> void
         def deactivate(stream)
+          route = routed_streams.delete(stream)
+          original_write = route&.original_write
           sc = stream.singleton_class
+
           sc.send(:remove_method, :write)
-          sc.send(:alias_method, :write, WRITE_WITHOUT_CLI_UI)
+          # Restore only a `write` we displaced; otherwise leave the class
+          # method exposed again. Prefer the compatibility alias when it is
+          # ours; otherwise restore the stored method without disturbing a
+          # foreign method using that name.
+          if original_write&.owner == sc
+            if route&.remove_compatibility_write
+              sc.send(:alias_method, :write, WRITE_WITHOUT_CLI_UI)
+            else
+              sc.send(:define_method, :write, original_write)
+            end
+          end
+          sc.send(:remove_method, WRITE_WITHOUT_CLI_UI) if route&.remove_compatibility_write
         end
 
         #: (io_like stream, Symbol streamname) -> void
         def activate(stream, streamname)
-          writer = StdoutRouter::Writer.new(stream, streamname)
-
-          raise if stream.respond_to?(WRITE_WITHOUT_CLI_UI)
-
-          stream.singleton_class.send(:alias_method, WRITE_WITHOUT_CLI_UI, :write)
+          original_write = stream.method(:write)
+          writer = StdoutRouter::Writer.new(stream, streamname, original_write)
+          installed_compatibility_write = false
+          compatibility_write = if stream.respond_to?(WRITE_WITHOUT_CLI_UI, true)
+            stream.method(WRITE_WITHOUT_CLI_UI)
+          end
+          remove_compatibility_write = compatibility_write&.owner == stream.singleton_class &&
+            compatibility_write == original_write
+
+          # Kept for callers that bypass the router through the alias itself,
+          # without overwriting an unrelated method using the same name.
+          unless compatibility_write
+            stream.singleton_class.send(:alias_method, WRITE_WITHOUT_CLI_UI, :write)
+            installed_compatibility_write = true
+            remove_compatibility_write = true
+          end
           stream.define_singleton_method(:write) do |*args|
             writer.write(*args)
           end
+          routed_streams[stream] = Route.new(original_write, !!remove_compatibility_write)
+        rescue
+          if installed_compatibility_write
+            stream.singleton_class.send(:remove_method, WRITE_WITHOUT_CLI_UI)
+          end
+          raise
         end
       end
     end