Thin wrapper around fastmcp Client for this codebase. Durable execution note: - Sandbox tools require `ctx.session_id` which is ultimately derived from MCP's stateful Streamable HTTP session id (`mcp-session-id` header). - To support resume after process kill, we must be able
| 35 | |
| 36 | |
| 37 | class MCPClient: |
| 38 | """ |
| 39 | Thin wrapper around fastmcp Client for this codebase. |
| 40 | |
| 41 | Durable execution note: |
| 42 | - Sandbox tools require `ctx.session_id` which is ultimately derived from MCP's |
| 43 | stateful Streamable HTTP session id (`mcp-session-id` header). |
| 44 | - To support resume after process kill, we must be able to reattach to the same |
| 45 | MCP session by re-sending the same `mcp-session-id`. |
| 46 | """ |
| 47 | |
| 48 | def __init__( |
| 49 | self, |
| 50 | url: str = None, |
| 51 | tool_config_fp: str = None, |
| 52 | *, |
| 53 | session_id: Optional[str] = None, |
| 54 | trace_path: Optional[str] = None, |
| 55 | replay_trace_path: Optional[str] = None, |
| 56 | connect: bool = True, |
| 57 | ): |
| 58 | hostname = os.getenv("MCP_HOSTNAME") or "localhost" |
| 59 | port = os.getenv("MCP_PORT") or 7002 |
| 60 | basepath = os.getenv("MCP_BASEPATH") or "/mcp" |
| 61 | self.server_url = url or f"http://{hostname}:{port}{basepath}" |
| 62 | self._client: Client | None = None |
| 63 | self._session_id: Optional[str] = session_id |
| 64 | |
| 65 | # Durable execution: tool-level trace/replay |
| 66 | self._trace_path: Optional[str] = trace_path |
| 67 | self._trace_lock = threading.Lock() |
| 68 | self._trace_seq = 0 |
| 69 | self._replay_events: Optional[List[Dict[str, Any]]] = None |
| 70 | self._replay_idx = 0 |
| 71 | self._replay_lock = threading.Lock() |
| 72 | if replay_trace_path: |
| 73 | self.enable_replay(replay_trace_path) |
| 74 | |
| 75 | # Thread lock for thread-safe access to the event loop |
| 76 | # This is necessary because event loops are not thread-safe |
| 77 | self._loop_lock = threading.Lock() |
| 78 | self._loop: Optional[asyncio.AbstractEventLoop] = None |
| 79 | |
| 80 | if connect: |
| 81 | self._loop = asyncio.new_event_loop() |
| 82 | asyncio.set_event_loop(self._loop) |
| 83 | self._loop.run_until_complete(self._init_client()) |
| 84 | |
| 85 | def _update_session_id_from_transport(self) -> None: |
| 86 | """ |
| 87 | Best-effort: capture current MCP session id from transport. |
| 88 | """ |
| 89 | try: |
| 90 | if not self._client: |
| 91 | return |
| 92 | transport = getattr(self._client, "transport", None) |
| 93 | get_sid = getattr(transport, "get_session_id", None) |
| 94 | if callable(get_sid): |
no outgoing calls