fix: wait for first result before closing stdin for SDK MCP servers

e2e-tests/test_sdk_mcp_tools.py

@@ -11,7 +11,9 @@
 from claude_agent_sdk import (
     ClaudeAgentOptions,
     ClaudeSDKClient,
+    ResultMessage,
     create_sdk_mcp_server,
+    query,
     tool,
 )
 
@@ -166,3 +168,48 @@ async def echo_tool(args: dict[str, Any]) -> dict[str, Any]:
             print(f"  [{type(message).__name__}] {message}")
 
     assert "echo" not in executions, "SDK MCP tool was executed"
+
+
+@pytest.mark.e2e
+@pytest.mark.asyncio
+async def test_string_prompt_with_sdk_mcp_tool():
+    """Test that string prompts work correctly with SDK MCP tools.
+
+    This test verifies the fix for: https://github.com/anthropics/claude-agent-sdk-python/issues/578
+    String prompts should work with SDK MCP tools without CLIConnectionError.
+    """
+    executions = []
+
+    @tool("echo", "Echo back the input text", {"text": str})
+    async def echo_tool(args: dict[str, Any]) -> dict[str, Any]:
+        """Echo back whatever text is provided."""
+        executions.append("echo")
+        return {"content": [{"type": "text", "text": f"Echo: {args['text']}"}]}
+
+    server = create_sdk_mcp_server(
+        name="test",
+        version="1.0.0",
+        tools=[echo_tool],
+    )
+
+    options = ClaudeAgentOptions(
+        mcp_servers={"test": server},
+        allowed_tools=["mcp__test__echo"],
+    )
+
+    # Use query() with string prompt (not ClaudeSDKClient)
+    # This is the exact scenario from issue #578
+    messages = []
+    async for msg in query(
+        prompt="Call the mcp__test__echo tool with text='hello world'",
+        options=options,
+    ):
+        messages.append(msg)
+        if isinstance(msg, ResultMessage):
+            break
+
+    # Verify the tool was executed (no CLIConnectionError)
+    assert "echo" in executions, "SDK MCP tool should have been executed"
+
+    # Verify we got a result message
+    assert any(isinstance(msg, ResultMessage) for msg in messages)

src/claude_agent_sdk/_internal/client.py

@@ -131,6 +131,29 @@ async def process_query(
                     "parent_tool_use_id": None,
                 }
                 await chosen_transport.write(json.dumps(user_message) + "\n")
+
+                # Keep stdin open until CLI completes MCP initialization
+                # This ensures bidirectional control protocol works for SDK MCP servers and hooks
+                if sdk_mcp_servers or configured_options.hooks:
+                    import logging
+
+                    import anyio
+
+                    logger = logging.getLogger(__name__)
+
+                    logger.debug(
+                        f"Waiting for first result before closing stdin "
+                        f"(sdk_mcp_servers={len(sdk_mcp_servers)}, has_hooks={bool(configured_options.hooks)})"
+                    )
+                    try:
+                        with anyio.move_on_after(query._stream_close_timeout):
+                            await query._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"
+                        )
+
                 await chosen_transport.end_input()
             elif isinstance(prompt, AsyncIterable) and query._tg:
                 # Stream input in background for async iterables

tests/test_sdk_mcp_integration.py

@@ -377,5 +377,83 @@ async def plain_tool(args: dict[str, Any]) -> dict[str, Any]:
     assert tools_by_name["read_only_tool"]["annotations"]["readOnlyHint"] is True
     assert tools_by_name["read_only_tool"]["annotations"]["openWorldHint"] is False
 
-    # Tool without annotations should not have the key
-    assert "annotations" not in tools_by_name["plain_tool"]
+
+@pytest.mark.asyncio
+async def test_string_prompt_with_sdk_mcp_servers():
+    """Test that string prompts work correctly with SDK MCP servers.
+
+    This test verifies the fix for: https://github.com/anthropics/claude-agent-sdk-python/issues/578
+    String prompts should wait for first result before closing stdin to allow
+    bidirectional control protocol communication for SDK MCP servers.
+    """
+    from unittest.mock import AsyncMock, MagicMock, patch
+
+    from claude_agent_sdk._internal.client import InternalClient
+
+    # Track tool executions
+    tool_executions = []
+
+    @tool("test_tool", "A test tool", {"input": str})
+    async def test_tool(args: dict[str, Any]) -> dict[str, Any]:
+        tool_executions.append({"name": "test_tool", "args": args})
+        return {"content": [{"type": "text", "text": f"Result: {args['input']}"}]}
+
+    server = create_sdk_mcp_server(name="test", tools=[test_tool])
+
+    # Create mock transport
+    mock_transport = MagicMock()
+    mock_transport.write = AsyncMock()
+    mock_transport.end_input = AsyncMock()
+    mock_transport.connect = AsyncMock()
+
+    # Create mock query with _first_result_event
+    mock_query = MagicMock()
+    mock_first_result_event = MagicMock()
+    mock_first_result_event.wait = AsyncMock()
+    mock_query._first_result_event = mock_first_result_event
+    mock_query._stream_close_timeout = 60.0
+    mock_query.start = AsyncMock()
+    mock_query.initialize = AsyncMock()
+    mock_query.close = AsyncMock()
+
+    # Mock receive_messages to yield a result (which should trigger _first_result_event)
+    async def mock_receive():
+        yield {
+            "type": "result",
+            "subtype": "success",
+            "duration_ms": 100,
+            "duration_api_ms": 50,
+            "is_error": False,
+            "num_turns": 1,
+            "session_id": "test-session",
+        }
+
+    mock_query.receive_messages = mock_receive
+
+    # Patch Query creation to return our mock
+    with patch("claude_agent_sdk._internal.client.Query", return_value=mock_query):
+        client = InternalClient()
+        options = ClaudeAgentOptions(mcp_servers={"test": server})
+
+        # Execute query with string prompt
+        messages = []
+        async for msg in client.process_query(
+            prompt="Test prompt",
+            options=options,
+            transport=mock_transport,
+        ):
+            messages.append(msg)
+
+    # Verify user message was written
+    assert mock_transport.write.called
+    written_data = mock_transport.write.call_args[0][0]
+    assert "Test prompt" in written_data
+
+    # Verify end_input was called
+    assert mock_transport.end_input.called
+
+    # Verify _first_result_event.wait() was called (the fix)
+    assert mock_first_result_event.wait.called, (
+        "String prompt with SDK MCP servers should wait for first result "
+        "before closing stdin (issue #578 fix verification)"
+    )