Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 17 additions & 4 deletions src/claude_agent_sdk/_internal/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ def __init__(
| None = None,
hooks: dict[str, list[dict[str, Any]]] | None = None,
sdk_mcp_servers: dict[str, "McpServer"] | None = None,
initialize_timeout: float = 60.0,
):
"""Initialize Query with transport and callbacks.

Expand All @@ -81,7 +82,9 @@ def __init__(
can_use_tool: Optional callback for tool permission requests
hooks: Optional hook configurations
sdk_mcp_servers: Optional SDK MCP server instances
initialize_timeout: Timeout in seconds for the initialize request
"""
self._initialize_timeout = initialize_timeout
self.transport = transport
self.is_streaming_mode = is_streaming_mode
self.can_use_tool = can_use_tool
Expand Down Expand Up @@ -140,7 +143,10 @@ async def initialize(self) -> dict[str, Any] | None:
"hooks": hooks_config if hooks_config else None,
}

response = await self._send_control_request(request)
# Use longer timeout for initialize since MCP servers may take time to start
response = await self._send_control_request(
request, timeout=self._initialize_timeout
)
self._initialized = True
self._initialization_result = response # Store for later access
return response
Expand Down Expand Up @@ -315,8 +321,15 @@ async def _handle_control_request(self, request: SDKControlRequest) -> None:
}
await self.transport.write(json.dumps(error_response) + "\n")

async def _send_control_request(self, request: dict[str, Any]) -> dict[str, Any]:
"""Send control request to CLI and wait for response."""
async def _send_control_request(
self, request: dict[str, Any], timeout: float = 60.0
) -> dict[str, Any]:
"""Send control request to CLI and wait for response.

Args:
request: The control request to send
timeout: Timeout in seconds to wait for response (default 60s)
"""
if not self.is_streaming_mode:
raise Exception("Control requests require streaming mode")

Expand All @@ -339,7 +352,7 @@ async def _send_control_request(self, request: dict[str, Any]) -> dict[str, Any]

# Wait for response
try:
with anyio.fail_after(60.0):
with anyio.fail_after(timeout):
await event.wait()

result = self.pending_control_results.pop(request_id)
Expand Down
8 changes: 8 additions & 0 deletions src/claude_agent_sdk/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,13 @@ async def _empty_stream() -> AsyncIterator[dict[str, Any]]:
if isinstance(config, dict) and config.get("type") == "sdk":
sdk_mcp_servers[name] = config["instance"] # type: ignore[typeddict-item]

# Calculate initialize timeout from CLAUDE_CODE_STREAM_CLOSE_TIMEOUT env var if set
# CLAUDE_CODE_STREAM_CLOSE_TIMEOUT is in milliseconds, convert to seconds
initialize_timeout_ms = int(
os.environ.get("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "60000")
)
initialize_timeout = max(initialize_timeout_ms / 1000.0, 60.0)

# Create Query to handle control protocol
self._query = Query(
transport=self._transport,
Expand All @@ -149,6 +156,7 @@ async def _empty_stream() -> AsyncIterator[dict[str, Any]]:
if self.options.hooks
else None,
sdk_mcp_servers=sdk_mcp_servers,
initialize_timeout=initialize_timeout,
)

# Start reading messages and initialize
Expand Down