From 37bb03836611d069d0133d8345052e72c406f17b Mon Sep 17 00:00:00 2001 From: Echolonius Date: Fri, 10 Jul 2026 21:03:20 -0500 Subject: [PATCH] fix: keep stdin open for can_use_tool permission responses wait_for_result_and_end_input() kept stdin open until the first result only when SDK MCP servers or hooks were present. A can_use_tool callback also needs the control protocol: the CLI sends can_use_tool requests mid-turn and the SDK writes the permission response back over stdin. With only can_use_tool configured (the query() + AsyncIterable path its own ValueError directs users to), stdin closed as soon as the input stream was exhausted, so every permission request failed with "Tool permission request failed: Error: Stream closed" and the callback was never invoked. Include can_use_tool in the keep-stdin-open condition, matching the existing behavior for SDK MCP servers and hooks. --- src/claude_agent_sdk/_internal/query.py | 22 ++--- tests/test_query.py | 107 ++++++++++++++++++++++++ 2 files changed, 119 insertions(+), 10 deletions(-) diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index 7aded9384..fa077b84b 100644 --- a/src/claude_agent_sdk/_internal/query.py +++ b/src/claude_agent_sdk/_internal/query.py @@ -809,18 +809,19 @@ 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. The control protocol - requires stdin to remain open for the entire conversation, so no - timeout is applied. The event is guaranteed to fire: either when the - result message arrives, or in _read_messages' finally block if the - process exits early. + If SDK MCP servers, hooks, or a can_use_tool callback require + bidirectional communication, keeps stdin open until the first result + arrives. The control protocol requires stdin to remain open for the + entire conversation, so no timeout is applied. The event is guaranteed + to fire: either when the result message arrives, or in _read_messages' + finally block if the process exits early. """ - if self.sdk_mcp_servers or self.hooks: + if self.sdk_mcp_servers or self.hooks or self.can_use_tool: logger.debug( "Waiting for first result before closing stdin " f"(sdk_mcp_servers={len(self.sdk_mcp_servers)}, " - f"has_hooks={bool(self.hooks)})" + f"has_hooks={bool(self.hooks)}, " + f"has_can_use_tool={bool(self.can_use_tool)})" ) await self._first_result_event.wait() @@ -829,8 +830,9 @@ async def wait_for_result_and_end_input(self) -> None: async def stream_input(self, stream: AsyncIterable[dict[str, Any]]) -> None: """Stream input messages to transport. - If SDK MCP servers or hooks are present, waits for the first result - before closing stdin to allow bidirectional control protocol communication. + If SDK MCP servers, hooks, or a can_use_tool callback are present, + waits for the first result before closing stdin to allow bidirectional + control protocol communication. """ try: async for message in stream: diff --git a/tests/test_query.py b/tests/test_query.py index 3ae65093d..0a317aa0b 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -309,6 +309,113 @@ async def mock_receive(): anyio.run(_test) + def test_streaming_prompt_with_can_use_tool_waits_for_result(self): + """A can_use_tool callback must keep stdin open until the first + result, exactly like SDK MCP servers and hooks: the CLI delivers + can_use_tool permission requests over the control protocol, and the + SDK writes the permission response back over stdin. + + The mock transport enforces the real CLI contract: the can_use_tool + control request arrives only after the input stream is exhausted, the + result is not produced until the permission response has been + written, and writing after end_input() fails (closed stdin). + Regression test: without waiting, end_input() fired as soon as the + input stream was exhausted and every permission response hit closed + stdin ("Tool permission request failed: Error: Stream closed"). + """ + + async def _test(): + from claude_agent_sdk import PermissionResultAllow + + permission_calls = [] + + async def on_permission(tool_name, tool_input, context): + permission_calls.append(tool_name) + return PermissionResultAllow() + + async def one_message(): + yield { + "type": "user", + "session_id": "", + "message": {"role": "user", "content": "Write a file"}, + "parent_tool_use_id": None, + } + + mock_transport = AsyncMock() + mock_transport.connect = AsyncMock() + mock_transport.close = AsyncMock() + mock_transport.is_ready = Mock(return_value=True) + + stdin_open = True + permission_response_written = anyio.Event() + writes = [] + + async def tracking_write(data): + if not stdin_open: + raise RuntimeError("stdin closed") + writes.append(data) + if "control_response" in data: + permission_response_written.set() + + async def tracking_end_input(): + nonlocal stdin_open + stdin_open = False + + mock_transport.write = tracking_write + mock_transport.end_input = tracking_end_input + + async def mock_receive(): + # The permission request arrives mid-turn, after the input + # stream is exhausted; the CLI cannot finish the turn until + # the SDK responds. + yield { + "type": "control_request", + "request_id": "perm_1", + "request": { + "subtype": "can_use_tool", + "tool_name": "Write", + "input": {"file_path": "x.txt", "content": "hi"}, + }, + } + with anyio.fail_after(5): + await permission_response_written.wait() + for msg in _ASSISTANT_AND_RESULT: + yield msg + + mock_transport.read_messages = mock_receive + + with ( + patch( + "claude_agent_sdk._internal.client.SubprocessCLITransport" + ) as mock_cls, + patch( + "claude_agent_sdk._internal.query.Query.initialize", + new_callable=AsyncMock, + ), + ): + mock_cls.return_value = mock_transport + + messages = [] + with anyio.fail_after(10): + async for msg in query( + prompt=one_message(), + options=ClaudeAgentOptions(can_use_tool=on_permission), + ): + messages.append(msg) + + assert permission_calls == ["Write"] + assert len(messages) == 2 + assert isinstance(messages[0], AssistantMessage) + assert isinstance(messages[1], ResultMessage) + + control_responses = [ + json.loads(w.rstrip("\n")) for w in writes if "control_response" in w + ] + assert len(control_responses) == 1 + assert control_responses[0]["response"]["response"]["behavior"] == "allow" + + anyio.run(_test) + def test_string_prompt_with_hooks_waits_for_result(self): """end_input() should wait for first result when hooks are configured, even without SDK MCP servers."""