OpenWorker: initial import

Imported from andrewyng/aisuite@1b4bbf303e
(contents of its platform/ directory, hoisted to the repo root).
Development history prior to this commit lives in that repository.

Co-authored-by: Devika <devikaverma11@gmail.com>

coworker/mcp/client.py

@@ -0,0 +1,143 @@
+"""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
+from contextlib import AsyncExitStack
+from typing import Any, 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
+
+
+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._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) -> _Conn:
+        """Return a live connection for `server`, connecting (once) if needed."""
+        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))
+            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 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) -> 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)
+                    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,
+                    )
+                    read, write = await stack.enter_async_context(stdio_client(params))
+                session = await stack.enter_async_context(ClientSession(read, write))
+                await session.initialize()
+                listed = await session.list_tools()
+                conn = _Conn(session, list(listed.tools))
+                if not ready.done():
+                    ready.set_result(conn)
+                await conn.shutdown.wait()
+        except Exception as exc:  # connection / init failure
+            if not ready.done():
+                ready.set_exception(exc)
+        finally:
+            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