fix: resolve cross-task cancel scope RuntimeError on async generator cleanup (#454)

[qing-ant]

Problem

When users break out of the async for loop over query(), Python may finalize the async generator in a different task than the one that created the task group. This causes close() to call TaskGroup.__aexit__() from a different task than start() called __aenter__(), triggering:

RuntimeError: Attempted to exit cancel scope in a different task than it was entered in

Fixes #454.

Root cause

The Query class was using anyio's TaskGroup with manual __aenter__/__aexit__ calls. anyio's cancel scopes have task affinity — they must be exited by the same async task that entered them. During async generator finalization, Python can schedule the generator's cleanup in a different task, violating this invariant.

Why PR #364 doesn't fix it

PR #364 introduces an "owner task pattern" that wraps the inner task group in a dedicated owner task. However, it still creates an outer task group (_outer_tg) using the same manual __aenter__/__aexit__ pattern, so the cross-task error just moves one level up. The tests in that PR call start() and close() from the same task, so they don't reproduce the actual failure scenario.

Solution

Replace anyio TaskGroup with asyncio.create_task() for background task management. asyncio.create_task() has no cancel scope, so close() can cancel tasks from any task context without triggering the RuntimeError.

Changes:

• query.py: Replace _tg (anyio TaskGroup) with _read_task (asyncio Task) and _child_tasks (set of asyncio Tasks). Add spawn_task() method as the replacement for _tg.start_soon().
• client.py / _internal/client.py: Update callers to use spawn_task() instead of _tg.start_soon().
• test_query.py: Add tests that reproduce the cross-task cleanup scenario.

Test plan

• All 356 existing tests pass
• New test test_close_from_different_task_does_not_raise verifies cross-task cleanup works
• New test test_close_from_same_task_still_works verifies normal cleanup still works
• Linting (ruff) and type checking (mypy) pass

[qing-ant]

E2E Test Results

Test Script

"""E2E test for PR #746 / Issue #454: cross-task cancel scope RuntimeError on async generator cleanup.

When breaking out of `async for` over query(), Python may finalize the async
generator in a different task than the one that entered the anyio cancel scope,
causing:
  RuntimeError: Attempted to exit cancel scope in a different task than it was entered in

This test breaks early from the async generator and checks stderr for the error.
"""

import asyncio
import sys
import io
import logging
import warnings


async def run_test():
    from claude_agent_sdk import query, ClaudeAgentOptions, ResultMessage

    # Capture warnings and stderr to detect the RuntimeError
    stderr_capture = io.StringIO()
    handler = logging.StreamHandler(stderr_capture)
    handler.setLevel(logging.DEBUG)
    logging.getLogger().addHandler(handler)

    old_stderr = sys.stderr
    sys.stderr = io.TextIOWrapper(io.BytesIO(), write_through=True)

    error_detected = False

    try:
        messages = []
        async for msg in query(
            prompt="Say hello in exactly 3 words",
            options=ClaudeAgentOptions(model="claude-sonnet-4-20250514"),
        ):
            messages.append(msg)
            if isinstance(msg, ResultMessage):
                result_preview = (msg.result or "")[:80]
                print(f"Got ResultMessage, breaking early. Result: {result_preview}...")
                break

        # Give the event loop a chance to process any pending callbacks/finalizers
        # that would trigger the cross-task cancel scope error
        await asyncio.sleep(0.5)

        # Force garbage collection to trigger async generator finalization
        import gc
        gc.collect()
        await asyncio.sleep(0.5)

    except Exception as e:
        print(f"Exception during query: {type(e).__name__}: {e}")
        error_detected = True

    finally:
        # Restore stderr and check for errors
        captured_stderr_bytes = sys.stderr.buffer.getvalue() if hasattr(sys.stderr, 'buffer') else b""
        sys.stderr = old_stderr
        captured_stderr = captured_stderr_bytes.decode("utf-8", errors="replace")
        captured_logs = stderr_capture.getvalue()

    # Check for the specific RuntimeError in captured output
    all_output = captured_stderr + captured_logs
    if "cancel scope" in all_output.lower() or "RuntimeError" in all_output:
        error_detected = True
        print(f"\n--- CAPTURED STDERR/LOGS ---")
        print(all_output.strip())
        print(f"--- END CAPTURED ---")

    return error_detected, messages


def main():
    # Also install a custom exception handler to catch "Task exception was never retrieved"
    exceptions_found = []

    def custom_exception_handler(loop, context):
        msg = context.get("message", "")
        exc = context.get("exception", None)
        detail = f"{msg}: {exc}" if exc else msg
        exceptions_found.append(detail)
        print(f"[Exception handler] {detail}", file=sys.__stderr__)

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.set_exception_handler(custom_exception_handler)

    try:
        error_detected, messages = loop.run_until_complete(run_test())
        # Give time for any deferred task exceptions
        loop.run_until_complete(asyncio.sleep(1.0))
    finally:
        # Run pending callbacks
        loop.run_until_complete(asyncio.sleep(0.2))
        loop.close()

    # Check for cross-task cancel scope errors in the exception handler output
    cancel_scope_errors = [e for e in exceptions_found if "cancel scope" in e.lower() or "RuntimeError" in e]
    any_task_exceptions = len(exceptions_found) > 0

    print(f"\nMessages received before break: {len(messages)}")
    print(f"Exception handler caught {len(exceptions_found)} exception(s)")
    for e in exceptions_found:
        print(f"  - {e}")

    if error_detected or cancel_scope_errors:
        print("\n>>> FAIL: RuntimeError about cancel scope detected <<<")
        sys.exit(1)
    elif any_task_exceptions:
        print(f"\n>>> FAIL: Task exceptions detected (not cancel scope, but still errors) <<<")
        sys.exit(1)
    else:
        print("\n>>> PASS: No cross-task cancel scope errors <<<")
        sys.exit(0)


if __name__ == "__main__":
    main()

Regression Test (main branch) - FAIL as expected

SDK installed from origin/main (76cb292). Breaking early from the async generator triggers the cross-task cancel scope RuntimeError:

[Exception handler] Task exception was never retrieved: Attempted to exit cancel scope in a different task than it was entered in
Got ResultMessage, breaking early. Result: Hello there friend!...
Traceback (most recent call last):
  File "/tmp/e2e-746-test.py", line 120, in <module>
    main()
  File "/tmp/e2e-746-test.py", line 91, in main
    error_detected, messages = loop.run_until_complete(run_test())
  File "asyncio/base_events.py", line 725, in run_until_complete
    return future.result()
  File "/tmp/e2e-746-test.py", line 46, in run_test
    await asyncio.sleep(0.5)
  File "asyncio/tasks.py", line 718, in sleep
    return await future
asyncio.exceptions.CancelledError: Cancelled by cancel scope 7f9c466835c0

Exit code: 1 (FAIL)

Fixed Branch Test (qing/fix-454-cross-task-cancel-scope) - PASS

SDK installed from the fix branch (b7dddce). Breaking early from the async generator works cleanly with no errors:

Got ResultMessage, breaking early. Result: Hello there friend!...

Messages received before break: 5
Exception handler caught 0 exception(s)

>>> PASS: No cross-task cancel scope errors <<<

Exit code: 0 (PASS)

Verdict

PASS - The fix correctly resolves the cross-task cancel scope RuntimeError. Replacing the anyio TaskGroup with direct asyncio.Task management eliminates the cancel scope that caused the cross-task finalization error when Python's async generator cleanup runs in a different task.

[bogini]

stamped 🐎

[tjni]

Hi @qing-ant, is it possible to resolve this in a way that continues to use anyio in order to support trio?