python-sdk
python-sdk copied to clipboard
[OpenAI Codex PR] Fix cleanup order issue with task groups
Summary
- fix the issue mentioned in #577
- add CompatTaskGroup helper that uses
asyncio.TaskGroupwhen available - use CompatTaskGroup across client and server transports
- update BaseSession to rely on CompatTaskGroup
- this PR is created by OpenAI Codex
Testing
-
ruff check --fix . -
pyright -p .(fails: venv not found and missing deps) -
pytest -q(fails: command not found) - test with the following example
import os, asyncio, json
from typing import Optional
from contextlib import AsyncExitStack
from mcp import ClientSession, StdioServerParameters
from mcp.types import TextContent
from mcp.client.stdio import stdio_client
class MCPClient:
def __init__(self, command: str, args: list[str], env: Optional[dict] = None):
self.session: Optional[ClientSession] = None
self.command, self.args, self.env = command, args, env
self._cleanup_lock = asyncio.Lock()
self.exit_stack: Optional[AsyncExitStack] = None
async def connect_to_server(self):
await self.cleanup()
self.exit_stack = AsyncExitStack()
server_params = StdioServerParameters(
command=self.command, args=self.args, env=self.env
)
stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
self.stdio, self.write = stdio_transport
self.session = await self.exit_stack.enter_async_context(
ClientSession(self.stdio, self.write)
)
await self.session.initialize()
async def cleanup(self):
if self.exit_stack:
async with self._cleanup_lock:
await self.exit_stack.aclose()
self.session = None
self.exit_stack = None
async def main():
cfg = {
"command": "npx",
"args": [
"-y",
"@modelcontextprotocol/server-filesystem",
"/Users/jhaoming/Desktop"
]
}
c1, c2 = MCPClient(**cfg), MCPClient(**cfg)
await c1.connect_to_server()
await c2.connect_to_server()
# Works (FILO)
await c2.cleanup()
await c1.cleanup()
# Fails (FIFO)
# await c1.cleanup() # <-- boom before pr, fixed after pr
# await c2.cleanup()
if __name__ == "__main__":
asyncio.run(main())
Breaking Changes
Types of changes
- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation update
Checklist
- [x] I have read the MCP Documentation
- [ ] My code follows the repository's style guidelines
- [ ] New and existing tests pass locally
- [x] I have added appropriate error handling
- [ ] I have added or updated documentation as needed