feat: add async context manager support for CopilotClient and Copilot…

python/README.md

@@ -23,6 +23,45 @@ python chat.py
 
 ## Quick Start
 
+### Using Context Managers (Recommended)
+
+The SDK supports Python's async context manager protocol for automatic resource cleanup:
+
+```python
+import asyncio
+from copilot import CopilotClient
+
+async def main():
+    # Client automatically starts on enter and cleans up on exit
+    async with CopilotClient() as client:
+        # Create a session with automatic cleanup
+        async with await client.create_session({"model": "gpt-4"}) as session:
+            # Wait for response using session.idle event
+            done = asyncio.Event()
+
+            def on_event(event):
+                if event.type.value == "assistant.message":
+                    print(event.data.content)
+                elif event.type.value == "session.idle":
+                    done.set()
+
+            session.on(on_event)
+
+            # Send a message and wait for completion
+            await session.send({"prompt": "What is 2+2?"})
+            await done.wait()
+
+        # Session automatically destroyed here
+
+    # Client automatically stopped here
+
+asyncio.run(main())
+```
+
+### Manual Resource Management
+
+You can also manage resources manually:
+
 ```python
 import asyncio
 from copilot import CopilotClient
@@ -73,6 +112,7 @@ async with await client.create_session({"model": "gpt-5"}) as session:
 - ✅ Session history with `get_messages()`
 - ✅ Type hints throughout
 - ✅ Async/await native
+- ✅ Async context manager support for automatic resource cleanup
 
 ## API Reference
 
@@ -157,6 +197,44 @@ unsubscribe()
 - `session.foreground` - A session became the foreground session in TUI
 - `session.background` - A session is no longer the foreground session
 
+### Context Manager Support
+
+Both `CopilotClient` and `CopilotSession` support Python's async context manager protocol for automatic resource cleanup. This is the recommended pattern as it ensures resources are properly cleaned up even if exceptions occur.
+
+**CopilotClient Context Manager:**
+
+```python
+async with CopilotClient() as client:
+    # Client automatically starts on enter
+    session = await client.create_session()
+    await session.send({"prompt": "Hello!"})
+    # Client automatically stops on exit, cleaning up all sessions
+```
+
+**CopilotSession Context Manager:**
+
+```python
+async with await client.create_session() as session:
+    await session.send({"prompt": "Hello!"})
+    # Session automatically destroyed on exit
+```
+
+**Nested Context Managers:**
+
+```python
+async with CopilotClient() as client:
+    async with await client.create_session() as session:
+        await session.send({"prompt": "Hello!"})
+    # Session destroyed here
+# Client stopped here
+```
+
+**Benefits:**
+- Prevents resource leaks by ensuring cleanup even if exceptions occur
+- Eliminates the need to manually call `stop()` and `destroy()`
+- Follows Python best practices for resource management
+- Particularly useful in batch operations and evaluations to prevent process accumulation
+
 ### Tools
 
 Define tools with automatic JSON schema generation using the `@define_tool` decorator and Pydantic models:

python/copilot/client.py

@@ -22,7 +22,8 @@
 from collections.abc import Callable
 from dataclasses import asdict, is_dataclass
 from pathlib import Path
-from typing import Any, cast
+from types import TracebackType
+from typing import Any, Optional, cast
 
 from .generated.rpc import ServerRpc
 from .generated.session_events import session_event_from_dict
@@ -212,6 +213,38 @@ def __init__(self, options: CopilotClientOptions | None = None):
         self._lifecycle_handlers_lock = threading.Lock()
         self._rpc: ServerRpc | None = None
 
+    async def __aenter__(self) -> "CopilotClient":
+        """
+        Enter the async context manager.
+
+        Automatically starts the CLI server and establishes a connection if not
+        already connected.
+
+        Returns:
+            The CopilotClient instance.
+
+        Example:
+            >>> async with CopilotClient() as client:
+            ...     session = await client.create_session()
+            ...     await session.send({"prompt": "Hello!"})
+        """
+        await self.start()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """
+        Exit the async context manager.
+
+        Performs graceful cleanup by destroying all active sessions and stopping
+        the CLI server.
+        """
+        await self.stop()
+
     @property
     def rpc(self) -> ServerRpc:
         """Typed server-scoped RPC methods."""

python/copilot/session.py

@@ -9,7 +9,8 @@
 import inspect
 import threading
 from collections.abc import Callable
-from typing import Any, cast
+from types import TracebackType
+from typing import Any, Optional, cast
 
 from .generated.rpc import SessionModelSwitchToParams, SessionRpc
 from .generated.session_events import SessionEvent, SessionEventType, session_event_from_dict
@@ -73,6 +74,7 @@ def __init__(self, session_id: str, client: Any, workspace_path: str | None = No
         self.session_id = session_id
         self._client = client
         self._workspace_path = workspace_path
+        self._destroyed = False
         self._event_handlers: set[Callable[[SessionEvent], None]] = set()
         self._event_handlers_lock = threading.Lock()
         self._tool_handlers: dict[str, ToolHandler] = {}
@@ -85,6 +87,35 @@ def __init__(self, session_id: str, client: Any, workspace_path: str | None = No
         self._hooks_lock = threading.Lock()
         self._rpc: SessionRpc | None = None
 
+    async def __aenter__(self) -> "CopilotSession":
+        """
+        Enter the async context manager.
+
+        Returns the session instance, ready for use. The session must already be
+        created (via CopilotClient.create_session or resume_session).
+
+        Returns:
+            The CopilotSession instance.
+
+        Example:
+            >>> async with await client.create_session() as session:
+            ...     await session.send({"prompt": "Hello!"})
+        """
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """
+        Exit the async context manager.
+
+        Automatically destroys the session and releases all associated resources.
+        """
+        await self.destroy()
+
     @property
     def rpc(self) -> SessionRpc:
         """Typed session-scoped RPC methods."""
@@ -487,20 +518,33 @@ async def disconnect(self) -> None:
 
         After calling this method, the session object can no longer be used.
 
+        This method is idempotent—calling it multiple times is safe and will
+        not raise an error if the session is already destroyed.
+
         Raises:
-            Exception: If the connection fails.
+            Exception: If the connection fails (on first destroy call).
 
         Example:
             >>> # Clean up when done — session can still be resumed later
             >>> await session.disconnect()
         """
-        await self._client.request("session.destroy", {"sessionId": self.session_id})
+        # Ensure that the check and update of _destroyed are atomic so that
+        # only the first caller proceeds to send the destroy RPC.
         with self._event_handlers_lock:
-            self._event_handlers.clear()
-        with self._tool_handlers_lock:
-            self._tool_handlers.clear()
-        with self._permission_handler_lock:
-            self._permission_handler = None
+            if self._destroyed:
+                return
+            self._destroyed = True
+
+        try:
+            await self._client.request("session.destroy", {"sessionId": self.session_id})
+        finally:
+            # Clear handlers even if the request fails
+            with self._event_handlers_lock:
+                self._event_handlers.clear()
+            with self._tool_handlers_lock:
+                self._tool_handlers.clear()
+            with self._permission_handler_lock:
+                self._permission_handler = None
 
     async def destroy(self) -> None:
         """

python/test_client.py

@@ -316,3 +316,49 @@ async def mock_request(method, params):
             assert captured["session.model.switchTo"]["modelId"] == "gpt-4.1"
         finally:
             await client.force_stop()
+
+
+class TestContextManager:
+    @pytest.mark.asyncio
+    async def test_client_context_manager_returns_self(self):
+        """Test that __aenter__ returns the client instance."""
+        client = CopilotClient({"cli_path": CLI_PATH})
+        returned_client = await client.__aenter__()
+        assert returned_client is client
+        await client.force_stop()
+
+    @pytest.mark.asyncio
+    async def test_client_aexit_returns_none(self):
+        """Test that __aexit__ returns None to propagate exceptions."""
+        client = CopilotClient({"cli_path": CLI_PATH})
+        await client.start()
+        result = await client.__aexit__(None, None, None)
+        assert result is None
+
+    @pytest.mark.asyncio
+    async def test_session_context_manager_returns_self(self):
+        """Test that session __aenter__ returns the session instance."""
+        client = CopilotClient({"cli_path": CLI_PATH})
+        await client.start()
+        try:
+            session = await client.create_session(
+                {"on_permission_request": PermissionHandler.approve_all}
+            )
+            returned_session = await session.__aenter__()
+            assert returned_session is session
+        finally:
+            await client.force_stop()
+
+    @pytest.mark.asyncio
+    async def test_session_aexit_returns_none(self):
+        """Test that session __aexit__ returns None to propagate exceptions."""
+        client = CopilotClient({"cli_path": CLI_PATH})
+        await client.start()
+        try:
+            session = await client.create_session(
+                {"on_permission_request": PermissionHandler.approve_all}
+            )
+            result = await session.__aexit__(None, None, None)
+            assert result is None
+        finally:
+            await client.force_stop()