diff --git a/e2e-tests/test_sdk_mcp_tools.py b/e2e-tests/test_sdk_mcp_tools.py index a5247fa3..42277808 100644 --- a/e2e-tests/test_sdk_mcp_tools.py +++ b/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) diff --git a/src/claude_agent_sdk/_internal/client.py b/src/claude_agent_sdk/_internal/client.py index 90f535fb..0befc913 100644 --- a/src/claude_agent_sdk/_internal/client.py +++ b/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 diff --git a/tests/test_sdk_mcp_integration.py b/tests/test_sdk_mcp_integration.py index 67f9e66e..9c7bd96a 100644 --- a/tests/test_sdk_mcp_integration.py +++ b/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)" + )