Class: ClaudeAgentSDK::Query

Inherits:
Object
  • Object
show all
Defined in:
lib/claude_agent_sdk/query.rb

Overview

Handles bidirectional control protocol on top of Transport

This class manages:

  • Control request/response routing
  • Hook callbacks
  • Tool permission callbacks
  • Message streaming
  • Initialization handshake

Defined Under Namespace

Classes: ThreadWaiter

Constant Summary collapse

CONTROL_REQUEST_TIMEOUT_ENV_VAR = 
'CLAUDE_AGENT_SDK_CONTROL_REQUEST_TIMEOUT_SECONDS'
DEFAULT_CONTROL_REQUEST_TIMEOUT_SECONDS = 
1200.0
DEFERRING_TASK_TYPES =

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).

%w[local_agent local_workflow].freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(transport:, is_streaming_mode:, can_use_tool: nil, hooks: nil, sdk_mcp_servers: nil, agents: nil, exclude_dynamic_sections: nil, skills: nil) ⇒ Query

Returns a new instance of Query.



65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
 
# File 'lib/claude_agent_sdk/query.rb', line 65

def initialize(transport:, is_streaming_mode:, can_use_tool: nil, hooks: nil, sdk_mcp_servers: nil, agents: nil,
               exclude_dynamic_sections: nil, skills: nil)
  @transport = transport
  @is_streaming_mode = is_streaming_mode
  @can_use_tool = can_use_tool
  @hooks = hooks || {}
  @sdk_mcp_servers = sdk_mcp_servers || {}
  @agents = agents
  @exclude_dynamic_sections = exclude_dynamic_sections
  @skills = skills

  # Control protocol state
  @pending_control_responses = {}
  @pending_control_results = {}
  @hook_callbacks = {}
  @hook_callback_timeouts = {}
  @next_callback_id = 0
  @request_counter = 0
  @request_counter_mutex = Mutex.new
  @inflight_control_request_tasks = {}

  # Message stream
  @message_queue = Async::Queue.new
  # Set when a run-ending result arrives (a result frame with no tasks in
  # flight) so the stdin-closing waiter can wake. Named for history — it
  # once tracked the literal first result.
  @first_result_received = 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
  @last_error_result_text = nil
  @first_result_condition = Async::Condition.new
  @task = nil
  @child_tasks = []
  @initialized = false
  @closed = false
  @initialization_result = nil
  @transcript_mirror_batcher = nil

  # Cross-thread close marshaling (see #close). Thread::Queue is the one
  # primitive that is both fiber-scheduler-aware on the reactor side and
  # thread-safe on the caller side (push -> scheduler#unblock).
  @close_requests = ::Thread::Queue.new
  @close_watcher = nil
  @owning_scheduler = nil
  # First-caller-wins guard for close_now (see there).
  @close_mutex = Mutex.new
  @close_started = false
end

Instance Attribute Details

#is_streaming_mode ⇒ Object (readonly)

Returns the value of attribute is_streaming_mode.



22
23
24
 
# File 'lib/claude_agent_sdk/query.rb', line 22

def is_streaming_mode
  @is_streaming_mode
end

#sdk_mcp_servers ⇒ Object (readonly)

Returns the value of attribute sdk_mcp_servers.



22
23
24
 
# File 'lib/claude_agent_sdk/query.rb', line 22

def sdk_mcp_servers
  @sdk_mcp_servers
end

#transport ⇒ Object (readonly)

Returns the value of attribute transport.



22
23
24
 
# File 'lib/claude_agent_sdk/query.rb', line 22

def transport
  @transport
end

Instance Method Details

#close ⇒ Object

Close the query and transport.

Callable from any thread. Async::Task#stop needs the reactor's Fiber.scheduler (per-thread), which foreign callers — FiberBoundary workers running a tool handler / hook that calls client.disconnect, or plain user threads — don't have; stopping from one raised NoMethodError and left the read/child tasks running. Such callers hand the close to the reactor-side watcher (spawned in #start) and wait for it to finish, so close semantics are identical regardless of the calling thread.



1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
 
# File 'lib/claude_agent_sdk/query.rb', line 1283

def close
  if @close_watcher&.alive? && !Fiber.scheduler.equal?(@owning_scheduler)
    marshal_close_to_reactor
  else
    # Same scheduler (reactor-side caller, including the watcher itself),
    # or no live watcher: when the reactor is gone its task fibers are
    # dead, so stopping them no longer touches Fiber.scheduler.
    close_now
  end
end

#get_context_usage ⇒ Hash

Get a breakdown of current context window usage by category.

Returns:

  • (Hash) —

    Context usage response with categories, totalTokens, maxTokens, etc.



1116
1117
1118
 
# File 'lib/claude_agent_sdk/query.rb', line 1116

def get_context_usage
  send_control_request({ subtype: 'get_context_usage' })
end

#get_mcp_status ⇒ Hash

Get current MCP server connection status (only works with streaming mode)

Returns:

  • (Hash) —

    MCP status information, including mcpServers list



1122
1123
1124
 
# File 'lib/claude_agent_sdk/query.rb', line 1122

def get_mcp_status
  send_control_request({ subtype: 'mcp_status' })
end

#initialize_protocol ⇒ Hash?

Initialize control protocol if in streaming mode

Returns:

  • (Hash, nil) —

    Initialize response with supported commands, or nil if not streaming



120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
 
# File 'lib/claude_agent_sdk/query.rb', line 120

def initialize_protocol
  return nil unless @is_streaming_mode

  # Build hooks configuration for initialization
  hooks_config = {}
  if @hooks && !@hooks.empty?
    @hooks.each do |event, matchers|
      next if matchers.nil? || matchers.empty?

      hooks_config[event] = []
      matchers.each do |matcher|
        callback_ids = []
        (matcher[:hooks] || []).each do |callback|
          callback_id = "hook_#{@next_callback_id}"
          @next_callback_id += 1
          @hook_callbacks[callback_id] = callback
          @hook_callback_timeouts[callback_id] = matcher[:timeout] if matcher[:timeout]
          callback_ids << callback_id
        end
        matcher_config = {
          matcher: matcher[:matcher],
          hookCallbackIds: callback_ids
        }
        # Wire field is literal "timeout" in SECONDS, per matcher,
        # omitted when absent (Python _internal/query.py parity — no
        # camelCase, no ms conversion). Local enforcement via
        # @hook_callback_timeouts stays as defense-in-depth for CLIs
        # that ignore the field.
        matcher_config[:timeout] = matcher[:timeout] if matcher[:timeout]
        hooks_config[event] << matcher_config
      end
    end
  end

  # Build agents dict for initialization
  agents_dict = nil
  if @agents
    agents_dict = @agents.transform_values do |agent_def|
      {
        description: agent_def.description,
        prompt: agent_def.prompt,
        tools: agent_def.tools,
        disallowedTools: agent_def.disallowed_tools,
        model: agent_def.model,
        skills: agent_def.skills,
        memory: agent_def.memory,
        mcpServers: agent_def.mcp_servers,
        initialPrompt: agent_def.initial_prompt,
        maxTurns: agent_def.max_turns,
        background: agent_def.background,
        effort: agent_def.effort,
        permissionMode: agent_def.permission_mode
      }.compact
    end
  end

  # Send initialize request
  request = {
    subtype: 'initialize',
    hooks: hooks_config.empty? ? nil : hooks_config,
    agents: agents_dict
  }
  request[:excludeDynamicSections] = @exclude_dynamic_sections unless @exclude_dynamic_sections.nil?
  # 'all' and omitted are equivalent at the wire level (no filter), so
  # only send the field when it's an explicit list (mirrors Python).
  request[:skills] = @skills if @skills.is_a?(Array)

  response = send_control_request(request)
  @initialized = true
  @initialization_result = response
  response
end

#interrupt ⇒ Object

Send interrupt control request



1127
1128
1129
 
# File 'lib/claude_agent_sdk/query.rb', line 1127

def interrupt
  send_control_request({ subtype: 'interrupt' })
end

#mirror_batches_dropped? ⇒ Boolean

True when the mirror dropped at least one batch (store copy incomplete). Meaningful after #close, which runs the final flush. Consulted by the resume-from-store teardown to decide whether the materialized temp dir holds the only copy of some turns and must be preserved.

Returns:

  • (Boolean) 


266
267
268
 
# File 'lib/claude_agent_sdk/query.rb', line 266

def mirror_batches_dropped?
  !!@transcript_mirror_batcher&.batches_dropped?
end

#receive_messages(&block) ⇒ Object

Receive SDK messages (not control messages)



1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
 
# File 'lib/claude_agent_sdk/query.rb', line 1258

def receive_messages(&block)
  return enum_for(:receive_messages) unless block

  # NOT Kernel#loop: it rescues StopIteration, so a user block leaking one
  # (e.g. calling .next on an exhausted Enumerator) would silently end
  # reception — ResultMessage dropped, the query reported as complete, and
  # on_error never fired. `while` propagates it like any other error.
  while true # rubocop:disable Style/InfiniteLoop
    message = @message_queue.dequeue
    break if message[:type] == 'end'
    raise message[:error] if message[:type] == 'error'

    block.call(message)
  end
end

#reconnect_mcp_server(server_name) ⇒ Object

Reconnect a failed MCP server

Parameters:

  • server_name (String) —

    Name of the MCP server to reconnect



1149
1150
1151
1152
1153
1154
 
# File 'lib/claude_agent_sdk/query.rb', line 1149

def reconnect_mcp_server(server_name)
  send_control_request({
                         subtype: 'mcp_reconnect',
                         serverName: server_name
                       })
end

#report_mirror_error(key, error) ⇒ Object

Synthesize a mirror_error system message and put it on the SDK message stream so consumers learn a mirror batch was dropped after exhausting retries. Non-blocking: the message queue is unbounded, so unlike the Python SDK there is no buffer-full drop path.



274
275
276
277
278
279
280
281
282
283
284
 
# File 'lib/claude_agent_sdk/query.rb', line 274

def report_mirror_error(key, error)
  session_id = key && (key['session_id'] || key[:session_id])
  @message_queue.enqueue(
    type: 'system',
    subtype: 'mirror_error',
    error: error,
    key: key,
    uuid: SecureRandom.uuid,
    session_id: session_id || ''
  )
end

#rewind_files(user_message_uuid) ⇒ Object

Rewind files to a previous checkpoint (v0.1.15+) Restores file state to what it was at the given user message Requires enable_file_checkpointing to be true in options

Parameters:

  • user_message_uuid (String) —

    The UUID of the UserMessage to rewind to



1180
1181
1182
1183
1184
1185
 
# File 'lib/claude_agent_sdk/query.rb', line 1180

def rewind_files(user_message_uuid)
  send_control_request({
                         subtype: 'rewind_files',
                         user_message_id: user_message_uuid
                       })
end

#set_model(model) ⇒ Object

Change the AI model



1140
1141
1142
1143
1144
1145
 
# File 'lib/claude_agent_sdk/query.rb', line 1140

def set_model(model)
  send_control_request({
                         subtype: 'set_model',
                         model: model
                       })
end

#set_permission_mode(mode) ⇒ Object

Change permission mode



1132
1133
1134
1135
1136
1137
 
# File 'lib/claude_agent_sdk/query.rb', line 1132

def set_permission_mode(mode)
  send_control_request({
                         subtype: 'set_permission_mode',
                         mode: mode
                       })
end

#set_transcript_mirror_batcher(batcher) ⇒ Object

Install the transcript-mirror batcher fed by transcript_mirror frames (Client mode with a session_store). nil disables mirroring.



258
259
260
 
# File 'lib/claude_agent_sdk/query.rb', line 258

def set_transcript_mirror_batcher(batcher)
  @transcript_mirror_batcher = batcher
end

#spawn_task(&block) ⇒ Object

Spawn a child task that is stopped by #close (mirrors the Python SDK's Query#spawn_task / _child_tasks). Used for background input streaming so a dying read loop or #close can never strand the stream task and hang the enclosing Async reactor.

NOTE: intentionally a partial mirror — Python prunes completed tasks via add_done_callback(_child_tasks.discard); here entries live until #close. Fine for the current one-shot call sites (max two tasks per Query); do not route per-request work (control handlers, per-turn streams) through this without adding completion-based removal.

Raises:



247
248
249
250
251
252
253
254
 
# File 'lib/claude_agent_sdk/query.rb', line 247

def spawn_task(&block)
  parent = Async::Task.current?
  raise CLIConnectionError, 'Query#spawn_task must be called inside an Async{} block' unless parent

  task = parent.async(&block)
  @child_tasks << task
  task
end

#start ⇒ Object

Start reading messages from transport.

Spawns read_messages as a direct child task of the current Async task and stores that child in @task. An earlier version wrapped task.async { read_messages } inside an outer Async do ... end and assigned the outer task to @task; the outer task completed almost immediately after spawning, so close's @task.stop never reached the actual read_messages fiber and the read loop kept running until the transport raised. Now @task.stop stops the read loop.

Must be called inside an Async{} block (matches query() which wraps its own internals in Async, and the documented Client#connect pattern). If invoked outside a reactor, raise a clear error rather than letting Async::Task.current raise an opaque "No async task available!" — earlier versions of this method appeared to work from synchronous callers but actually hung indefinitely because the outer Async{} root task waited for read_messages to finish, which never happens for a live Client.

Raises:



211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
 
# File 'lib/claude_agent_sdk/query.rb', line 211

def start
  return if @task

  parent = Async::Task.current?
  raise CLIConnectionError, 'Query#start must be called inside an Async{} block (e.g. wrap Client#connect in Async{...})' unless parent

  @owning_scheduler = Fiber.scheduler
  @task = parent.async { read_messages }
  # Reactor-side agent for #close calls arriving from foreign threads
  # (FiberBoundary callbacks, plain user threads): Async::Task#stop needs
  # the owning thread's Fiber.scheduler, so the off-thread caller hands the
  # whole close over and waits. Transient: must never keep the reactor
  # alive, and is stopped automatically when the parent task finishes.
  # One-shot: after serving a close it is done; a reactor-side close wakes
  # it via @close_requests.close (pop -> nil) so it exits without serving.
  @close_watcher = parent.async(transient: true) do
    if (reply = @close_requests.pop)
      begin
        close
      ensure
        reply << true
      end
    end
  end
end

#stop_task(task_id) ⇒ Object

Stop a running background task

Parameters:

  • task_id (String) —

    The ID of the task to stop



1169
1170
1171
1172
1173
1174
 
# File 'lib/claude_agent_sdk/query.rb', line 1169

def stop_task(task_id)
  send_control_request({
                         subtype: 'stop_task',
                         task_id: task_id
                       })
end

#stream_input(stream) ⇒ Object

Stream input messages to transport. NOTE: iteration runs on the reactor (the deliberate FiberBoundary carve-out — see fiber_boundary.rb): scheduler-aware blocking (Thread::Queue#pop, sleep, socket IO) parks only this task; CPU-bound or scheduler-opaque work in the enumerator must be moved to a producer Thread by the user.



1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
 
# File 'lib/claude_agent_sdk/query.rb', line 1214

def stream_input(stream)
  wrote_message = false
  stream.each do |message|
    break if @closed
    serialized = message.is_a?(Hash) ? JSON.generate(message) : message.to_s
    writeln(serialized)
    wrote_message = true
  end
rescue StandardError => e
  # Log error but don't raise
  warn "Error streaming input: #{e.message}"
ensure
  # Three teardown shapes:
  # - #close in progress (@closed, Async::Stop unwinding): do nothing —
  #   the transport is about to be closed, and waiting on
  #   @first_result_condition inside a stopping fiber could suspend
  #   teardown. Mirrors Python, where cancellation skips this entirely.
  # - A turn is in flight (some message reached the CLI): hold stdin
  #   open until its first result so hooks/SDK MCP control replies can
  #   still be written (no timeout — the result or process exit is
  #   guaranteed to signal).
  # - No complete message ever reached the CLI (empty stream, or the
  #   stream raised before the first write): no result can ever arrive,
  #   so waiting would park query() forever beside an idle CLI. Close
  #   stdin so the CLI sees EOF and exits. Deliberate improvement over
  #   Python, which leaves stdin open and hangs on this path.
  unless @closed
    if wrote_message
      wait_for_result_and_end_input
    else
      @transport.end_input
    end
  end
end

#toggle_mcp_server(server_name, enabled) ⇒ Object

Enable or disable an MCP server

Parameters:

  • server_name (String) —

    Name of the MCP server

  • enabled (Boolean) —

    Whether to enable or disable



1159
1160
1161
1162
1163
1164
1165
 
# File 'lib/claude_agent_sdk/query.rb', line 1159

def toggle_mcp_server(server_name, enabled)
  send_control_request({
                         subtype: 'mcp_toggle',
                         serverName: server_name,
                         enabled: enabled
                       })
end

#wait_for_result_and_end_input ⇒ Object

Wait for a run-ending result before closing stdin when hooks or SDK MCP servers may still need to exchange control messages with the CLI. The control protocol requires stdin to stay open for the entire turn (hook replies, can_use_tool replies and SDK MCP tool results are all written to stdin), so no timeout is applied — closing stdin mid-turn silently broke hooks/MCP on turns longer than the old 60s bound (mirrors Python SDK commit c3d96cb). A result frame ends one turn, not necessarily the run: while background tasks are in flight the result branch withholds the signal (Python #1088/#1103), and each deferring task's completion wakes the parent for a follow-up turn that ends in another result. The condition is guaranteed to be signaled: by the result branch in read_messages once no tasks are in flight, or by its ensure block when the process exits early.



1200
1201
1202
1203
1204
1205
1206
1207
 
# File 'lib/claude_agent_sdk/query.rb', line 1200

def wait_for_result_and_end_input
  if !@first_result_received &&
     ((@sdk_mcp_servers && !@sdk_mcp_servers.empty?) || (@hooks && !@hooks.empty?))
    @first_result_condition.wait
  end
ensure
  @transport.end_input
end

#write(string) ⇒ Object 



1253
1254
1255
 
# File 'lib/claude_agent_sdk/query.rb', line 1253

def write(string)
  @transport.write(string)
end

#writeln(string) ⇒ Object 



1249
1250
1251
 
# File 'lib/claude_agent_sdk/query.rb', line 1249

def writeln(string)
  write string.end_with?("\n") ? string : "#{string}\n"
end