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, callback_scheduling: :thread, callback_wrapper: 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
117
118
 
# 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, callback_scheduling: :thread, callback_wrapper: nil)
  @transport = transport
  @is_streaming_mode = is_streaming_mode
  @can_use_tool = can_use_tool
  @hooks = hooks || {}
  @sdk_mcp_servers = sdk_mcp_servers || {}
  @callback_scheduling = callback_scheduling || :thread
  @callback_wrapper = callback_wrapper
  @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.



1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
 
# File 'lib/claude_agent_sdk/query.rb', line 1342

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.



1175
1176
1177
 
# File 'lib/claude_agent_sdk/query.rb', line 1175

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



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

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



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
192
193
 
# File 'lib/claude_agent_sdk/query.rb', line 122

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



1186
1187
1188
 
# File 'lib/claude_agent_sdk/query.rb', line 1186

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) 


268
269
270
 
# File 'lib/claude_agent_sdk/query.rb', line 268

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

#receive_messages(&block) ⇒ Object

Receive SDK messages (not control messages)



1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
 
# File 'lib/claude_agent_sdk/query.rb', line 1317

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



1208
1209
1210
1211
1212
1213
 
# File 'lib/claude_agent_sdk/query.rb', line 1208

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 (timeouts immediately, other failures after up to three attempts). Non-blocking: the message queue is unbounded, so unlike the Python SDK there is no buffer-full drop path.



277
278
279
280
281
282
283
284
285
286
287
 
# File 'lib/claude_agent_sdk/query.rb', line 277

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



1239
1240
1241
1242
1243
1244
 
# File 'lib/claude_agent_sdk/query.rb', line 1239

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



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

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

#set_permission_mode(mode) ⇒ Object

Change permission mode



1191
1192
1193
1194
1195
1196
 
# File 'lib/claude_agent_sdk/query.rb', line 1191

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.



260
261
262
 
# File 'lib/claude_agent_sdk/query.rb', line 260

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:



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

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:



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

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



1228
1229
1230
1231
1232
1233
 
# File 'lib/claude_agent_sdk/query.rb', line 1228

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.



1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
 
# File 'lib/claude_agent_sdk/query.rb', line 1273

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



1218
1219
1220
1221
1222
1223
1224
 
# File 'lib/claude_agent_sdk/query.rb', line 1218

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.



1259
1260
1261
1262
1263
1264
1265
1266
 
# File 'lib/claude_agent_sdk/query.rb', line 1259

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 



1312
1313
1314
 
# File 'lib/claude_agent_sdk/query.rb', line 1312

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

#writeln(string) ⇒ Object 



1308
1309
1310
 
# File 'lib/claude_agent_sdk/query.rb', line 1308

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