claude-agent-sdk 1.2.0 → 1.2.1

data/lib/claude_agent_sdk/query.rb

@@ -9,6 +9,7 @@ require 'timeout'
 require_relative 'transport'
 require_relative 'errors'
 require_relative 'cancellation_signal'
+require_relative 'query/run_lifecycle'
 
 module ClaudeAgentSDK
   # Handles bidirectional control protocol on top of Transport
@@ -30,37 +31,22 @@ module ClaudeAgentSDK
     # @api private
     attr_reader :initialization_result
 
+    # When stdin may close (RunLifecycle). Read by specs.
+    #
+    # @api private
+    attr_reader :run_lifecycle
+
     CONTROL_REQUEST_TIMEOUT_ENV_VAR = 'CLAUDE_AGENT_SDK_CONTROL_REQUEST_TIMEOUT_SECONDS'
     DEFAULT_CONTROL_REQUEST_TIMEOUT_SECONDS = 1200.0
 
-    # Task types whose completion runs a follow-up turn, and which therefore
-    # may still need the control channel after the turn's result frame.
-    #
-    # Mirrors the set the CLI itself holds a result back for, which is
-    # narrower than its notion of "delegated agent work". The types left out
-    # are left out on purpose:
-    #   - background shells and monitors run indefinitely by design, so
-    #     deferring the close on one withholds it forever rather than briefly;
-    #   - teammates are long-lived too — their status stays running for their
-    #     whole lifetime, so they never settle the ledger;
-    #   - remote agents can be long-running monitors the CLI likewise refuses
-    #     to wait on.
-    # Anything added here must be a type that reliably reaches a terminal
-    # status, or it will hang the query (see #track_task_lifecycle).
-    DEFERRING_TASK_TYPES = %w[local_agent local_workflow].freeze
-
     # The CLI's own wait for background work once stdin is closed; the SDK
     # bounds its wait for the CLI's "idle" by the same value
-    # (#arm_run_end_ceiling, Python #1279).
+    # (RunLifecycle#arm, Python #1279).
     RUN_END_CEILING_ENV_VAR = 'CLAUDE_CODE_PRINT_BG_WAIT_CEILING_MS'
     DEFAULT_RUN_END_CEILING_MS = 600_000
     # The longest ceiling honored (~24.8 days), as in the TypeScript and
     # Python SDKs, whose timers cannot run longer.
     MAX_RUN_END_CEILING_MS = (2**31) - 1
-    # Frame types that mark a main-thread turn under way (when they carry no
-    # parent_tool_use_id), and the states that do not re-arm the ceiling.
-    TURN_FRAME_TYPES = %w[assistant stream_event].freeze
-    NON_RUNNING_SESSION_STATES = %w[idle requires_action].freeze
 
     # Read CLAUDE_CODE_PRINT_BG_WAIT_CEILING_MS from where the CLI gets it:
     # +options_env+ (ClaudeAgentOptions#env) overrides the inherited
@@ -76,33 +62,6 @@ module ClaudeAgentSDK
       digits.match?(/\A\d+\z/) ? Integer(digits, 10) : DEFAULT_RUN_END_CEILING_MS
     end
 
-    # One run's end, set at most once (the Python SDK's per-run anyio.Event).
-    # A waiter holds the object it started waiting on, so a run that ends and
-    # is then reopened (#reopen_run swaps in a fresh RunEnd) still releases the
-    # waiters the ended run woke, while later waits wait for the new run.
-    # Reactor-only, like every caller.
-    class RunEnd
-      def initialize
-        @ended = false
-        @condition = Async::Condition.new
-      end
-
-      def ended?
-        @ended
-      end
-
-      def end!
-        return if @ended
-
-        @ended = true
-        @condition.signal
-      end
-
-      def wait
-        @condition.wait until @ended
-      end
-    end
-
     # Apply ClaudeAgentOptions#verbatim_prompts to one outgoing user message
     # (Python #1269's stamp_user_message). Off: returns +message+ unchanged.
     # On: returns a new Hash with `client_composed: true`, dropping any
@@ -164,7 +123,7 @@ module ClaudeAgentSDK
                    exclude_dynamic_sections: nil, system_prompt_snapshot: nil, skills: nil,
                    forward_subagent_text: false, agent_progress_summaries: nil,
                    callback_scheduling: :thread, callback_wrapper: nil, verbatim_prompts: false,
-                   run_end_ceiling_ms: DEFAULT_RUN_END_CEILING_MS)
+                   run_end_ceiling_ms: DEFAULT_RUN_END_CEILING_MS, sleeper: nil)
       @transport = transport
       @is_streaming_mode = is_streaming_mode
       @can_use_tool = can_use_tool
@@ -179,7 +138,6 @@ module ClaudeAgentSDK
       @forward_subagent_text = forward_subagent_text
       @agent_progress_summaries = agent_progress_summaries
       @verbatim_prompts = verbatim_prompts
-      @run_end_ceiling_ms = run_end_ceiling_ms
 
       # Control protocol state
       @pending_control_responses = {}
@@ -197,36 +155,14 @@ module ClaudeAgentSDK
       @message_queue = Async::Queue.new
       # Set once #receive_messages has consumed the read loop's end sentinel.
       @stream_ended = false
-      # Ends when the run is over, so the stdin-closing waiter can wake; see
-      # #read_messages and @inflight_tasks below (Python #1088, #1190/#1279).
-      # Work the CLI takes up after the run ended swaps in a fresh one
-      # (#reopen_run).
-      @run_end = RunEnd.new
-      @result_received = false
-      # The CLI's latest session_state_changed state, or nil while it sends
-      # none (a CLI too old to honor CLAUDE_CODE_SDK_READS_SESSION_STATE). A
-      # CLI that reports state stays "running" while a background agent is
-      # live or its completion is still to be handled, and reports "idle" once
-      # no further turn is owed.
-      @session_state = nil
-      # Ends the run if no new turn starts within the ceiling after a result
-      # (#arm_run_end_ceiling). The generation tells a sleeper that woke after
-      # it was cleared or re-armed to stand down.
-      @run_end_ceiling_task = nil
-      @run_end_ceiling_generation = 0
-      # A main-thread turn is under way (its assistant/stream_event frames
-      # have started and its result has not arrived): the ceiling counts only
-      # the wait between turns, so it is not armed meanwhile.
-      @turn_in_progress = false
-      # Set once stdin is closed or the reader is gone: the run then stays
-      # ended, since nothing can wait on a reopened one.
-      @run_final = false
-      # Task IDs of started-but-not-finished deferring tasks. A result frame
-      # only ends one turn, not the run: a background task keeps running past
-      # it and still needs stdin for hook/SDK-MCP control responses (Python
-      # #1088/#1103), so a result that arrives while this set is non-empty
-      # must not close stdin.
-      @inflight_tasks = Set.new
+      # When stdin may close: the end of the run (Python #1088, #1190/#1279).
+      # +sleeper+ replaces the ceiling's timer in specs; the default reads @task
+      # when it is called, which is after #start.
+      @run_lifecycle = RunLifecycle.new(
+        ceiling_ms: run_end_ceiling_ms, sleeper: sleeper || method(:sleep_on_read_task),
+        bidirectional_needs: -> { bidirectional_needs? }, closed: -> { @closed },
+        control_requests_in_flight: -> { !@inflight_control_request_tasks.empty? }
+      )
       # Set to the result payload when the most recent message is a result
       # with is_error=true. Used to replace the generic "exit code 1"
       # ProcessError with a ResultError carrying what the CLI already
@@ -294,7 +230,7 @@ module ClaudeAgentSDK
           # A Hash stands for the AgentDefinition with the same attributes
           # (the options signature allows either). Built through .new: the
           # user wrote it, so a misspelled key raises as on the typed class.
-          agent_def = AgentDefinition.new(agent_def) if agent_def.is_a?(Hash)
+          agent_def = OptionForms.agent_definition(agent_def)
           {
             description: agent_def.description,
             prompt: agent_def.prompt,
@@ -488,43 +424,27 @@ module ClaudeAgentSDK
           @transcript_mirror_batcher&.enqueue(message[:filePath] || message[:file_path], message[:entries] || [])
           next
         else
-          # Track task lifecycle frames so results can tell "one turn ended"
-          # apart from "the run is done" (Python #1088/#1103).
-          if msg_type == 'system'
-            had_tasks_in_flight = !@inflight_tasks.empty?
-            track_task_lifecycle(message)
-            # The ceiling left the last tracked agent alone; the wait between
-            # turns starts over now that it settled.
-            rearm_run_end_ceiling_between_turns if had_tasks_in_flight && @inflight_tasks.empty?
-            if message[:subtype] == 'session_state_changed'
-              on_session_state(message[:state])
-              # Frames the CLI sent only because the transport asked for them
-              # (CLAUDE_CODE_SDK_READS_SESSION_STATE); the caller did not opt
-              # in, so they never reach the stream, observers or the parser.
-              next if message[:sdk_host_only] == true
-            end
-          end
+          session_state_frame = msg_type == 'system' && message[:subtype] == 'session_state_changed'
+          # Flush the mirror before signaling/yielding the result so a
+          # consumer observing the result sees an up-to-date store for the turn.
+          flush_transcript_mirror if msg_type == 'result'
+          # Every frame that gets here drives the run's end (task lifecycle,
+          # session state, results, main-thread turns), whether or not it goes
+          # on to the stream.
+          @run_lifecycle.frame(message)
+          # Frames the CLI sent only because the transport asked for them
+          # (CLAUDE_CODE_SDK_READS_SESSION_STATE); the caller did not opt in,
+          # so they never reach the stream, observers or the parser.
+          next if session_state_frame && message[:sdk_host_only] == true
 
           if msg_type == 'result'
-            # Flush the mirror before signaling/yielding the result so a
-            # consumer observing the result sees an up-to-date store for the turn.
-            flush_transcript_mirror
-            on_result
             @last_error_result = message[:is_error] ? message : nil
-          elsif !(msg_type == 'system' && message[:subtype] == 'session_state_changed')
+          elsif !session_state_frame
             # Anything other than the post-turn session_state_changed marker
             # means the conversation moved on; a ProcessError now is a fresh
             # crash, not the expected exit from a prior error result. Mirrors
             # the Python/TypeScript SDK reset logic.
             @last_error_result = nil
-            # A main-thread turn is under way, so the ceiling stops (it counts
-            # only the wait between turns, as the CLI's does) and the run
-            # reopens even if the ceiling ended it while no state changed.
-            if TURN_FRAME_TYPES.include?(msg_type) && message[:parent_tool_use_id].nil?
-              @turn_in_progress = true
-              reopen_run
-              clear_run_end_ceiling
-            end
           end
           # Regular SDK messages go to the queue
           @message_queue.enqueue(message)
@@ -576,189 +496,12 @@ module ClaudeAgentSDK
         # Unblock the stdin-closing waiter so it doesn't stall on early exit;
         # with the reader gone the run stays ended. Also stops a pending
         # ceiling sleeper, which would otherwise keep the reactor alive.
-        @run_final = true
-        end_run
+        @run_lifecycle.reader_gone
         # Always signal end of stream
         @message_queue.enqueue({ type: 'end' })
       end
     end
 
-    # Track in-flight tasks from `system` task lifecycle frames.
-    #
-    # `task_started` marks a task in flight; `task_notification` or a
-    # `task_updated` patch with a terminal status clears it. Terminal
-    # completion can arrive as either frame (not every terminal task emits a
-    # notification), so both are handled; Set deletion keeps the pair
-    # idempotent.
-    #
-    # This is a mitigation, not a complete answer to Python #1088. An empty
-    # set means "nothing we know of is running", which is not the same as
-    # "the run is over": a task that settles *before* the turn's result frame
-    # leaves the set empty at that result, so stdin closes even though the
-    # completion may still wake the parent for a continuation turn. What this
-    # does fix is the common ordering, where the task outlives the turn that
-    # spawned it.
-    #
-    # Only delegated agent work is tracked (DEFERRING_TASK_TYPES). A
-    # background *shell* is also reported through these frames, but it may
-    # never reach a terminal status, and the CLI in stream-json mode only
-    # exits on stdin EOF — tracking one would withhold the close forever.
-    #
-    # `background_tasks_changed` is deliberately not consumed, in either
-    # direction: its payload is the live *background* set, while a subagent
-    # is registered in the foreground and only flips to backgrounded later
-    # without a second task_started, so narrowing against the snapshot would
-    # drop an agent that goes on to outlive its turn, and widening from it
-    # could admit an id no later frame ever clears (observer agents suppress
-    # both their start and terminal frames).
-    def track_task_lifecycle(message)
-      task_id = message[:task_id]
-      return if task_id.nil? || task_id.to_s.empty?
-
-      case message[:subtype]
-      when 'task_started'
-        @inflight_tasks.add(task_id) if DEFERRING_TASK_TYPES.include?(message[:task_type])
-      when 'task_notification'
-        @inflight_tasks.delete(task_id)
-      when 'task_updated'
-        patch = message[:patch]
-        status = patch.is_a?(Hash) ? patch[:status] : nil
-        @inflight_tasks.delete(task_id) if TERMINAL_TASK_STATUSES.include?(status)
-      end
-    end
-
-    # A result ends a turn, not necessarily the run: a background agent that
-    # finished just before it still wakes the session for another turn, whose
-    # hook, permission and SDK MCP requests need stdin (Python #1190/#1279). A
-    # CLI that reports session state stays "running" while such a turn is
-    # owed, so wait for "idle" (some hosts send it just before the result).
-    # Without state events the result is all there is to go on.
-    def on_result
-      @result_received = true
-      @turn_in_progress = false
-      if @session_state.nil? || @session_state == 'idle' || !bidirectional_needs?
-        maybe_end_run
-      elsif @session_state != 'requires_action'
-        # While the SDK is still answering a request the ceiling waits for
-        # the "running" that follows.
-        arm_run_end_ceiling
-      end
-    end
-
-    # Track the CLI's session_state_changed state (Python #1279).
-    def on_session_state(state)
-      @session_state = state
-      if state == 'idle'
-        maybe_end_run if @result_received
-        return
-      end
-      # Work the CLI took up after the run ended (a finished background task
-      # woke it) reopens the run until the next "idle".
-      reopen_run
-      if state == 'requires_action'
-        # The host is answering a request; stdin must outlast it.
-        clear_run_end_ceiling
-      else
-        rearm_run_end_ceiling_between_turns
-      end
-    end
-
-    # End the run unless a tracked background task is still in flight: such
-    # a task may still need hook/SDK-MCP control responses over stdin
-    # (Python #1088), and its completion wakes the parent for a follow-up
-    # turn whose result (or "idle") ends the run then. A CLI that reports
-    # session state never reports "idle" with an agent still live, so this
-    # matters for CLIs that report "idle" at every turn end or not at all.
-    def maybe_end_run
-      end_run if @inflight_tasks.empty?
-    end
-
-    # The run is over: wake the stdin-closing waiter. Idempotent.
-    #
-    # The run is marked ended BEFORE the ceiling is cleared: stopping the
-    # sleeper task yields to whatever else is ready, and a #stream_input task
-    # that writes its next message in that gap must find the run already
-    # ended, so that #reopen_run gives the message a run of its own. Cleared
-    # first, the message joined the run that was about to end, and stdin
-    # closed before the message's own run had produced a frame.
-    def end_run
-      @run_end.end!
-      clear_run_end_ceiling
-    end
-
-    # Reopen an ended run for work that started after it ended. A waiter the
-    # ended run already woke still closes stdin; this makes a later wait
-    # (stream_input's, once its prompts are all written) wait for the new
-    # work too. Once stdin is closed, or the reader is gone, the run stays
-    # ended.
-    def reopen_run
-      @run_end = RunEnd.new if @run_end.ended? && !@run_final
-    end
-
-    # End the run anyway once the ceiling passes with no new turn. The CLI's
-    # own background-wait ceiling only counts once stdin is closed, so without
-    # this, work that never finishes would hold "running", and stdin, open
-    # forever. It counts only the wait between turns: restarted at each result
-    # and whenever the CLI reports "running" again, cleared by main-thread
-    # turn activity and by "requires_action", never armed mid-turn.
-    #
-    # The sleeper is a child of the read task, NOT a #spawn_task entry: it is re-armed at every
-    # result and every "running", and @child_tasks never prunes. The read
-    # task's stop cascades to it, and every exit path clears it (#end_run in
-    # the read loop's ensure, #wait_for_result_and_end_input's ensure), so a
-    # pending sleeper can never keep the enclosing reactor alive.
-    def arm_run_end_ceiling
-      clear_run_end_ceiling
-      # A frame read while close is under way (the result branch flushes the
-      # mirror first) must not leave a sleeper behind either.
-      return if @run_end_ceiling_ms <= 0 || @run_end.ended? || @run_final || @closed ||
-                @turn_in_progress || !bidirectional_needs?
-
-      generation = @run_end_ceiling_generation
-      seconds = [@run_end_ceiling_ms, MAX_RUN_END_CEILING_MS].min / 1000.0
-      @run_end_ceiling_task = @task.async do
-        sleep seconds
-        end_run_at_ceiling(generation)
-      end
-    end
-
-    # Restart the ceiling if the run is between turns, past a result, with
-    # the CLI still reporting work ("running").
-    def rearm_run_end_ceiling_between_turns
-      return unless @result_received
-      return if @session_state.nil? || NON_RUNNING_SESSION_STATES.include?(@session_state)
-
-      arm_run_end_ceiling
-    end
-
-    def end_run_at_ceiling(generation)
-      # Cleared or re-armed while this sleeper was already waking up.
-      return unless generation == @run_end_ceiling_generation
-
-      # Detach before ending the run: #end_run clears the ceiling, and
-      # clearing it must not stop the task running this very method.
-      @run_end_ceiling_task = nil
-      # A tracked background agent still running may still need stdin for its
-      # hook, permission and SDK MCP requests (Python #1088), so it is not cut
-      # off; the ceiling starts over once it settles (#read_messages).
-      return unless @inflight_tasks.empty?
-
-      # A control request the SDK is still answering (a slow hook or SDK MCP
-      # tool) must be able to write its reply, so the clock starts over rather
-      # than closing stdin under it. Ruby-only guard: Python relies on the CLI
-      # reporting requires_action for every such request.
-      return arm_run_end_ceiling unless @inflight_control_request_tasks.empty?
-
-      end_run
-    end
-
-    def clear_run_end_ceiling
-      @run_end_ceiling_generation += 1
-      task = @run_end_ceiling_task
-      @run_end_ceiling_task = nil
-      task&.stop
-    end
-
     # Whether the CLI may still send control requests that need a reply.
     #
     # SDK MCP servers, hooks and the can_use_tool permission callback are all
@@ -777,6 +520,17 @@ module ClaudeAgentSDK
       !@sdk_mcp_servers.empty? || !@hooks.empty? || !!@can_use_tool
     end
 
+    # The run-end ceiling's timer (RunLifecycle's sleeper): runs +on_wake+ once
+    # +seconds+ have passed. A child of the read task, NOT a #spawn_task entry:
+    # it is re-armed at every result and every "running", and @child_tasks
+    # never prunes. The read task's stop cascades to it.
+    def sleep_on_read_task(seconds, &on_wake)
+      @task.async do
+        sleep seconds
+        on_wake.call
+      end
+    end
+
     # Flush the transcript-mirror batcher, swallowing errors — a mirror failure
     # must never propagate into the read loop or its teardown.
     def flush_transcript_mirror
@@ -1744,7 +1498,7 @@ module ClaudeAgentSDK
     # wait BETWEEN turns is bounded: if the CLI still reports "running"
     # CLAUDE_CODE_PRINT_BG_WAIT_CEILING_MS (10 minutes by default, 0 for no
     # limit) after a result with no new turn, the run ends anyway
-    # (#arm_run_end_ceiling). A turn under way, a request the SDK is still
+    # (RunLifecycle#arm). A turn under way, a request the SDK is still
     # answering and a tracked background agent still in flight (Python #1088)
     # stop that clock. The run is guaranteed to end: as above, or in
     # read_messages' ensure when the process exits early.
@@ -1755,10 +1509,9 @@ module ClaudeAgentSDK
     # control requests from that later turn can find stdin closed.
     # Single-message and String prompts are fully covered.
     def wait_for_result_and_end_input
-      @run_end.wait if bidirectional_needs?
+      @run_lifecycle.wait
     ensure
-      @run_final = true
-      clear_run_end_ceiling
+      @run_lifecycle.stdin_closing
       @transport.end_input
     end
 
@@ -1781,9 +1534,7 @@ module ClaudeAgentSDK
         line = Query.serialize_user_message(message, @verbatim_prompts)
         # This message owes a run of its own, result included: an earlier
         # one having ended does not end it.
-        reopen_run
-        @result_received = false
-        clear_run_end_ceiling
+        @run_lifecycle.message_will_write
         writeln(line)
         wrote_message = true
       end
@@ -1819,7 +1570,7 @@ module ClaudeAgentSDK
         if wrote_message
           wait_for_result_and_end_input
         else
-          @run_final = true
+          @run_lifecycle.empty_input
           @transport.end_input
         end
       end