Fix string prompt closing stdin before MCP server init completes (#630)

Author: omid-ant

src/claude_agent_sdk/_internal/client.py

@@ -1,5 +1,6 @@
 """Internal client implementation."""
 
+import json
 from collections.abc import AsyncIterable, AsyncIterator
 from dataclasses import asdict, replace
 from typing import Any
@@ -122,16 +123,14 @@ async def process_query(
             if isinstance(prompt, str):
                 # For string prompts, write user message to stdin after initialize
                 # (matching TypeScript SDK behavior)
-                import json
-
                 user_message = {
                     "type": "user",
                     "session_id": "",
                     "message": {"role": "user", "content": prompt},
                     "parent_tool_use_id": None,
                 }
                 await chosen_transport.write(json.dumps(user_message) + "\n")
-                await chosen_transport.end_input()
+                await query.wait_for_result_and_end_input()
             elif isinstance(prompt, AsyncIterable) and query._tg:
                 # Stream input in background for async iterables
                 query._tg.start_soon(query.stream_input, prompt)

src/claude_agent_sdk/_internal/query.py

@@ -227,6 +227,9 @@ async def _read_messages(self) -> None:
             # Put error in stream so iterators can handle it
             await self._message_send.send({"type": "error", "error": str(e)})
         finally:
+            # Unblock any waiters (e.g. string-prompt path waiting for first
+            # result) so they don't stall for the full timeout on early exit.
+            self._first_result_event.set()
             # Always signal end of stream
             await self._message_send.send({"type": "end"})
 
@@ -608,6 +611,24 @@ async def stop_task(self, task_id: str) -> None:
             }
         )
 
+    async def wait_for_result_and_end_input(self) -> None:
+        """Wait for the first result (if needed) then close stdin.
+
+        If SDK MCP servers or hooks require bidirectional communication,
+        keeps stdin open until the first result arrives (or timeout).
+        Otherwise closes stdin immediately.
+        """
+        if self.sdk_mcp_servers or self.hooks:
+            logger.debug(
+                "Waiting for first result before closing stdin "
+                f"(sdk_mcp_servers={len(self.sdk_mcp_servers)}, "
+                f"has_hooks={bool(self.hooks)})"
+            )
+            with anyio.move_on_after(self._stream_close_timeout):
+                await self._first_result_event.wait()
+
+        await self.transport.end_input()
+
     async def stream_input(self, stream: AsyncIterable[dict[str, Any]]) -> None:
         """Stream input messages to transport.
 
@@ -620,25 +641,7 @@ async def stream_input(self, stream: AsyncIterable[dict[str, Any]]) -> None:
                     break
                 await self.transport.write(json.dumps(message) + "\n")
 
-            # If we have SDK MCP servers or hooks that need bidirectional communication,
-            # wait for first result before closing the channel
-            has_hooks = bool(self.hooks)
-            if self.sdk_mcp_servers or has_hooks:
-                logger.debug(
-                    f"Waiting for first result before closing stdin "
-                    f"(sdk_mcp_servers={len(self.sdk_mcp_servers)}, has_hooks={has_hooks})"
-                )
-                try:
-                    with anyio.move_on_after(self._stream_close_timeout):
-                        await self._first_result_event.wait()
-                        logger.debug("Received first result, closing input stream")
-                except Exception:
-                    logger.debug(
-                        "Timed out waiting for first result, closing input stream"
-                    )
-
-            # After all messages sent (and result received if needed), end input
-            await self.transport.end_input()
+            await self.wait_for_result_and_end_input()
         except Exception as e:
             logger.debug(f"Error streaming input: {e}")