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

Author: mdworsky

src/claude_agent_sdk/_internal/query.py

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

src/claude_agent_sdk/client.py

@@ -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,
@@ -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