python-sdk icon indicating copy to clipboard operation
python-sdk copied to clipboard

[OpenAI Codex PR] Fix cleanup order issue with task groups

Open jhaoming-oai opened this issue 8 months ago • 0 comments

Summary

  • fix the issue mentioned in #577
  • add CompatTaskGroup helper that uses asyncio.TaskGroup when 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

Additional context

jhaoming-oai avatar May 21 '25 16:05 jhaoming-oai