Skip to content

Commit 09b03e9

Browse files
authored
fix(claude-agent-sdk): propagate genuine task cancellation (#882)
## Summary - Propagate `asyncio.CancelledError` when the current task has a pending cancellation request, while keeping suppression for transport-originated cancellation where asyncio exposes the task cancellation counter. - On Python 3.10, where asyncio has no cancellation counter, preserve `CancelledError` rather than risk masking genuine task cancellation. - Add cassette-backed timeout and explicit cancellation coverage for both `query()` and `receive_response()`. - Remove fake-only cancellation tests in favor of replay through the real Claude Agent SDK transport. ## Validation Python 3.12 and 3.10, using both supported Claude Agent SDK pins: - `latest`: 50 passed - `0.1.10`: 43 passed, 7 skipped Refs SDK-497.
1 parent 964b6eb commit 09b03e9

2 files changed

Lines changed: 57 additions & 83 deletions

File tree

‎py/src/braintrust/integrations/claude_agent_sdk/test_claude_agent_sdk.py‎

Lines changed: 48 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -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-
18491834
def _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(

‎py/src/braintrust/integrations/claude_agent_sdk/tracing.py‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1310,10 +1310,15 @@ async def _stream_messages_with_tracing(
13101310
except asyncio.CancelledError:
13111311
# The CancelledError may come from the subprocess transport
13121312
# (e.g., anyio internal cleanup when subagents complete) rather
1313-
# than a genuine external cancellation. We suppress it here so
1314-
# the response stream ends cleanly. If the caller genuinely
1315-
# cancelled the task, they still have pending cancellation
1316-
# requests that will fire at their next await point.
1313+
# than a genuine external cancellation. Suppress only the transport
1314+
# case; asyncio tracks genuine cancellation requests on the current
1315+
# task, including cancellations injected by wait_for(). Python 3.10
1316+
# has no cancellation counter, so preserve the exception there rather
1317+
# than risk hiding a genuine task cancellation.
1318+
task = asyncio.current_task()
1319+
if task is not None and (not hasattr(task, "cancelling") or task.cancelling()):
1320+
finish_request_tracker(log_output=True)
1321+
raise
13171322
finish_request_tracker(log_output=True)
13181323
else:
13191324
finish_request_tracker(log_output=True)

0 commit comments

Comments
 (0)