Make the router's routing state authoritative instead of sniffed

Disabling the router restored the original write method but left the
write_without_cli_ui marker method installed on the stream. Since
enabled? checked for that marker, the router reported itself enabled
while routing was actually inactive, and a later enable refused to do
anything. The process was stuck: output permanently bypassed routing,
and Capture#run's assert_enabled! still passed while captures silently
collected nothing.

with_enabled compounded this by always disabling on exit, so entering
that scope in a process that had already enabled the router (the common
production state) permanently broke routing.

Routing state now lives in a private stream => original-write registry
rather than being inferred from a marker method on the stream, so
nothing on the stream can make enabled? lie:

* enable/disable act per stream instead of all-or-nothing, so a process
  that ends up half-routed (e.g. $stdout reassigned after enable, or
  $stderr aliased to $stdout) can recover. with_enabled unroutes
  exactly the streams it routed, leaving an already-enabled router
  alone.
* Writer holds the stream's pre-routing write method, so it no longer
  dispatches through write_without_cli_ui. A write already in flight on
  another thread survives a concurrent disable.
* deactivate restores only a write it displaced. Previously every
  disable left a singleton copy of the class's write behind, shadowing
  the class method for the life of the stream.
* write_without_routing is the supported way to bypass the router.
  Callers currently reach for the write_without_cli_ui alias directly
  (dev's cd/edit commands do, unguarded); that alias only exists while
  the stream is routed, and now really does disappear on disable.
* NotEnabled moves out of class << self, where it was only reachable as
  StdoutRouter.singleton_class::NotEnabled - i.e. nobody could rescue
  the error Capture#run raises.
* ensure_activated is now just enable, which is idempotent.

Co-authored-by: River <river@shopify.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: 3 people

lib/cli/ui/stdout_router.rb

@@ -11,11 +11,19 @@
 module CLI
   module UI
     module StdoutRouter
+      # Defined here rather than on the singleton class, where callers had no
+      # way to name it.
+      NotEnabled = Class.new(StandardError)
+
       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.method(:write))
           @stream = stream
           @name = name
+          @original_write = original_write
         end
 
         #: (*Object args) -> Integer
@@ -38,7 +46,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
@@ -309,8 +317,6 @@ def synchronize(&block)
       class << self
         WRITE_WITHOUT_CLI_UI = :write_without_cli_ui
 
-        NotEnabled = Class.new(StandardError)
-
         #: io_like?
         attr_accessor :duplicate_output_to
 
@@ -339,62 +345,100 @@ 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
 
+        # Routes whichever of $stdout/$stderr isn't routed yet; 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
+          !activate_current_streams.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 = [$stdout, $stderr].uniq.select { |stream| enabled?(stream) }
+          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]
+          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 `write` it had before routing. 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, Method]
+        def routed_streams
+          @routed_streams ||= {}.compare_by_identity
+        end
+
+        #: -> Array[io_like]
+        def activate_current_streams
+          activated = [] #: Array[io_like]
+          { stdout: $stdout, stderr: $stderr }.each do |streamname, stream|
+            next if enabled?(stream)
+
+            activate(stream, streamname)
+            activated << stream
+          end
+          activated
+        end
+
         #: (io_like stream) -> void
         def deactivate(stream)
+          original_write = routed_streams.delete(stream)
           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. The alias binds the method, not the name, so
+          # it survives removing WRITE_WITHOUT_CLI_UI.
+          sc.send(:alias_method, :write, WRITE_WITHOUT_CLI_UI) if original_write&.owner == sc
+          sc.send(:remove_method, WRITE_WITHOUT_CLI_UI)
         end
 
         #: (io_like stream, Symbol streamname) -> void
         def activate(stream, streamname)
-          writer = StdoutRouter::Writer.new(stream, streamname)
+          raise("#{streamname} is already routed") if stream.respond_to?(WRITE_WITHOUT_CLI_UI)
 
-          raise if stream.respond_to?(WRITE_WITHOUT_CLI_UI)
+          original_write = stream.method(:write)
+          writer = StdoutRouter::Writer.new(stream, streamname, original_write)
 
+          # Kept for callers that bypass the router through the alias itself.
           stream.singleton_class.send(:alias_method, WRITE_WITHOUT_CLI_UI, :write)
           stream.define_singleton_method(:write) do |*args|
             writer.write(*args)
           end
+          routed_streams[stream] = original_write
         end
       end
     end

test/cli/ui/stdout_router_test.rb

@@ -24,6 +24,145 @@ def test_current_id
         end
       end
 
+      def test_router_can_reenable_after_disable
+        # capture_io swaps in fresh StringIO streams, so enable/disable here
+        # can't interfere with a router installed on the real stdio streams.
+        capture_io do
+          assert(StdoutRouter.enable)
+          assert($stdout.respond_to?(:write_without_cli_ui))
+          assert(StdoutRouter.disable)
+
+          refute_predicate(StdoutRouter, :enabled?)
+          refute($stdout.respond_to?(:write_without_cli_ui))
+
+          assert(StdoutRouter.enable)
+          assert_predicate(StdoutRouter, :enabled?)
+
+          cap = StdoutRouter::Capture.new { print('hi') }
+          cap.run
+          assert_equal('hi', cap.stdout)
+        ensure
+          StdoutRouter.disable
+        end
+      end
+
+      def test_capture_raises_when_the_router_is_disabled
+        capture_io do
+          assert_raises(StdoutRouter::NotEnabled) do
+            StdoutRouter::Capture.new {}.run
+          end
+        end
+      end
+
+      def test_disable_leaves_the_streams_as_it_found_them
+        capture_io do
+          StdoutRouter.enable
+          StdoutRouter.disable
+
+          assert_empty($stdout.singleton_methods)
+          assert_empty($stderr.singleton_methods)
+        end
+      end
+
+      def test_disable_restores_a_pre_existing_singleton_write
+        capture_io do
+          written = []
+          $stdout.define_singleton_method(:write) do |*args|
+            written << args.join
+            super(*args)
+          end
+
+          StdoutRouter.enable
+          $stdout.write('routed')
+          StdoutRouter.disable
+          $stdout.write('plain')
+
+          assert_equal(['routed', 'plain'], written)
+          assert_equal([:write], $stdout.singleton_methods)
+        end
+      end
+
+      def test_enable_recovers_a_half_routed_process
+        capture_io do
+          StdoutRouter.enable
+          $stderr = StringIO.new
+
+          assert(StdoutRouter.enable)
+          assert(StdoutRouter.enabled?($stderr))
+          assert(StdoutRouter.disable)
+          refute_predicate(StdoutRouter, :enabled?)
+        ensure
+          StdoutRouter.disable
+        end
+      end
+
+      def test_enable_handles_stderr_aliased_to_stdout
+        capture_io do
+          $stderr = $stdout
+
+          assert(StdoutRouter.enable)
+          assert(StdoutRouter.disable)
+          assert_empty($stdout.singleton_methods)
+        end
+      end
+
+      def test_in_flight_writes_survive_disable
+        capture_io do
+          StdoutRouter.enable
+          in_flight = $stdout.method(:write)
+          StdoutRouter.disable
+
+          in_flight.call('late')
+          assert_includes($stdout.string, 'late')
+        end
+      end
+
+      def test_write_without_routing_bypasses_the_router
+        capture_io do
+          StdoutRouter.enable
+          cap = StdoutRouter::Capture.new { StdoutRouter.write_without_routing($stdout, 'bypass') }
+          cap.run
+
+          assert_equal('', cap.stdout)
+          assert_includes($stdout.string, 'bypass')
+        ensure
+          StdoutRouter.disable
+        end
+      end
+
+      def test_write_without_routing_falls_back_to_write_when_disabled
+        capture_io do
+          StdoutRouter.write_without_routing($stdout, 'unrouted')
+
+          assert_includes($stdout.string, 'unrouted')
+        end
+      end
+
+      def test_with_enabled_preserves_already_enabled_router
+        capture_io do
+          StdoutRouter.enable
+          StdoutRouter.with_enabled {}
+
+          assert_predicate(StdoutRouter, :enabled?)
+          cap = StdoutRouter::Capture.new { print('still routed') }
+          cap.run
+          assert_equal('still routed', cap.stdout)
+        ensure
+          StdoutRouter.disable
+        end
+      end
+
+      def test_with_enabled_unroutes_only_what_it_routed
+        capture_io do
+          StdoutRouter.with_enabled do
+            assert_predicate(StdoutRouter, :enabled?)
+          end
+
+          refute_predicate(StdoutRouter, :enabled?)
+          assert_empty($stdout.singleton_methods)
+        end
+      end
+
       def test_frame_can_autoload_after_router_is_enabled
         script = <<~RUBY
           require 'stringio'