Use CLAUDE_CODE_STREAM_CLOSE_TIMEOUT (if present) to override initialize() timeout (#354)

This commit is contained in:
Michael Dworsky 2025-11-21 16:46:22 -05:00 committed by GitHub
parent f446e3e42a
commit 41ceacd807
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 25 additions and 4 deletions

View file

@ -72,6 +72,7 @@ class Query:
| 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.
@ -81,7 +82,9 @@ class Query:
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
@ -140,7 +143,10 @@ class Query:
"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
@ -315,8 +321,15 @@ class Query:
}
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")
@ -339,7 +352,7 @@ class Query:
# 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)

View file

@ -140,6 +140,13 @@ class ClaudeSDKClient:
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,
@ -149,6 +156,7 @@ class ClaudeSDKClient:
if self.options.hooks
else None,
sdk_mcp_servers=sdk_mcp_servers,
initialize_timeout=initialize_timeout,
)
# Start reading messages and initialize