@@ -1831,21 +1831,6 @@ async def __aexit__(self, *args):
18311831 return None
18321832
18331833
1834- class FakeCancelledClaudeSDKClient (FakeClaudeSDKClient ):
1835- """Simulates the real error: messages are yielded, then CancelledError on stream close.
1836-
1837- The real error comes from anyio's MemoryObjectReceiveStream when the Claude
1838- subprocess connection closes — the internal ``await receive_event.wait()``
1839- raises ``asyncio.CancelledError`` even though the asyncio Task itself is
1840- *not* being cancelled.
1841- """
1842-
1843- async def receive_response (self ):
1844- for message in self .messages :
1845- yield message
1846- raise asyncio .CancelledError
1847-
1848-
18491834def _make_fake_sdk_mcp_tool_class ():
18501835 class FakeSdkMcpTool :
18511836 def __init__ (self , name , description , input_schema , handler , ** kwargs ):
@@ -1863,85 +1848,69 @@ def _clear_tool_span_tracker() -> None:
18631848 delattr (_thread_local , "tool_span_tracker" )
18641849
18651850
1851+ @pytest .mark .skipif (not CLAUDE_SDK_AVAILABLE , reason = "Claude Agent SDK not installed" )
18661852@pytest .mark .asyncio
18671853@pytest .mark .parametrize (
1868- "messages " ,
1854+ "entrypoint " ,
18691855 [
1870- # No messages: CancelledError fires immediately on iteration.
1871- pytest .param ([], id = "empty_stream" ),
1872- # CancelledError at stream close after a complete response.
1856+ "receive_response" ,
18731857 pytest .param (
1874- [AssistantMessage (content = [TextBlock ("The answer is 42." )]), ResultMessage ()], id = "after_messages"
1858+ "query" ,
1859+ marks = pytest .mark .skipif (
1860+ getattr (claude_agent_sdk , "__version__" , None ) == "0.1.10" ,
1861+ reason = "the 0.1.10 query cassette covers the client protocol, not one-shot query()" ,
1862+ ),
18751863 ),
1876- # CancelledError between messages, before the ResultMessage arrives.
1877- pytest .param ([AssistantMessage (content = [TextBlock ("Partial answer." )])], id = "mid_stream" ),
18781864 ],
18791865)
1880- async def test_receive_response_suppresses_unexpected_cancelled_error (memory_logger , messages ):
1881- """CancelledError from the transport is suppressed; messages yielded before it are kept and logged."""
1866+ @pytest .mark .parametrize ("cancellation" , ["timeout" , "task_cancel" ])
1867+ async def test_stream_propagates_genuine_cancellation (memory_logger , entrypoint , cancellation ):
1868+ """Both public stream entrypoints propagate cancellation from a blocked real SDK read."""
18821869 assert not memory_logger .pop ()
18831870
1884- wrapped_client_class = _create_client_wrapper_class (FakeCancelledClaudeSDKClient )
1885- client = wrapped_client_class ()
1886- client ._WrappedClaudeSDKClient__client .messages = messages # type: ignore[attr-defined]
1887-
1888- await client .query ("Delegate this task." )
1889- received = []
1890- async for message in client .receive_response ():
1891- received .append (message )
1892-
1893- # All messages yielded before the CancelledError should be received.
1894- assert received == messages
1895-
1896- spans = memory_logger .pop ()
1897- task_spans = find_spans_by_type (spans , SpanTypeAttribute .TASK )
1898- assert len (task_spans ) == 1
1899- task_span = task_spans [0 ]
1900- assert task_span ["span_attributes" ]["name" ] == "Claude Agent"
1901- assert task_span .get ("error" ) is None
1902- if messages :
1903- # Output should still be logged despite the CancelledError.
1904- assert task_span ["output" ]["role" ] == "assistant"
1905- assert len (find_spans_by_type (spans , SpanTypeAttribute .LLM )) == 1
1906-
1907-
1908- @pytest .mark .asyncio
1909- async def test_genuine_task_cancel_propagates_after_receive_response (memory_logger ):
1910- """When the asyncio Task is genuinely cancelled, CancelledError propagates
1911- at the caller's next await — not swallowed forever.
1912-
1913- This verifies that suppressing CancelledError inside the generator does not
1914- permanently disarm a real ``task.cancel()`` on Python 3.12+ (where the
1915- cancellation counter is decremented when the exception is caught). On
1916- Python < 3.12 the behaviour is the same because the task cancel flag is
1917- sticky until the CancelledError propagates.
1918- """
1919- assert not memory_logger .pop ()
1871+ options = claude_agent_sdk .ClaudeAgentOptions (
1872+ model = TEST_MODEL ,
1873+ permission_mode = "bypassPermissions" ,
1874+ )
1875+ transport = make_cassette_transport (
1876+ cassette_name = "test_auto_claude_agent_sdk" ,
1877+ prompt = "" ,
1878+ options = options ,
1879+ pause_before_sdk_messages = True ,
1880+ )
19201881
1921- wrapped_client_class = _create_client_wrapper_class (FakeCancelledClaudeSDKClient )
1922- client = wrapped_client_class ()
1923- client ._WrappedClaudeSDKClient__client .messages = [ # type: ignore[attr-defined]
1924- AssistantMessage (content = [TextBlock ("Hello." )]),
1925- ResultMessage (),
1926- ]
1882+ if entrypoint == "query" :
1883+ response_stream = _create_query_wrapper_function (claude_agent_sdk .query )(
1884+ prompt = "Say hi" ,
1885+ options = options ,
1886+ transport = transport ,
1887+ )
1888+ else :
1889+ with _patched_claude_sdk (wrap_client = True ):
1890+ async with claude_agent_sdk .ClaudeSDKClient (options = options , transport = transport ) as client :
1891+ await client .query ("Say hi" )
1892+ response_stream = client .receive_response ()
1893+ await _assert_stream_cancellation (response_stream , transport , cancellation )
1894+ return
19271895
1928- await client . query ( "Hi" )
1896+ await _assert_stream_cancellation ( response_stream , transport , cancellation )
19291897
1930- async def _drain_and_sleep ():
1931- """Consume the stream, then await something else."""
1932- async for _ in client .receive_response ():
1933- pass
1934- # This is the "next await" after the generator ends.
1935- await asyncio .sleep (0 )
19361898
1937- task = asyncio .ensure_future (_drain_and_sleep ())
1938- # Let the task start and begin iterating.
1899+ async def _assert_stream_cancellation (response_stream : Any , transport : Any , cancellation : str ) -> None :
1900+ read_task = asyncio .create_task (anext (response_stream ))
1901+ await asyncio .wait_for (transport .wait_until_sdk_message_blocked (), timeout = 1 )
19391902 await asyncio .sleep (0 )
1940- # Genuinely cancel the task from outside.
1941- task .cancel ()
1903+ assert not read_task .done (), "the replayed SDK message should keep the response read blocked"
1904+
1905+ if cancellation == "timeout" :
1906+ with pytest .raises (asyncio .TimeoutError ):
1907+ await asyncio .wait_for (read_task , timeout = 0.01 )
1908+ else :
1909+ read_task .cancel ()
1910+ with pytest .raises (asyncio .CancelledError ):
1911+ await read_task
19421912
1943- with pytest .raises (asyncio .CancelledError ):
1944- await task
1913+ await response_stream .aclose ()
19451914
19461915
19471916@pytest .mark .parametrize (
0 commit comments