From 916d05aa0e612653db6f4e89053870ab6e07b8f5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 15 Feb 2026 12:25:19 +0000 Subject: [PATCH 1/3] Initial plan From 3c38965a416b64f454146a9697a7038ef92c580a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 15 Feb 2026 12:30:39 +0000 Subject: [PATCH 2/3] Add async context manager support for CopilotClient and CopilotSession Co-authored-by: Sumanth007 <61139248+Sumanth007@users.noreply.github.com> --- python/README.md | 78 +++++++++++++++ python/copilot/client.py | 49 ++++++++++ python/copilot/session.py | 47 +++++++++ python/e2e/test_context_managers.py | 142 ++++++++++++++++++++++++++++ python/test_client.py | 43 +++++++++ 5 files changed, 359 insertions(+) create mode 100644 python/e2e/test_context_managers.py diff --git a/python/README.md b/python/README.md index 7aa11e1abc..f1cff745f8 100644 --- a/python/README.md +++ b/python/README.md @@ -14,6 +14,45 @@ uv pip install -e ".[dev]" ## 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-5"}) 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 @@ -56,6 +95,7 @@ asyncio.run(main()) - ✅ Session history with `get_messages()` - ✅ Type hints throughout - ✅ Async/await native +- ✅ Async context manager support for automatic resource cleanup ## API Reference @@ -140,6 +180,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: diff --git a/python/copilot/client.py b/python/copilot/client.py index 11669ddc9c..4b8ba31e50 100644 --- a/python/copilot/client.py +++ b/python/copilot/client.py @@ -21,6 +21,7 @@ import threading from dataclasses import asdict, is_dataclass from pathlib import Path +from types import TracebackType from typing import Any, Callable, Optional, cast from .generated.rpc import ServerRpc @@ -206,6 +207,54 @@ def __init__(self, options: Optional[CopilotClientOptions] = None): self._lifecycle_handlers_lock = threading.Lock() self._rpc: Optional[ServerRpc] = 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], + ) -> bool: + """ + Exit the async context manager. + + Performs graceful cleanup by destroying all active sessions and stopping + the CLI server. If cleanup errors occur, they are logged but do not + prevent the context from exiting. + + Args: + exc_type: The type of exception that occurred, if any. + exc_val: The exception instance that occurred, if any. + exc_tb: The traceback of the exception that occurred, if any. + + Returns: + False to propagate any exception that occurred in the context. + """ + try: + await self.stop() + except Exception as e: + # Log the error but don't raise - we want cleanup to always complete + import logging + + logging.warning(f"Error during CopilotClient cleanup: {e}") + return False + @property def rpc(self) -> ServerRpc: """Typed server-scoped RPC methods.""" diff --git a/python/copilot/session.py b/python/copilot/session.py index d7bd1a3f43..00225f57c1 100644 --- a/python/copilot/session.py +++ b/python/copilot/session.py @@ -8,6 +8,7 @@ import asyncio import inspect import threading +from types import TracebackType from typing import Any, Callable, Optional from .generated.rpc import SessionRpc @@ -82,6 +83,52 @@ def __init__(self, session_id: str, client: Any, workspace_path: Optional[str] = self._hooks_lock = threading.Lock() self._rpc: Optional[SessionRpc] = 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], + ) -> bool: + """ + Exit the async context manager. + + Automatically destroys the session and releases all associated resources. + If an error occurs during cleanup, it is logged but does not prevent the + context from exiting. + + Args: + exc_type: The type of exception that occurred, if any. + exc_val: The exception instance that occurred, if any. + exc_tb: The traceback of the exception that occurred, if any. + + Returns: + False to propagate any exception that occurred in the context. + """ + try: + await self.destroy() + except Exception as e: + # Log the error but don't raise - we want cleanup to always complete + import logging + + logging.warning(f"Error during CopilotSession cleanup: {e}") + return False + @property def rpc(self) -> SessionRpc: """Typed session-scoped RPC methods.""" diff --git a/python/e2e/test_context_managers.py b/python/e2e/test_context_managers.py new file mode 100644 index 0000000000..9f451f9b1b --- /dev/null +++ b/python/e2e/test_context_managers.py @@ -0,0 +1,142 @@ +"""E2E Context Manager Tests""" + +import pytest + +from copilot import CopilotClient + +from .testharness import CLI_PATH + +pytestmark = pytest.mark.asyncio(loop_scope="module") + + +class TestCopilotClientContextManager: + async def test_should_auto_start_and_cleanup_with_context_manager(self): + """Test that CopilotClient context manager auto-starts and cleans up.""" + async with CopilotClient({"cli_path": CLI_PATH}) as client: + assert client.get_state() == "connected" + # Verify we can use the client + pong = await client.ping("test") + assert pong.message == "pong: test" + + # After exiting context, client should be disconnected + assert client.get_state() == "disconnected" + + async def test_should_create_session_in_context(self): + """Test creating and using a session within client context.""" + async with CopilotClient({"cli_path": CLI_PATH}) as client: + session = await client.create_session({"model": "fake-test-model"}) + assert session.session_id + + # Verify session is usable + messages = await session.get_messages() + assert len(messages) > 0 + assert messages[0].type.value == "session.start" + + # After exiting context, verify cleanup happened + assert client.get_state() == "disconnected" + + async def test_should_cleanup_multiple_sessions(self): + """Test that all sessions are cleaned up when client context exits.""" + async with CopilotClient({"cli_path": CLI_PATH}) as client: + session1 = await client.create_session() + session2 = await client.create_session() + session3 = await client.create_session() + + assert session1.session_id + assert session2.session_id + assert session3.session_id + + # All sessions should be cleaned up + assert client.get_state() == "disconnected" + + async def test_should_propagate_exceptions(self): + """Test that exceptions within context are propagated.""" + with pytest.raises(ValueError, match="test error"): + async with CopilotClient({"cli_path": CLI_PATH}) as client: + assert client.get_state() == "connected" + raise ValueError("test error") + + # Client should still be cleaned up even after exception + assert client.get_state() == "disconnected" + + async def test_should_handle_cleanup_errors_gracefully(self): + """Test that cleanup errors don't prevent context from exiting.""" + async with CopilotClient({"cli_path": CLI_PATH}) as client: + await client.create_session() + + # Kill the process to force cleanup to fail + if client._process: + client._process.kill() + + # Context should still exit successfully despite cleanup errors + assert client.get_state() == "disconnected" + + +class TestCopilotSessionContextManager: + async def test_should_cleanup_session_with_context_manager(self): + """Test that CopilotSession context manager cleans up session.""" + client = CopilotClient({"cli_path": CLI_PATH}) + await client.start() + + try: + async with await client.create_session() as session: + assert session.session_id + # Send a message to verify session is working + await session.send({"prompt": "Hello!"}) + + # After exiting context, session should be destroyed + with pytest.raises(Exception, match="Session not found"): + await session.get_messages() + finally: + await client.force_stop() + + async def test_should_propagate_exceptions_in_session_context(self): + """Test that exceptions within session context are propagated.""" + client = CopilotClient({"cli_path": CLI_PATH}) + await client.start() + + try: + with pytest.raises(ValueError, match="test session error"): + async with await client.create_session() as session: + assert session.session_id + raise ValueError("test session error") + + # Session should still be cleaned up after exception + with pytest.raises(Exception, match="Session not found"): + await session.get_messages() + finally: + await client.force_stop() + + async def test_nested_context_managers(self): + """Test using nested context managers for client and session.""" + async with CopilotClient({"cli_path": CLI_PATH}) as client: + async with await client.create_session() as session: + assert session.session_id + await session.send({"prompt": "Test message"}) + + # Session should be cleaned up + with pytest.raises(Exception, match="Session not found"): + await session.get_messages() + + # Client should be cleaned up + assert client.get_state() == "disconnected" + + async def test_multiple_sequential_session_contexts(self): + """Test creating multiple sessions sequentially with context managers.""" + async with CopilotClient({"cli_path": CLI_PATH}) as client: + # First session + async with await client.create_session() as session1: + session1_id = session1.session_id + await session1.send({"prompt": "First session"}) + + # Second session (after first is cleaned up) + async with await client.create_session() as session2: + session2_id = session2.session_id + await session2.send({"prompt": "Second session"}) + + # Both sessions should be different + assert session1_id != session2_id + + # First session should be destroyed + with pytest.raises(Exception, match="Session not found"): + await session1.get_messages() diff --git a/python/test_client.py b/python/test_client.py index 7b4af8c0f0..17c585ae05 100644 --- a/python/test_client.py +++ b/python/test_client.py @@ -147,3 +147,46 @@ def test_use_logged_in_user_with_cli_url_raises(self): CopilotClient( {"cli_url": "localhost:8080", "use_logged_in_user": False, "log_level": "error"} ) + + +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_false(self): + """Test that __aexit__ returns False to propagate exceptions.""" + client = CopilotClient({"cli_path": CLI_PATH}) + await client.start() + result = await client.__aexit__(None, None, None) + assert result is False + + @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() + returned_session = await session.__aenter__() + assert returned_session is session + finally: + await client.force_stop() + + @pytest.mark.asyncio + async def test_session_aexit_returns_false(self): + """Test that session __aexit__ returns False to propagate exceptions.""" + client = CopilotClient({"cli_path": CLI_PATH}) + await client.start() + try: + session = await client.create_session() + result = await session.__aexit__(None, None, None) + assert result is False + finally: + await client.force_stop() + From ebbc3a7faf5c66ed2b6149e7d76c1761806043fc Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 15 Feb 2026 12:35:06 +0000 Subject: [PATCH 3/3] Move logging import to top of files (code review feedback) Co-authored-by: Sumanth007 <61139248+Sumanth007@users.noreply.github.com> --- python/copilot/client.py | 3 +-- python/copilot/session.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/python/copilot/client.py b/python/copilot/client.py index 4b8ba31e50..c88c5afa76 100644 --- a/python/copilot/client.py +++ b/python/copilot/client.py @@ -14,6 +14,7 @@ import asyncio import inspect +import logging import os import re import subprocess @@ -250,8 +251,6 @@ async def __aexit__( await self.stop() except Exception as e: # Log the error but don't raise - we want cleanup to always complete - import logging - logging.warning(f"Error during CopilotClient cleanup: {e}") return False diff --git a/python/copilot/session.py b/python/copilot/session.py index 00225f57c1..969205fd94 100644 --- a/python/copilot/session.py +++ b/python/copilot/session.py @@ -7,6 +7,7 @@ import asyncio import inspect +import logging import threading from types import TracebackType from typing import Any, Callable, Optional @@ -124,8 +125,6 @@ async def __aexit__( await self.destroy() except Exception as e: # Log the error but don't raise - we want cleanup to always complete - import logging - logging.warning(f"Error during CopilotSession cleanup: {e}") return False