From 32fa7803b7b811709288434722c6066e49719bad Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 5 Oct 2026 19:10:14 -0400 Subject: [PATCH 1/2] fix(claude-agent-sdk): propagate genuine task cancellation Keep transport-originated cancellation suppression while preserving task cancellation for timeouts and explicit cancel requests. Refs SDK-497. --- .../claude_agent_sdk/test_claude_agent_sdk.py | 127 +++++++----------- .../integrations/claude_agent_sdk/tracing.py | 11 +- 2 files changed, 55 insertions(+), 83 deletions(-) diff --git a/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py b/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py index 304e5fbda..8c41fb49b 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py @@ -1831,21 +1831,6 @@ async def __aexit__(self, *args): return None -class FakeCancelledClaudeSDKClient(FakeClaudeSDKClient): - """Simulates the real error: messages are yielded, then CancelledError on stream close. - - The real error comes from anyio's MemoryObjectReceiveStream when the Claude - subprocess connection closes — the internal ``await receive_event.wait()`` - raises ``asyncio.CancelledError`` even though the asyncio Task itself is - *not* being cancelled. - """ - - async def receive_response(self): - for message in self.messages: - yield message - raise asyncio.CancelledError - - def _make_fake_sdk_mcp_tool_class(): class FakeSdkMcpTool: def __init__(self, name, description, input_schema, handler, **kwargs): @@ -1863,85 +1848,69 @@ def _clear_tool_span_tracker() -> None: delattr(_thread_local, "tool_span_tracker") +@pytest.mark.skipif(not CLAUDE_SDK_AVAILABLE, reason="Claude Agent SDK not installed") @pytest.mark.asyncio @pytest.mark.parametrize( - "messages", + "entrypoint", [ - # No messages: CancelledError fires immediately on iteration. - pytest.param([], id="empty_stream"), - # CancelledError at stream close after a complete response. + "receive_response", pytest.param( - [AssistantMessage(content=[TextBlock("The answer is 42.")]), ResultMessage()], id="after_messages" + "query", + marks=pytest.mark.skipif( + getattr(claude_agent_sdk, "__version__", None) == "0.1.10", + reason="the 0.1.10 query cassette covers the client protocol, not one-shot query()", + ), ), - # CancelledError between messages, before the ResultMessage arrives. - pytest.param([AssistantMessage(content=[TextBlock("Partial answer.")])], id="mid_stream"), ], ) -async def test_receive_response_suppresses_unexpected_cancelled_error(memory_logger, messages): - """CancelledError from the transport is suppressed; messages yielded before it are kept and logged.""" +@pytest.mark.parametrize("cancellation", ["timeout", "task_cancel"]) +async def test_stream_propagates_genuine_cancellation(memory_logger, entrypoint, cancellation): + """Both public stream entrypoints propagate cancellation from a blocked real SDK read.""" assert not memory_logger.pop() - wrapped_client_class = _create_client_wrapper_class(FakeCancelledClaudeSDKClient) - client = wrapped_client_class() - client._WrappedClaudeSDKClient__client.messages = messages # type: ignore[attr-defined] - - await client.query("Delegate this task.") - received = [] - async for message in client.receive_response(): - received.append(message) - - # All messages yielded before the CancelledError should be received. - assert received == messages - - spans = memory_logger.pop() - task_spans = find_spans_by_type(spans, SpanTypeAttribute.TASK) - assert len(task_spans) == 1 - task_span = task_spans[0] - assert task_span["span_attributes"]["name"] == "Claude Agent" - assert task_span.get("error") is None - if messages: - # Output should still be logged despite the CancelledError. - assert task_span["output"]["role"] == "assistant" - assert len(find_spans_by_type(spans, SpanTypeAttribute.LLM)) == 1 - - -@pytest.mark.asyncio -async def test_genuine_task_cancel_propagates_after_receive_response(memory_logger): - """When the asyncio Task is genuinely cancelled, CancelledError propagates - at the caller's next await — not swallowed forever. - - This verifies that suppressing CancelledError inside the generator does not - permanently disarm a real ``task.cancel()`` on Python 3.12+ (where the - cancellation counter is decremented when the exception is caught). On - Python < 3.12 the behaviour is the same because the task cancel flag is - sticky until the CancelledError propagates. - """ - assert not memory_logger.pop() + options = claude_agent_sdk.ClaudeAgentOptions( + model=TEST_MODEL, + permission_mode="bypassPermissions", + ) + transport = make_cassette_transport( + cassette_name="test_auto_claude_agent_sdk", + prompt="", + options=options, + pause_before_sdk_messages=True, + ) - wrapped_client_class = _create_client_wrapper_class(FakeCancelledClaudeSDKClient) - client = wrapped_client_class() - client._WrappedClaudeSDKClient__client.messages = [ # type: ignore[attr-defined] - AssistantMessage(content=[TextBlock("Hello.")]), - ResultMessage(), - ] + if entrypoint == "query": + response_stream = _create_query_wrapper_function(claude_agent_sdk.query)( + prompt="Say hi", + options=options, + transport=transport, + ) + else: + with _patched_claude_sdk(wrap_client=True): + async with claude_agent_sdk.ClaudeSDKClient(options=options, transport=transport) as client: + await client.query("Say hi") + response_stream = client.receive_response() + await _assert_stream_cancellation(response_stream, transport, cancellation) + return - await client.query("Hi") + await _assert_stream_cancellation(response_stream, transport, cancellation) - async def _drain_and_sleep(): - """Consume the stream, then await something else.""" - async for _ in client.receive_response(): - pass - # This is the "next await" after the generator ends. - await asyncio.sleep(0) - task = asyncio.ensure_future(_drain_and_sleep()) - # Let the task start and begin iterating. +async def _assert_stream_cancellation(response_stream: Any, transport: Any, cancellation: str) -> None: + read_task = asyncio.create_task(anext(response_stream)) + await asyncio.wait_for(transport.wait_until_sdk_message_blocked(), timeout=1) await asyncio.sleep(0) - # Genuinely cancel the task from outside. - task.cancel() + assert not read_task.done(), "the replayed SDK message should keep the response read blocked" + + if cancellation == "timeout": + with pytest.raises(TimeoutError): + await asyncio.wait_for(read_task, timeout=0.01) + else: + read_task.cancel() + with pytest.raises(asyncio.CancelledError): + await read_task - with pytest.raises(asyncio.CancelledError): - await task + await response_stream.aclose() @pytest.mark.parametrize( diff --git a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py index e42763b64..1a44e29bb 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py @@ -1310,10 +1310,13 @@ async def _stream_messages_with_tracing( except asyncio.CancelledError: # The CancelledError may come from the subprocess transport # (e.g., anyio internal cleanup when subagents complete) rather - # than a genuine external cancellation. We suppress it here so - # the response stream ends cleanly. If the caller genuinely - # cancelled the task, they still have pending cancellation - # requests that will fire at their next await point. + # than a genuine external cancellation. Suppress only the transport + # case; asyncio tracks genuine cancellation requests on the current + # task, including cancellations injected by wait_for(). + task = asyncio.current_task() + if task is not None and task.cancelling(): + finish_request_tracker(log_output=True) + raise finish_request_tracker(log_output=True) else: finish_request_tracker(log_output=True) From 5345df8edae461bda668fce7fd7723df2e98103a Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 5 Oct 2026 19:17:42 -0400 Subject: [PATCH 2/2] fix(claude-agent-sdk): support cancellation on Python 3.10 --- .../integrations/claude_agent_sdk/test_claude_agent_sdk.py | 2 +- py/src/braintrust/integrations/claude_agent_sdk/tracing.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py b/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py index 8c41fb49b..1023411e2 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py @@ -1903,7 +1903,7 @@ async def _assert_stream_cancellation(response_stream: Any, transport: Any, canc assert not read_task.done(), "the replayed SDK message should keep the response read blocked" if cancellation == "timeout": - with pytest.raises(TimeoutError): + with pytest.raises(asyncio.TimeoutError): await asyncio.wait_for(read_task, timeout=0.01) else: read_task.cancel() diff --git a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py index 1a44e29bb..ee4b19f39 100644 --- a/py/src/braintrust/integrations/claude_agent_sdk/tracing.py +++ b/py/src/braintrust/integrations/claude_agent_sdk/tracing.py @@ -1312,9 +1312,11 @@ async def _stream_messages_with_tracing( # (e.g., anyio internal cleanup when subagents complete) rather # than a genuine external cancellation. Suppress only the transport # case; asyncio tracks genuine cancellation requests on the current - # task, including cancellations injected by wait_for(). + # task, including cancellations injected by wait_for(). Python 3.10 + # has no cancellation counter, so preserve the exception there rather + # than risk hiding a genuine task cancellation. task = asyncio.current_task() - if task is not None and task.cancelling(): + if task is not None and (not hasattr(task, "cancelling") or task.cancelling()): finish_request_tracker(log_output=True) raise finish_request_tracker(log_output=True)