@@ -309,6 +309,113 @@ async def mock_receive():
309309
310310 anyio .run (_test )
311311
312+ def test_streaming_prompt_with_can_use_tool_waits_for_result (self ):
313+ """A can_use_tool callback must keep stdin open until the first
314+ result, exactly like SDK MCP servers and hooks: the CLI delivers
315+ can_use_tool permission requests over the control protocol, and the
316+ SDK writes the permission response back over stdin.
317+
318+ The mock transport enforces the real CLI contract: the can_use_tool
319+ control request arrives only after the input stream is exhausted, the
320+ result is not produced until the permission response has been
321+ written, and writing after end_input() fails (closed stdin).
322+ Regression test: without waiting, end_input() fired as soon as the
323+ input stream was exhausted and every permission response hit closed
324+ stdin ("Tool permission request failed: Error: Stream closed").
325+ """
326+
327+ async def _test ():
328+ from claude_agent_sdk import PermissionResultAllow
329+
330+ permission_calls = []
331+
332+ async def on_permission (tool_name , tool_input , context ):
333+ permission_calls .append (tool_name )
334+ return PermissionResultAllow ()
335+
336+ async def one_message ():
337+ yield {
338+ "type" : "user" ,
339+ "session_id" : "" ,
340+ "message" : {"role" : "user" , "content" : "Write a file" },
341+ "parent_tool_use_id" : None ,
342+ }
343+
344+ mock_transport = AsyncMock ()
345+ mock_transport .connect = AsyncMock ()
346+ mock_transport .close = AsyncMock ()
347+ mock_transport .is_ready = Mock (return_value = True )
348+
349+ stdin_open = True
350+ permission_response_written = anyio .Event ()
351+ writes = []
352+
353+ async def tracking_write (data ):
354+ if not stdin_open :
355+ raise RuntimeError ("stdin closed" )
356+ writes .append (data )
357+ if "control_response" in data :
358+ permission_response_written .set ()
359+
360+ async def tracking_end_input ():
361+ nonlocal stdin_open
362+ stdin_open = False
363+
364+ mock_transport .write = tracking_write
365+ mock_transport .end_input = tracking_end_input
366+
367+ async def mock_receive ():
368+ # The permission request arrives mid-turn, after the input
369+ # stream is exhausted; the CLI cannot finish the turn until
370+ # the SDK responds.
371+ yield {
372+ "type" : "control_request" ,
373+ "request_id" : "perm_1" ,
374+ "request" : {
375+ "subtype" : "can_use_tool" ,
376+ "tool_name" : "Write" ,
377+ "input" : {"file_path" : "x.txt" , "content" : "hi" },
378+ },
379+ }
380+ with anyio .fail_after (5 ):
381+ await permission_response_written .wait ()
382+ for msg in _ASSISTANT_AND_RESULT :
383+ yield msg
384+
385+ mock_transport .read_messages = mock_receive
386+
387+ with (
388+ patch (
389+ "claude_agent_sdk._internal.client.SubprocessCLITransport"
390+ ) as mock_cls ,
391+ patch (
392+ "claude_agent_sdk._internal.query.Query.initialize" ,
393+ new_callable = AsyncMock ,
394+ ),
395+ ):
396+ mock_cls .return_value = mock_transport
397+
398+ messages = []
399+ with anyio .fail_after (10 ):
400+ async for msg in query (
401+ prompt = one_message (),
402+ options = ClaudeAgentOptions (can_use_tool = on_permission ),
403+ ):
404+ messages .append (msg )
405+
406+ assert permission_calls == ["Write" ]
407+ assert len (messages ) == 2
408+ assert isinstance (messages [0 ], AssistantMessage )
409+ assert isinstance (messages [1 ], ResultMessage )
410+
411+ control_responses = [
412+ json .loads (w .rstrip ("\n " )) for w in writes if "control_response" in w
413+ ]
414+ assert len (control_responses ) == 1
415+ assert control_responses [0 ]["response" ]["response" ]["behavior" ] == "allow"
416+
417+ anyio .run (_test )
418+
312419 def test_string_prompt_with_hooks_waits_for_result (self ):
313420 """end_input() should wait for first result when hooks are configured,
314421 even without SDK MCP servers."""
0 commit comments