mirror of
https://github.com/andrewyng/openworker.git
synced 2026-09-11 06:30:25 +00:00
Rohit C Prasad
476032e603
verify() replaces cached-yes ensure for explicit tests; dead connections tear down and reconnect. Connected rows show 'tested ⟨when⟩' so the click has a visible result.
224 lines
9.2 KiB
Python
224 lines
9.2 KiB
Python
"""MCPManager — our own thin async MCP client over the official `mcp` SDK.
|
|
|
|
Async-native (no `nest_asyncio`, no second event loop): each server runs in a dedicated
|
|
asyncio task that opens the transport + `ClientSession`, keeps them alive until shutdown,
|
|
then closes them in the *same* task — required because the SDK's transports use anyio cancel
|
|
scopes that must be entered and exited on one task. Tool calls are awaited from any task on
|
|
the same loop, which is safe.
|
|
|
|
Tool execution from the (sync) ToolRegistry bridges back here via
|
|
`run_coroutine_threadsafe` — see `coworker/mcp/tools.py`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import tempfile
|
|
from contextlib import AsyncExitStack
|
|
from typing import Any, IO, Optional
|
|
|
|
from mcp import ClientSession, StdioServerParameters
|
|
from mcp.client.stdio import stdio_client
|
|
from mcp.client.streamable_http import streamablehttp_client
|
|
|
|
from .config import MCPServerDef
|
|
|
|
|
|
_STDERR_TAIL_LINES = 20
|
|
_STDERR_TAIL_CHARS = 1500
|
|
|
|
|
|
def _read_tail(errfile: Optional[IO[str]]) -> Optional[str]:
|
|
"""Last few lines of a captured stderr file — the crash evidence, not the log."""
|
|
if errfile is None:
|
|
return None
|
|
try:
|
|
errfile.seek(0)
|
|
text = errfile.read()
|
|
except (OSError, ValueError):
|
|
return None
|
|
lines = [ln for ln in text.strip().splitlines() if ln.strip()]
|
|
if not lines:
|
|
return None
|
|
return "\n".join(lines[-_STDERR_TAIL_LINES:])[-_STDERR_TAIL_CHARS:]
|
|
|
|
|
|
class _Conn:
|
|
def __init__(self, session: ClientSession, tools: list[Any]) -> None:
|
|
self.session = session
|
|
self.tools = tools # list[mcp.types.Tool]
|
|
self.shutdown = asyncio.Event()
|
|
|
|
|
|
class MCPManager:
|
|
"""Owns persistent MCP connections keyed by server name; lazy-connects on demand."""
|
|
|
|
def __init__(self, secrets: Any = None) -> None:
|
|
self._conns: dict[str, _Conn] = {}
|
|
self._tasks: dict[str, asyncio.Task] = {}
|
|
self._stderr_tails: dict[str, str] = {}
|
|
self._lock = asyncio.Lock()
|
|
# SecretStore for OAuth servers' token persistence (mcp/oauth.py); lazy default
|
|
# so library/CLI construction without secrets keeps working.
|
|
self._secrets = secrets
|
|
|
|
async def ensure(self, server: MCPServerDef, *, interactive: bool = False) -> _Conn:
|
|
"""Return a live connection for `server`, connecting (once) if needed.
|
|
|
|
`interactive=True` (explicit connect actions only) lets an OAuth server run
|
|
the browser sign-in flow; the default refuses it — stored tokens and silent
|
|
refresh still work, but a server that insists on re-authorization raises
|
|
InteractiveAuthRequired instead of hijacking the user's browser.
|
|
"""
|
|
async with self._lock:
|
|
existing = self._conns.get(server.name)
|
|
if existing is not None:
|
|
return existing
|
|
ready: asyncio.Future = asyncio.get_running_loop().create_future()
|
|
self._tasks[server.name] = asyncio.create_task(
|
|
self._serve(server, ready, interactive=interactive)
|
|
)
|
|
conn = await ready # propagates connection errors
|
|
self._conns[server.name] = conn
|
|
return conn
|
|
|
|
async def tools(self, server: MCPServerDef) -> list[Any]:
|
|
return (await self.ensure(server)).tools
|
|
|
|
async def verify(self, server: MCPServerDef, *, interactive: bool = False) -> _Conn:
|
|
"""A REAL health check for explicit Test actions. `ensure` returns a cached
|
|
connection untouched, which made Test-on-Live a silent no-op that could not
|
|
detect a dead server (owner-hit 2026-08-21). Here a cached connection is
|
|
round-tripped (tools/list, refreshing the tool set); a dead one is torn
|
|
down and reconnected fresh."""
|
|
conn = self._conns.get(server.name)
|
|
if conn is not None:
|
|
try:
|
|
listed = await asyncio.wait_for(conn.session.list_tools(), timeout=20)
|
|
conn.tools = list(listed.tools)
|
|
return conn
|
|
except Exception:
|
|
conn.shutdown.set()
|
|
task = self._tasks.pop(server.name, None)
|
|
if task is not None:
|
|
try:
|
|
await asyncio.wait_for(asyncio.shield(task), timeout=5)
|
|
except Exception:
|
|
task.cancel()
|
|
self._conns.pop(server.name, None) # _serve pops too; belt and braces
|
|
return await self.ensure(server, interactive=interactive)
|
|
|
|
def last_stderr(self, name: str) -> Optional[str]:
|
|
"""Stderr tail from the most recent failed startup of `name`, if any."""
|
|
return self._stderr_tails.get(name)
|
|
|
|
async def call(
|
|
self, name: str, tool: str, arguments: Optional[dict[str, Any]]
|
|
) -> Any:
|
|
conn = self._conns.get(name)
|
|
if conn is None:
|
|
raise RuntimeError(f"MCP server not connected: {name}")
|
|
result = await conn.session.call_tool(tool, arguments or {})
|
|
return _result_payload(result)
|
|
|
|
async def aclose(self) -> None:
|
|
for conn in self._conns.values():
|
|
conn.shutdown.set()
|
|
for task in list(self._tasks.values()):
|
|
try:
|
|
await asyncio.wait_for(asyncio.shield(task), timeout=5)
|
|
except (asyncio.TimeoutError, Exception):
|
|
task.cancel()
|
|
self._conns.clear()
|
|
self._tasks.clear()
|
|
|
|
# -- per-server lifecycle (one task owns enter+exit) ------------------------
|
|
async def _serve(
|
|
self, server: MCPServerDef, ready: asyncio.Future, *, interactive: bool = False
|
|
) -> None:
|
|
errfile = None
|
|
try:
|
|
async with AsyncExitStack() as stack:
|
|
if server.transport == "http":
|
|
if not server.url:
|
|
raise ValueError(
|
|
f"MCP server '{server.name}' is http but has no url"
|
|
)
|
|
auth = None
|
|
if server.auth == "oauth":
|
|
from ..secrets import SecretStore
|
|
from .oauth import build_auth
|
|
|
|
if self._secrets is None:
|
|
self._secrets = SecretStore()
|
|
auth = build_auth(
|
|
server.name,
|
|
server.url,
|
|
self._secrets,
|
|
interactive=interactive,
|
|
)
|
|
read, write, *_ = await stack.enter_async_context(
|
|
streamablehttp_client(
|
|
server.url, headers=server.headers or None, auth=auth
|
|
)
|
|
)
|
|
else:
|
|
if not server.command:
|
|
raise ValueError(
|
|
f"MCP server '{server.name}' is stdio but has no command"
|
|
)
|
|
params = StdioServerParameters(
|
|
command=server.command,
|
|
args=server.args,
|
|
env=server.env or None,
|
|
cwd=server.cwd,
|
|
)
|
|
# Capture the child's stderr so a startup crash leaves evidence
|
|
# the UI can show (the SDK needs a real file descriptor here).
|
|
errfile = tempfile.TemporaryFile(
|
|
mode="w+", encoding="utf-8", errors="replace"
|
|
)
|
|
read, write = await stack.enter_async_context(
|
|
stdio_client(params, errlog=errfile)
|
|
)
|
|
session = await stack.enter_async_context(ClientSession(read, write))
|
|
await session.initialize()
|
|
listed = await session.list_tools()
|
|
conn = _Conn(session, list(listed.tools))
|
|
self._stderr_tails.pop(server.name, None)
|
|
if not ready.done():
|
|
ready.set_result(conn)
|
|
await conn.shutdown.wait()
|
|
except Exception as exc: # connection / init failure
|
|
tail = _read_tail(errfile)
|
|
if tail:
|
|
self._stderr_tails[server.name] = tail
|
|
if not ready.done():
|
|
ready.set_exception(exc)
|
|
finally:
|
|
if errfile is not None:
|
|
try:
|
|
errfile.close()
|
|
except OSError:
|
|
pass
|
|
self._conns.pop(server.name, None)
|
|
self._tasks.pop(server.name, None)
|
|
|
|
|
|
def _result_payload(result: Any) -> Any:
|
|
"""Flatten a CallToolResult into something the engine can serialize for the model."""
|
|
texts: list[str] = []
|
|
for block in getattr(result, "content", None) or []:
|
|
text = getattr(block, "text", None)
|
|
if text is not None:
|
|
texts.append(text)
|
|
else: # non-text content (image/resource) — describe it
|
|
texts.append(f"[{getattr(block, 'type', 'content')}]")
|
|
body = "\n".join(texts)
|
|
if getattr(result, "isError", False):
|
|
return {"error": body or "MCP tool error"}
|
|
structured = getattr(result, "structuredContent", None)
|
|
if structured is not None and not body:
|
|
return structured
|
|
return body
|