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

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.



48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
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
 
# File 'lib/claude_agent_sdk/query.rb', line 48

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
  @first_result_received = false
  @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.



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

def is_streaming_mode
  @is_streaming_mode
end

#sdk_mcp_servers ⇒ Object (readonly)

Returns the value of attribute sdk_mcp_servers.



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

def sdk_mcp_servers
  @sdk_mcp_servers
end

#transport ⇒ Object (readonly)

Returns the value of attribute transport.



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

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.



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

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.



1036
1037
1038
 
# File 'lib/claude_agent_sdk/query.rb', line 1036

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



1042
1043
1044
 
# File 'lib/claude_agent_sdk/query.rb', line 1042

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



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
119
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
 
# File 'lib/claude_agent_sdk/query.rb', line 94

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



1047
1048
1049
 
# File 'lib/claude_agent_sdk/query.rb', line 1047

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) 


240
241
242
 
# File 'lib/claude_agent_sdk/query.rb', line 240

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

#receive_messages(&block) ⇒ Object

Receive SDK messages (not control messages)



1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
 
# File 'lib/claude_agent_sdk/query.rb', line 1174

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



1069
1070
1071
1072
1073
1074
 
# File 'lib/claude_agent_sdk/query.rb', line 1069

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.



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

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



1100
1101
1102
1103
1104
1105
 
# File 'lib/claude_agent_sdk/query.rb', line 1100

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



1060
1061
1062
1063
1064
1065
 
# File 'lib/claude_agent_sdk/query.rb', line 1060

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

#set_permission_mode(mode) ⇒ Object

Change permission mode



1052
1053
1054
1055
1056
1057
 
# File 'lib/claude_agent_sdk/query.rb', line 1052

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.



232
233
234
 
# File 'lib/claude_agent_sdk/query.rb', line 232

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:



221
222
223
224
225
226
227
228
 
# File 'lib/claude_agent_sdk/query.rb', line 221

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:



185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
 
# File 'lib/claude_agent_sdk/query.rb', line 185

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



1089
1090
1091
1092
1093
1094
 
# File 'lib/claude_agent_sdk/query.rb', line 1089

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.



1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
 
# File 'lib/claude_agent_sdk/query.rb', line 1130

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



1079
1080
1081
1082
1083
1084
1085
 
# File 'lib/claude_agent_sdk/query.rb', line 1079

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 the first 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). The condition is guaranteed to be signaled: by the result branch in read_messages, or by its ensure block when the process exits early.



1116
1117
1118
1119
1120
1121
1122
1123
 
# File 'lib/claude_agent_sdk/query.rb', line 1116

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 



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

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

#writeln(string) ⇒ Object 



1165
1166
1167
 
# File 'lib/claude_agent_sdk/query.rb', line 1165

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