fix: Resolve RuntimeError on async generator cleanup

This fixes the "Attempted to exit cancel scope in a different task than
it was entered in" error that occurs when sequential query() calls are made.

The fix uses two mechanisms:
1. A CancelScope for the reader task that can be cancelled from any context
2. Suppressing the RuntimeError that occurs when task group __aexit__()
   is called from a different task than __aenter__()

Also adds aclosing() wrapper and GeneratorExit handling in process_query()
for proper async generator cleanup.

Based on PR anthropics#527: anthropics#527

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: brianantonelli

src/claude_agent_sdk/_internal/client.py

@@ -1,6 +1,8 @@
 """Internal client implementation."""
 
+import logging
 from collections.abc import AsyncIterable, AsyncIterator
+from contextlib import aclosing
 from dataclasses import replace
 from typing import Any
 
@@ -15,6 +17,8 @@
 from .transport import Transport
 from .transport.subprocess_cli import SubprocessCLITransport
 
+logger = logging.getLogger(__name__)
+
 
 class InternalClient:
     """Internal client implementation."""
@@ -117,8 +121,14 @@ async def process_query(
             # For string prompts, the prompt is already passed via CLI args
 
             # Yield parsed messages
-            async for data in query.receive_messages():
-                yield parse_message(data)
-
+            # Use aclosing() for proper async generator cleanup
+            async with aclosing(query.receive_messages()) as messages:
+                async for data in messages:
+                    yield parse_message(data)
+
+        except GeneratorExit:
+            # Handle early termination of the async generator gracefully
+            # This occurs when the caller breaks out of the async for loop
+            logger.debug("process_query generator closed early by caller")
         finally:
             await query.close()

src/claude_agent_sdk/_internal/query.py

@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING, Any
 
 import anyio
+from anyio.abc import CancelScope
 from mcp.types import (
     CallToolRequest,
     CallToolRequestParams,
@@ -113,6 +114,15 @@ def __init__(
             float(os.environ.get("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "60000")) / 1000.0
         )  # Convert ms to seconds
 
+        # Cancel scope for the reader task - can be cancelled from any task context
+        # This fixes the RuntimeError when async generator cleanup happens in a different task
+        self._reader_cancel_scope: CancelScope | None = None
+        self._reader_task_started = anyio.Event()
+
+        # Track whether we entered the task group in this task
+        # Used to determine if we can safely call __aexit__()
+        self._tg_entered_in_current_task = False
+
     async def initialize(self) -> dict[str, Any] | None:
         """Initialize control protocol if in streaming mode.
 
@@ -158,11 +168,33 @@ async def initialize(self) -> dict[str, Any] | None:
         return response
 
     async def start(self) -> None:
-        """Start reading messages from transport."""
+        """Start reading messages from transport.
+
+        This method starts background tasks for reading messages. The task lifecycle
+        is managed using a CancelScope that can be safely cancelled from any async
+        task context, avoiding the RuntimeError that occurs when task group
+        __aexit__() is called from a different task than __aenter__().
+        """
         if self._tg is None:
+            # Create a task group for spawning background tasks
             self._tg = anyio.create_task_group()
             await self._tg.__aenter__()
-            self._tg.start_soon(self._read_messages)
+            self._tg_entered_in_current_task = True
+
+            # Start the reader with its own cancel scope that can be cancelled safely
+            self._tg.start_soon(self._read_messages_with_cancel_scope)
+
+    async def _read_messages_with_cancel_scope(self) -> None:
+        """Wrapper for _read_messages that sets up a cancellable scope.
+
+        This wrapper creates a CancelScope that can be cancelled from any task
+        context, solving the issue where async generator cleanup happens in a
+        different task than where the task group was entered.
+        """
+        self._reader_cancel_scope = anyio.CancelScope()
+        self._reader_task_started.set()
+        with self._reader_cancel_scope:
+            await self._read_messages()
 
     async def _read_messages(self) -> None:
         """Read messages from transport and route them."""
@@ -604,15 +636,66 @@ async def receive_messages(self) -> AsyncIterator[dict[str, Any]]:
             yield message
 
     async def close(self) -> None:
-        """Close the query and transport."""
+        """Close the query and transport.
+
+        This method safely cleans up resources, handling the case where cleanup
+        happens in a different async task context than where start() was called.
+        This commonly occurs during async generator cleanup (e.g., when breaking
+        out of an `async for` loop or when asyncio.run() shuts down).
+
+        The fix uses two mechanisms:
+        1. A CancelScope for the reader task that can be cancelled from any context
+        2. Suppressing the RuntimeError that occurs when task group __aexit__()
+           is called from a different task than __aenter__()
+        """
+        if self._closed:
+            return
         self._closed = True
-        if self._tg:
+
+        # Cancel the reader task via its cancel scope (safe from any task context)
+        if self._reader_cancel_scope is not None:
+            self._reader_cancel_scope.cancel()
+
+        # Handle task group cleanup
+        if self._tg is not None:
+            # Always cancel the task group's scope to stop any running tasks
             self._tg.cancel_scope.cancel()
-            # Wait for task group to complete cancellation
-            with suppress(anyio.get_cancelled_exc_class()):
-                await self._tg.__aexit__(None, None, None)
+
+            # Try to properly exit the task group, but handle the case where
+            # we're in a different task context than where __aenter__() was called
+            try:
+                with suppress(anyio.get_cancelled_exc_class()):
+                    await self._tg.__aexit__(None, None, None)
+            except RuntimeError as e:
+                # Handle "Attempted to exit cancel scope in a different task"
+                # This happens during async generator cleanup when Python's GC
+                # runs the finally block in a different task context.
+                if "different task" in str(e):
+                    logger.debug(
+                        "Task group cleanup skipped due to cross-task context "
+                        "(this is expected during async generator cleanup)"
+                    )
+                else:
+                    raise
+            finally:
+                self._tg = None
+                self._tg_entered_in_current_task = False
+
         await self.transport.close()
 
+    # Make Query an async context manager
+    async def __aenter__(self) -> "Query":
+        """Enter async context - starts reading messages."""
+        await self.start()
+        return self
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
+    ) -> bool:
+        """Exit async context - closes the query."""
+        await self.close()
+        return False
+
     # Make Query an async iterator
     def __aiter__(self) -> AsyncIterator[dict[str, Any]]:
         """Return async iterator for messages."""