feat: add async context manager support for CopilotClient and Copilot…#475
feat: add async context manager support for CopilotClient and Copilot…#475Sumanth007 wants to merge 3 commits intogithub:mainfrom
Conversation
…Session with automatic resource cleanup
There was a problem hiding this comment.
Pull request overview
Adds async context manager (async with) support to the Python SDK’s CopilotClient and CopilotSession so callers can get automatic startup/teardown, and updates Python docs + tests to encourage/verify the pattern.
Changes:
- Implemented
__aenter__/__aexit__onCopilotClientto auto-start()and cleanup viastop(). - Implemented
__aenter__/__aexit__onCopilotSessionto auto-cleanup viadestroy(). - Added unit/E2E tests and updated
python/README.mdto recommendasync withusage.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| python/copilot/client.py | Adds async context manager methods and cleanup logging for the client lifecycle. |
| python/copilot/session.py | Adds async context manager methods for session lifecycle cleanup. |
| python/test_client.py | Adds unit tests for context manager return values / exception propagation behavior. |
| python/e2e/test_context_managers.py | Adds E2E coverage for client/session context manager behavior and cleanup semantics. |
| python/README.md | Documents and recommends async with usage patterns for safer resource cleanup. |
python/e2e/test_context_managers.py
Outdated
| 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!"}) | ||
|
|
There was a problem hiding this comment.
Same issue here: this test creates a standalone CopilotClient without the E2E proxy env/cwd setup, but then calls session.send(...). That will bypass the replay proxy and can fail depending on auth/network availability. Prefer using ctx.client (or pass env=ctx.get_env() and cwd=ctx.work_dir).
python/copilot/client.py
Outdated
| await self.stop() | ||
| except Exception as e: | ||
| # Log the error but don't raise - we want cleanup to always complete | ||
| logging.warning(f"Error during CopilotClient cleanup: {e}") |
There was a problem hiding this comment.
The cleanup exception logging uses logging.warning(f"..."), which drops the traceback and eagerly formats the message. For better diagnostics, consider using logging.exception(...) or logging.warning(..., exc_info=True) (and pass the exception as an arg rather than f-string formatting).
| logging.warning(f"Error during CopilotClient cleanup: {e}") | |
| logging.warning( | |
| "Error during CopilotClient cleanup: %s", | |
| e, | |
| exc_info=True, | |
| ) |
python/copilot/session.py
Outdated
| except Exception as e: | ||
| # Log the error but don't raise - we want cleanup to always complete | ||
| logging.warning(f"Error during CopilotSession cleanup: {e}") |
There was a problem hiding this comment.
The cleanup exception logging uses logging.warning(f"..."), which omits the traceback and eagerly formats the message. Consider logging.exception(...) or logging.warning(..., exc_info=True) (and pass the exception as an argument) to make cleanup failures diagnosable.
| except Exception as e: | |
| # Log the error but don't raise - we want cleanup to always complete | |
| logging.warning(f"Error during CopilotSession cleanup: {e}") | |
| except Exception: | |
| # Log the error but don't raise - we want cleanup to always complete | |
| logging.warning("Error during CopilotSession cleanup", exc_info=True) |
| try: | ||
| await self.stop() | ||
| except Exception as e: | ||
| # Log the error but don't raise - we want cleanup to always complete | ||
| logging.warning(f"Error during CopilotClient cleanup: {e}") | ||
| return False |
There was a problem hiding this comment.
CopilotClient.__aexit__ calls stop() but ignores the returned StopError list. This means session-destroy failures during cleanup are silently dropped, contradicting the docstring claim that cleanup errors are logged. Consider capturing the return value from stop() and logging each StopError (or aggregating/raising in a controlled way) so cleanup failures are visible to users of the context manager.
| try: | ||
| await self.destroy() | ||
| except Exception as e: | ||
| # Log the error but don't raise - we want cleanup to always complete | ||
| logging.warning(f"Error during CopilotSession cleanup: {e}") | ||
| return False |
There was a problem hiding this comment.
CopilotSession.__aexit__ destroys the session but the owning CopilotClient still keeps the session in its internal _sessions map. When the client later stops (including via CopilotClient.__aexit__), it will attempt to destroy this already-destroyed session again, producing spurious cleanup errors (and keeping references alive longer than needed). Consider making session destruction unregister from the client (e.g., via an owner reference/callback) or making destroy() idempotent so repeated cleanup does not error.
python/e2e/test_context_managers.py
Outdated
| 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" |
There was a problem hiding this comment.
These E2E tests construct new CopilotClient instances without the replay-proxy environment (ctx.get_env() / COPILOT_API_URL) or isolated cwd. Because the tests call session.send(...), they will take the real network/auth path instead of the snapshot proxy and are likely to be flaky or fail in CI. Align with the other snapshot-based E2E tests by using the ctx fixture (e.g., ctx.client) or by passing env=ctx.get_env() and cwd=ctx.work_dir when creating the client.
| class TestCopilotClientContextManager: | ||
| async def test_should_auto_start_and_cleanup_with_context_manager(self, ctx): | ||
| """Test that CopilotClient context manager auto-starts and cleans up.""" | ||
| async with ctx.client 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, ctx): | ||
| """Test creating and using a session within client context.""" | ||
| async with ctx.client 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, ctx): | ||
| """Test that all sessions are cleaned up when client context exits.""" | ||
| async with ctx.client 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, ctx): | ||
| """Test that exceptions within context are propagated.""" | ||
| with pytest.raises(ValueError, match="test error"): | ||
| async with ctx.client 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, ctx): | ||
| """Test that cleanup errors don't prevent context from exiting.""" | ||
| async with ctx.client 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" |
There was a problem hiding this comment.
The tests in TestCopilotClientContextManager use the shared fixture client (ctx.client) within a context manager that calls stop on exit. This creates a pattern where each test restarts the shared fixture's client. While this may work due to start/stop being somewhat idempotent, it's an anti-pattern that mixes fixture-managed lifecycle with context-manager-managed lifecycle. Consider creating fresh CopilotClient instances for each test instead of reusing ctx.client, passing ctx.cli_path and ctx.get_env() to properly configure them with the test harness proxy.
python/README.md
Outdated
| # 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: |
There was a problem hiding this comment.
The example uses "gpt-5" which is not a real model as of January 2025. Consider using "gpt-4" to match other examples in the file (lines 432, 447), or use "gpt-4o" for a more current example.
python/copilot/session.py
Outdated
| if self._destroyed: | ||
| return | ||
|
|
||
| try: | ||
| await self._client.request("session.destroy", {"sessionId": self.session_id}) | ||
| finally: | ||
| # Mark as destroyed and clear handlers even if the request fails | ||
| self._destroyed = True |
There was a problem hiding this comment.
The check and set of self._destroyed is not atomic, which could lead to a race condition if multiple tasks call destroy() concurrently. While the current implementation is mostly safe due to the try/finally block, consider using a lock or making this atomic to prevent the RPC request from being sent multiple times. For example: use self._event_handlers_lock to protect both the flag check and the RPC call, or use an asyncio.Lock.
| if self._destroyed: | |
| return | |
| try: | |
| await self._client.request("session.destroy", {"sessionId": self.session_id}) | |
| finally: | |
| # Mark as destroyed and clear handlers even if the request fails | |
| self._destroyed = True | |
| # 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: | |
| 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 |
… for CopilotClient and CopilotSession
| pytestmark = pytest.mark.asyncio(loop_scope="module") | ||
|
|
||
|
|
There was a problem hiding this comment.
The E2E tests require snapshot files in the test/snapshots/context_managers/ directory to function properly. The test harness expects YAML snapshot files for each test to replay CAPI responses. The snapshot directory test/snapshots/context_managers/ and corresponding YAML files (e.g., should_auto_start_and_cleanup_with_context_manager.yaml, should_create_session_in_context.yaml, etc.) are missing.
The tests will fail without these snapshot files because the replaying proxy uses them to mock API responses. You'll need to either:
- Create the snapshot directory and YAML files with appropriate test data
- Generate the snapshots by running the tests in record mode first
- Copy and adapt existing snapshot files from similar tests
| pytestmark = pytest.mark.asyncio(loop_scope="module") | |
| SNAPSHOT_DIR = os.path.join( | |
| os.path.dirname(os.path.dirname(__file__)), | |
| "test", | |
| "snapshots", | |
| "context_managers", | |
| ) | |
| pytestmark = [ | |
| pytest.mark.asyncio(loop_scope="module"), | |
| pytest.mark.skipif( | |
| not os.path.isdir(SNAPSHOT_DIR), | |
| reason="context_managers snapshots are missing; run E2E in record mode to generate them", | |
| ), | |
| ] |
1 participant
Session with automatic resource cleanup #341
This pull request introduces async context manager support for both
CopilotClientandCopilotSession, enabling automatic resource cleanup and following Python best practices for resource management. The documentation is updated to recommend this pattern, and comprehensive tests ensure correct behavior and error handling. These changes make it easier and safer to use the SDK, especially in batch or evaluation scenarios.Async context manager support and resource management:
__aenter__and__aexit__) toCopilotClientandCopilotSessionfor automatic startup, cleanup, and error handling, ensuring resources are always released—even if exceptions occur. [1] [2]python/README.mdto document and recommend the context manager usage pattern, including example code and benefits, and highlighted this feature in the capabilities list. [1] [2] [3]Testing and validation:
python/e2e/test_context_managers.pyto verify correct context manager behavior, including automatic cleanup, error propagation, and nested context scenarios.python/test_client.pyto confirm that context manager methods return the correct values and propagate exceptions as expected.Internal improvements:
loggingandTracebackTypein relevant files to support error logging and context manager implementation. [1] [2]