feat: OpenMesh 基础平台与 MD/PDF 转换技能
- 后端: coworker 智能体框架, WS API, 文件上传, 附件处理 - 前端: Open WebUI, 文件全量走 upload API (含 MD/TXT/JSON 等文本类) - 技能: md-to-office (pandoc + wkhtmltopdf) - 修复: 上传文件路径丢失, Agent 搜索浪费, 输出文件跑到 uploads/ - 打包: PyInstaller one-dir, 预打包 pandoc/wkhtmltopdf/chromium
This commit is contained in:
29
coworker/mcp/__init__.py
Normal file
29
coworker/mcp/__init__.py
Normal file
@@ -0,0 +1,29 @@
|
||||
"""MCP integration — our own async client on the official `mcp` SDK.
|
||||
|
||||
Public API: config loading/mutation, the connection manager, and tool wrapping.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .client import MCPManager
|
||||
from .config import (
|
||||
MCPServerDef,
|
||||
delete_global_server,
|
||||
load_mcp_servers,
|
||||
patch_global_server,
|
||||
put_global_server,
|
||||
read_global,
|
||||
)
|
||||
from .tools import build_callables, tool_name
|
||||
|
||||
__all__ = [
|
||||
"MCPManager",
|
||||
"MCPServerDef",
|
||||
"load_mcp_servers",
|
||||
"read_global",
|
||||
"put_global_server",
|
||||
"patch_global_server",
|
||||
"delete_global_server",
|
||||
"build_callables",
|
||||
"tool_name",
|
||||
]
|
||||
223
coworker/mcp/client.py
Normal file
223
coworker/mcp/client.py
Normal file
@@ -0,0 +1,223 @@
|
||||
"""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
|
||||
150
coworker/mcp/config.py
Normal file
150
coworker/mcp/config.py
Normal file
@@ -0,0 +1,150 @@
|
||||
"""MCP server config — the standard `mcpServers` JSON, layered global + workspace.
|
||||
|
||||
Global: ~/.config/coworker/mcp.json
|
||||
Workspace: <workspace>/.coworker/mcp.json (overrides global on name clash,
|
||||
but only after the user trusts that workspace — same gate as
|
||||
repository `allowed_commands`)
|
||||
|
||||
Paste-compatible with Claude Desktop / Cursor / Codex. `${VAR}` refs in command/args/env/
|
||||
url/headers are resolved at load time via the SecretStore (env + local `.env`). REST edits
|
||||
target the **global** file.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from ..secrets import SecretStore, state_dir
|
||||
|
||||
_HTTP_TYPES = {"http", "https", "sse", "streamable-http", "streamable_http"}
|
||||
|
||||
|
||||
@dataclass
|
||||
class MCPServerDef:
|
||||
name: str
|
||||
transport: str # "stdio" | "http"
|
||||
command: Optional[str] = None
|
||||
args: list[str] = field(default_factory=list)
|
||||
env: dict[str, str] = field(default_factory=dict)
|
||||
cwd: Optional[str] = None
|
||||
url: Optional[str] = None
|
||||
headers: dict[str, str] = field(default_factory=dict)
|
||||
enabled: bool = True
|
||||
include_tools: Optional[list[str]] = None
|
||||
exclude_tools: Optional[list[str]] = None
|
||||
requires_approval: bool = True
|
||||
# "oauth" → browser OAuth 2.1 + PKCE with Dynamic Client Registration (mcp/oauth.py).
|
||||
# HTTP transport only; tokens live in the SecretStore, never in this file.
|
||||
auth: Optional[str] = None
|
||||
|
||||
|
||||
def global_mcp_path() -> Path:
|
||||
return state_dir() / "mcp.json"
|
||||
|
||||
|
||||
def _read(path: Path) -> dict[str, Any]:
|
||||
try:
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return {}
|
||||
|
||||
|
||||
def _config_paths(
|
||||
workspace: Optional[str | Path], *, workspace_trusted: bool
|
||||
) -> list[Path]:
|
||||
"""Config files to merge. Workspace MCP is executable provenance (stdio spawn),
|
||||
so an untrusted repo's `.coworker/mcp.json` is never read — cloning alone must
|
||||
not be enough to define processes that run at session open.
|
||||
"""
|
||||
paths = [global_mcp_path()]
|
||||
if workspace and workspace_trusted:
|
||||
paths.append(Path(workspace).expanduser() / ".coworker" / "mcp.json")
|
||||
return paths
|
||||
|
||||
|
||||
def _parse(name: str, raw: dict[str, Any], secrets: SecretStore) -> MCPServerDef:
|
||||
raw = secrets.resolve(raw) # resolve ${VAR} everywhere before building the def
|
||||
declared = str(raw.get("type", "")).lower()
|
||||
is_http = declared in _HTTP_TYPES or bool(raw.get("url"))
|
||||
return MCPServerDef(
|
||||
name=name,
|
||||
transport="http" if is_http else "stdio",
|
||||
command=raw.get("command"),
|
||||
args=list(raw.get("args", []) or []),
|
||||
env={str(k): str(v) for k, v in (raw.get("env") or {}).items()},
|
||||
cwd=raw.get("cwd"),
|
||||
url=raw.get("url"),
|
||||
headers={str(k): str(v) for k, v in (raw.get("headers") or {}).items()},
|
||||
enabled=bool(raw.get("enabled", True)),
|
||||
include_tools=raw.get("include_tools"),
|
||||
exclude_tools=raw.get("exclude_tools"),
|
||||
requires_approval=bool(raw.get("requires_approval", True)),
|
||||
auth=(str(raw["auth"]).lower() if raw.get("auth") else None),
|
||||
)
|
||||
|
||||
|
||||
def load_mcp_servers(
|
||||
workspace: Optional[str | Path] = None,
|
||||
*,
|
||||
secrets: Optional[SecretStore] = None,
|
||||
workspace_trusted: bool = False,
|
||||
) -> list[MCPServerDef]:
|
||||
"""Merge global + (when trusted) workspace `mcpServers` into parsed server defs.
|
||||
|
||||
Only trusted workspaces contribute — the same consent boundary as repository
|
||||
``allowed_commands`` — and **global wins on name clash**, so even a trusted repo
|
||||
cannot silently redefine a global server by reusing its name. ``${VAR}`` refs in
|
||||
a workspace def are resolved from the user's env, which is acceptable only because
|
||||
the workspace is trusted; untrusted workspaces are never read.
|
||||
"""
|
||||
secrets = secrets or SecretStore()
|
||||
merged: dict[str, dict[str, Any]] = {}
|
||||
for path in _config_paths(workspace, workspace_trusted=workspace_trusted):
|
||||
for name, raw in (_read(path).get("mcpServers") or {}).items():
|
||||
if isinstance(raw, dict):
|
||||
merged.setdefault(name, raw) # global first → global wins on clash
|
||||
return [_parse(name, raw, secrets) for name, raw in merged.items()]
|
||||
|
||||
|
||||
# -- raw global-file mutation (REST) -------------------------------------------
|
||||
def read_global() -> dict[str, dict[str, Any]]:
|
||||
"""Raw `mcpServers` map from the global file (no `${VAR}` resolution)."""
|
||||
return dict(_read(global_mcp_path()).get("mcpServers") or {})
|
||||
|
||||
|
||||
def _write_global(servers: dict[str, dict[str, Any]]) -> None:
|
||||
path = global_mcp_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = path.with_name(path.name + ".tmp")
|
||||
tmp.write_text(json.dumps({"mcpServers": servers}, indent=2), encoding="utf-8")
|
||||
tmp.replace(path)
|
||||
|
||||
|
||||
def put_global_server(name: str, config: dict[str, Any]) -> None:
|
||||
servers = read_global()
|
||||
servers[name] = config
|
||||
_write_global(servers)
|
||||
|
||||
|
||||
def patch_global_server(name: str, changes: dict[str, Any]) -> bool:
|
||||
servers = read_global()
|
||||
if name not in servers:
|
||||
return False
|
||||
merged = {**servers[name], **changes}
|
||||
# A None value DELETES the key (there is no other way to remove one through a
|
||||
# merge patch) — used by the OPE-136 trust migration to drop `requires_approval`.
|
||||
servers[name] = {k: v for k, v in merged.items() if v is not None}
|
||||
_write_global(servers)
|
||||
return True
|
||||
|
||||
|
||||
def delete_global_server(name: str) -> bool:
|
||||
servers = read_global()
|
||||
if name not in servers:
|
||||
return False
|
||||
del servers[name]
|
||||
_write_global(servers)
|
||||
return True
|
||||
349
coworker/mcp/oauth.py
Normal file
349
coworker/mcp/oauth.py
Normal file
@@ -0,0 +1,349 @@
|
||||
"""Browser OAuth for remote MCP servers (OAuth 2.1 + PKCE + Dynamic Client Registration).
|
||||
|
||||
The official SDK's `OAuthClientProvider` drives the whole spec flow — protected-resource
|
||||
metadata discovery, DCR, PKCE, token refresh — as an httpx auth plugged into the
|
||||
streamable-HTTP transport. We supply its three integration points:
|
||||
|
||||
- token persistence → the SecretStore (profile `mcp-oauth:<server>`; 0600 file,
|
||||
never the mcp.json config, which is plain text and paste-shareable)
|
||||
- redirect → open the system browser at the authorize URL
|
||||
- callback → the sidecar's loopback `GET /mcp/oauth/callback` resolves a
|
||||
single-slot pending future (one interactive sign-in at a time — the flow is
|
||||
user-driven, so concurrency is meaningless)
|
||||
|
||||
DCR means there is no client id/secret registered anywhere up front — nothing for the
|
||||
ocw-connect broker to hold, so unlike the managed connectors this flow is fully local.
|
||||
First server: Granola (https://mcp.granola.ai/mcp).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import secrets
|
||||
import time
|
||||
from typing import Any, Optional
|
||||
|
||||
from mcp.client.auth import OAuthClientProvider, TokenStorage
|
||||
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken
|
||||
|
||||
from ..secrets import SecretStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PROFILE_PREFIX = "mcp-oauth:"
|
||||
CALLBACK_PATH = "/mcp/oauth/callback"
|
||||
# How long the connect waits for the user to finish the browser sign-in.
|
||||
FLOW_TIMEOUT_SECONDS = 300
|
||||
|
||||
CLIENT_NAME = "OpenWorker"
|
||||
|
||||
|
||||
def redirect_base() -> str:
|
||||
"""The sidecar's own loopback origin — the DCR-registered redirect must match it."""
|
||||
port = os.environ.get("COWORKER_PORT") or "8765"
|
||||
return f"http://127.0.0.1:{port}"
|
||||
|
||||
|
||||
def _profile(name: str) -> str:
|
||||
return PROFILE_PREFIX + name
|
||||
|
||||
|
||||
class SecretStoreTokenStorage(TokenStorage):
|
||||
"""SDK TokenStorage over our SecretStore: one profile per server holding the token
|
||||
set and the DCR-issued client registration (re-used across sign-ins)."""
|
||||
|
||||
def __init__(self, server_name: str, secrets: SecretStore) -> None:
|
||||
self._name = server_name
|
||||
self._secrets = secrets
|
||||
|
||||
def _data(self) -> dict[str, Any]:
|
||||
return self._secrets.get(_profile(self._name)) or {}
|
||||
|
||||
def _merge(self, patch: dict[str, Any]) -> None:
|
||||
self._secrets.put(_profile(self._name), {**self._data(), **patch})
|
||||
|
||||
async def get_tokens(self) -> Optional[OAuthToken]:
|
||||
data = self._data()
|
||||
raw = data.get("tokens")
|
||||
if not raw:
|
||||
return None
|
||||
try:
|
||||
tok = OAuthToken.model_validate(raw)
|
||||
except Exception:
|
||||
return None
|
||||
# SDK flaw (mcp 1.29): `_initialize()` loads stored tokens but never computes
|
||||
# `token_expiry_time`, and `is_token_valid()` treats None expiry as valid
|
||||
# forever — so an hour-old access token is sent as-is, the server 401s, and
|
||||
# the SDK's 401 branch goes straight to FULL re-authorization without trying
|
||||
# the refresh token. Non-interactive contexts must refuse the browser, so
|
||||
# every session said "sign-in required" while explicit connects appeared to
|
||||
# work (owner-hit 2026-08-21, DLAI Redshift). Countermeasure lives here, in
|
||||
# storage: when the stored token is past the lifetime we recorded at save
|
||||
# time (unknown age = stale), return the token set WITHOUT the access token —
|
||||
# `is_token_valid()` then fails on its own terms and the SDK runs the
|
||||
# refresh-token grant FIRST, which self-heals silently (no browser).
|
||||
if tok.expires_in is not None:
|
||||
issued = data.get("tokens_issued_at")
|
||||
if isinstance(issued, (int, float)):
|
||||
remaining = int(issued + tok.expires_in - time.time())
|
||||
else:
|
||||
remaining = -1
|
||||
tok = tok.model_copy(update={"expires_in": remaining})
|
||||
if remaining <= 60 and tok.refresh_token:
|
||||
tok = tok.model_copy(update={"access_token": ""})
|
||||
return tok
|
||||
|
||||
async def set_tokens(self, tokens: OAuthToken) -> None:
|
||||
self._merge(
|
||||
{
|
||||
"tokens": tokens.model_dump(mode="json", exclude_none=True),
|
||||
"tokens_issued_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
|
||||
async def get_client_info(self) -> Optional[OAuthClientInformationFull]:
|
||||
raw = self._data().get("client_info")
|
||||
if not raw:
|
||||
return None
|
||||
try:
|
||||
return OAuthClientInformationFull.model_validate(raw)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
async def set_client_info(self, info: OAuthClientInformationFull) -> None:
|
||||
self._merge({"client_info": info.model_dump(mode="json", exclude_none=True)})
|
||||
|
||||
|
||||
class InteractiveAuthRequired(RuntimeError):
|
||||
"""The server wants a browser sign-in, but this context must not open one.
|
||||
|
||||
Interactive OAuth (browser + loopback wait) is an explicit-connect-only
|
||||
privilege: a background context that hit this — an engine turn, a tools
|
||||
listing — raises instead, and the caller skips the server. Without this, a
|
||||
server whose refresh token the vendor rejected (Atlassian rotates them
|
||||
aggressively) would hijack the user's browser from ANY code path that
|
||||
touched it — owner-hit 2026-07-20: an authorize page opened at app launch.
|
||||
"""
|
||||
|
||||
|
||||
def is_auth_required(exc: BaseException) -> bool:
|
||||
"""True if InteractiveAuthRequired is anywhere in the exception tree — the SDK
|
||||
transport runs in anyio task groups, so it often arrives wrapped in an
|
||||
ExceptionGroup (or chained as a cause) rather than bare."""
|
||||
if isinstance(exc, InteractiveAuthRequired):
|
||||
return True
|
||||
for sub in getattr(exc, "exceptions", None) or []: # ExceptionGroup
|
||||
if is_auth_required(sub):
|
||||
return True
|
||||
cause = exc.__cause__ or exc.__context__
|
||||
return is_auth_required(cause) if cause is not None else False
|
||||
|
||||
|
||||
def is_http_auth_error(exc: BaseException) -> bool:
|
||||
"""True if an HTTP 401/403 is anywhere in the exception tree — an anonymous
|
||||
connect hit a server that wants credentials, so the fix is sign-in (switch
|
||||
the entry to `auth: oauth`), not a different config. Same tree walk as
|
||||
is_auth_required: the transport's task groups wrap and chain freely."""
|
||||
status = getattr(getattr(exc, "response", None), "status_code", None)
|
||||
if status in (401, 403):
|
||||
return True
|
||||
for sub in getattr(exc, "exceptions", None) or []: # ExceptionGroup
|
||||
if is_http_auth_error(sub):
|
||||
return True
|
||||
cause = exc.__cause__ or exc.__context__
|
||||
return is_http_auth_error(cause) if cause is not None else False
|
||||
|
||||
|
||||
# -- single-slot interactive flow ------------------------------------------------
|
||||
_pending: Optional[asyncio.Future] = None
|
||||
# The last authorize URL we sent the user to — surfaced over REST so the GUI can offer
|
||||
# a "reopen sign-in page" link if the browser popup was lost.
|
||||
last_authorize_url: Optional[str] = None
|
||||
# The `state` the SDK put in the current authorize URL. The SDK itself re-checks the
|
||||
# returned state (mcp.client.auth.oauth2 compare_digest), so this is NOT the CSRF guard —
|
||||
# it's a loopback gate: without it any local caller could hit /mcp/oauth/callback with a
|
||||
# bogus code and consume the single pending future, aborting the user's real sign-in
|
||||
# (which then finds no pending flow). Matching state here rejects that stray callback and
|
||||
# leaves the flow waiting for the genuine one.
|
||||
_expected_state: Optional[str] = None
|
||||
|
||||
|
||||
def _state_from_url(url: str) -> Optional[str]:
|
||||
"""Pull the `state` query param out of an authorize URL (None if absent)."""
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
values = parse_qs(urlsplit(url).query).get("state")
|
||||
return values[0] if values else None
|
||||
|
||||
|
||||
def deliver_callback(code: str, state: Optional[str]) -> bool:
|
||||
"""Called by the loopback route. Resolves the waiting flow; False if none waits.
|
||||
|
||||
A callback whose `state` doesn't match the pending flow's is ignored (returns False)
|
||||
WITHOUT consuming the pending future, so a stray/forged local hit can't abort a live
|
||||
sign-in — only the browser redirect carrying the SDK's own state resolves it.
|
||||
"""
|
||||
global _pending
|
||||
if _pending is None or _pending.done():
|
||||
return False
|
||||
# Only enforce when we actually captured a state for this flow; a flow with no state
|
||||
# in its authorize URL falls back to the prior accept-any behavior.
|
||||
if _expected_state is not None and (
|
||||
state is None or not secrets.compare_digest(state, _expected_state)
|
||||
):
|
||||
return False
|
||||
pending, _pending = _pending, None
|
||||
pending.set_result((code, state))
|
||||
return True
|
||||
|
||||
|
||||
async def _open_browser(url: str) -> None:
|
||||
global last_authorize_url, _expected_state
|
||||
last_authorize_url = url
|
||||
_expected_state = _state_from_url(url)
|
||||
import webbrowser
|
||||
|
||||
logger.info("mcp oauth: opening browser for sign-in")
|
||||
await asyncio.get_running_loop().run_in_executor(None, webbrowser.open, url)
|
||||
|
||||
|
||||
async def _refuse_browser(url: str) -> None:
|
||||
"""Non-interactive redirect handler: never open a browser, but keep the URL so
|
||||
the GUI's "reopen sign-in page" affordance still works after the refusal."""
|
||||
global last_authorize_url
|
||||
last_authorize_url = url
|
||||
raise InteractiveAuthRequired(
|
||||
"sign-in required — reconnect this server from its page"
|
||||
)
|
||||
|
||||
|
||||
async def _refuse_callback() -> tuple[str, Optional[str]]:
|
||||
raise InteractiveAuthRequired(
|
||||
"sign-in required — reconnect this server from its page"
|
||||
)
|
||||
|
||||
|
||||
async def _wait_for_callback() -> tuple[str, Optional[str]]:
|
||||
global _pending, _expected_state
|
||||
if _pending is not None and not _pending.done():
|
||||
_pending.cancel() # a stale flow lost its browser tab; the new one wins
|
||||
_pending = asyncio.get_running_loop().create_future()
|
||||
try:
|
||||
return await asyncio.wait_for(_pending, timeout=FLOW_TIMEOUT_SECONDS)
|
||||
except asyncio.TimeoutError:
|
||||
raise RuntimeError(
|
||||
"sign-in timed out — the browser window was not completed in "
|
||||
f"{FLOW_TIMEOUT_SECONDS // 60} minutes"
|
||||
)
|
||||
finally:
|
||||
_pending = None
|
||||
_expected_state = None # don't let this flow's state gate the next one
|
||||
|
||||
|
||||
class _MetadataSeededProvider(OAuthClientProvider):
|
||||
"""OAuthClientProvider that persists the discovered authorization-server
|
||||
metadata and re-seeds it on load. Without this the SDK's pre-request refresh
|
||||
grant runs BEFORE discovery and falls back to <origin>/token — a 404 on
|
||||
vendors whose real endpoint lives elsewhere (data.dlai.link uses
|
||||
/api/auth/mcp/token), which turned every silent refresh into a full re-auth
|
||||
demand (owner-hit 2026-08-21, with the stale-expiry flaw above)."""
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self._ocw_storage: SecretStoreTokenStorage = kwargs.get("storage") or self.context.storage # type: ignore[assignment]
|
||||
|
||||
async def _initialize(self) -> None:
|
||||
await super()._initialize()
|
||||
raw = self._ocw_storage._data().get("oauth_metadata")
|
||||
if raw and self.context.oauth_metadata is None:
|
||||
try:
|
||||
from mcp.shared.auth import OAuthMetadata
|
||||
|
||||
self.context.oauth_metadata = OAuthMetadata.model_validate(raw)
|
||||
except Exception:
|
||||
pass # stale/incompatible cache: discovery will refill it
|
||||
if self.context.oauth_metadata is None and self._ocw_storage._data().get(
|
||||
"tokens"
|
||||
):
|
||||
# No cache yet (tokens predate this fix): one best-effort fetch from the
|
||||
# standard well-known location, so the refresh grant can target the real
|
||||
# token endpoint on the very next request. Cached on success; any failure
|
||||
# falls back to the SDK's own (post-401) discovery.
|
||||
try:
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
from mcp.shared.auth import OAuthMetadata
|
||||
|
||||
pr = urlparse(self.context.server_url)
|
||||
url = f"{pr.scheme}://{pr.netloc}/.well-known/oauth-authorization-server"
|
||||
async with httpx.AsyncClient(timeout=10) as c:
|
||||
r = await c.get(url, headers={"Accept": "application/json"})
|
||||
if r.status_code == 200:
|
||||
self.context.oauth_metadata = OAuthMetadata.model_validate(r.json())
|
||||
self._persist_metadata()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _persist_metadata(self) -> None:
|
||||
md = self.context.oauth_metadata
|
||||
if md is not None:
|
||||
try:
|
||||
self._ocw_storage._merge(
|
||||
{"oauth_metadata": md.model_dump(mode="json", exclude_none=True)}
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("could not persist oauth metadata", exc_info=True)
|
||||
|
||||
async def _handle_token_response(self, response: Any) -> None:
|
||||
await super()._handle_token_response(response)
|
||||
self._persist_metadata()
|
||||
|
||||
async def _handle_refresh_response(self, response: Any) -> bool:
|
||||
ok = await super()._handle_refresh_response(response)
|
||||
if ok:
|
||||
self._persist_metadata()
|
||||
return ok
|
||||
|
||||
|
||||
def build_auth(
|
||||
server_name: str,
|
||||
server_url: str,
|
||||
secrets: SecretStore,
|
||||
*,
|
||||
interactive: bool = True,
|
||||
) -> OAuthClientProvider:
|
||||
"""The httpx auth for one OAuth MCP server (pass as streamablehttp_client(auth=…)).
|
||||
|
||||
`interactive=False` still uses stored tokens and silent refresh, but the moment
|
||||
the SDK wants a browser authorization it raises InteractiveAuthRequired instead
|
||||
of opening one — only explicit connect actions pass True.
|
||||
"""
|
||||
metadata = OAuthClientMetadata.model_validate(
|
||||
{
|
||||
"client_name": CLIENT_NAME,
|
||||
"redirect_uris": [redirect_base() + CALLBACK_PATH],
|
||||
"grant_types": ["authorization_code", "refresh_token"],
|
||||
"response_types": ["code"],
|
||||
# Public client: DCR issues no secret a native app could keep anyway.
|
||||
"token_endpoint_auth_method": "none",
|
||||
}
|
||||
)
|
||||
return _MetadataSeededProvider(
|
||||
server_url=server_url,
|
||||
client_metadata=metadata,
|
||||
storage=SecretStoreTokenStorage(server_name, secrets),
|
||||
redirect_handler=_open_browser if interactive else _refuse_browser,
|
||||
callback_handler=_wait_for_callback if interactive else _refuse_callback,
|
||||
)
|
||||
|
||||
|
||||
def has_tokens(server_name: str, secrets: SecretStore) -> bool:
|
||||
return bool((secrets.get(_profile(server_name)) or {}).get("tokens"))
|
||||
|
||||
|
||||
def sign_out(server_name: str, secrets: SecretStore) -> bool:
|
||||
"""Forget tokens AND the DCR registration; next connect runs a fresh flow."""
|
||||
return secrets.delete(_profile(server_name))
|
||||
110
coworker/mcp/tools.py
Normal file
110
coworker/mcp/tools.py
Normal file
@@ -0,0 +1,110 @@
|
||||
"""Turn MCP tools into ToolRegistry-ready callables.
|
||||
|
||||
Each MCP tool becomes a sync callable (so it fits the registry's `execute` contract, which
|
||||
the engine already runs via `asyncio.to_thread`). The callable bridges back to the live
|
||||
async session on the server loop via `run_coroutine_threadsafe`. We attach `ToolMetadata`
|
||||
(category="mcp", `requires_approval` per config) so the PermissionEngine gates it, and an
|
||||
explicit OpenAI schema built straight from the MCP `inputSchema` for fidelity.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
from typing import Any, Awaitable, Callable
|
||||
|
||||
import aisuite as ai
|
||||
|
||||
from .config import MCPServerDef
|
||||
|
||||
CallAsync = Callable[[str, dict[str, Any]], Awaitable[Any]]
|
||||
|
||||
_NAME_OK = re.compile(r"[^a-zA-Z0-9_-]")
|
||||
_MAX_NAME = 64 # OpenAI function-name limit
|
||||
|
||||
|
||||
def tool_name(server: str, tool: str) -> str:
|
||||
"""`mcp__<server>__<tool>`, sanitized to OpenAI's `[A-Za-z0-9_-]{1,64}` rule."""
|
||||
base = f"mcp__{_NAME_OK.sub('_', server)}__{_NAME_OK.sub('_', tool)}"
|
||||
if len(base) > _MAX_NAME:
|
||||
base = base[:_MAX_NAME]
|
||||
return base
|
||||
|
||||
|
||||
def _openai_schema(name: str, mcp_tool: Any) -> dict[str, Any]:
|
||||
params = getattr(mcp_tool, "inputSchema", None) or {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
}
|
||||
description = (getattr(mcp_tool, "description", None) or "")[:1024]
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {"name": name, "description": description, "parameters": params},
|
||||
}
|
||||
|
||||
|
||||
def _filtered(mcp_tools: list[Any], server: MCPServerDef) -> list[Any]:
|
||||
out = mcp_tools
|
||||
if server.include_tools is not None:
|
||||
allow = set(server.include_tools)
|
||||
out = [t for t in out if t.name in allow]
|
||||
if server.exclude_tools:
|
||||
block = set(server.exclude_tools)
|
||||
out = [t for t in out if t.name not in block]
|
||||
return out
|
||||
|
||||
|
||||
def build_callables(
|
||||
server: MCPServerDef,
|
||||
mcp_tools: list[Any],
|
||||
call_async: CallAsync,
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
*,
|
||||
timeout: float = 120.0,
|
||||
) -> list[Callable[..., Any]]:
|
||||
"""Wrap a server's (filtered) MCP tools as registry-ready callables."""
|
||||
callables: list[Callable[..., Any]] = []
|
||||
for mcp_tool in _filtered(mcp_tools, server):
|
||||
name = tool_name(server.name, mcp_tool.name)
|
||||
remote = mcp_tool.name
|
||||
|
||||
def _invoke(_remote: str = remote, **kwargs: Any) -> Any:
|
||||
future = asyncio.run_coroutine_threadsafe(call_async(_remote, kwargs), loop)
|
||||
return future.result(timeout)
|
||||
|
||||
# We attach the schema + metadata explicitly (rather than via `ai.tool`, which would
|
||||
# try to derive a schema from this `**kwargs` wrapper): the registry reads both attrs.
|
||||
_invoke.__name__ = name
|
||||
_invoke.__doc__ = (
|
||||
getattr(mcp_tool, "description", None)
|
||||
or f"MCP tool {remote} from {server.name}"
|
||||
)
|
||||
_invoke.__aisuite_tool_metadata__ = ai.ToolMetadata(
|
||||
name=name,
|
||||
category="mcp",
|
||||
risk_level="medium",
|
||||
capabilities=[server.name],
|
||||
requires_approval=server.requires_approval,
|
||||
)
|
||||
_invoke.__coworker_schema__ = _openai_schema(name, mcp_tool)
|
||||
# OPE-136 finding 4: where this call actually goes, for the approval card's
|
||||
# scope chip. From the server DEF (user-authored config), never from anything
|
||||
# the server itself claims. http → the remote host; stdio → a local process.
|
||||
_invoke.__coworker_mcp_destination__ = {
|
||||
"transport": server.transport,
|
||||
"host": _server_host(server),
|
||||
}
|
||||
callables.append(_invoke)
|
||||
return callables
|
||||
|
||||
|
||||
def _server_host(server: MCPServerDef) -> str:
|
||||
"""The hostname an HTTP server's calls reach (lowercased), "" for stdio/unparseable."""
|
||||
if not server.url:
|
||||
return ""
|
||||
try:
|
||||
from urllib.parse import urlparse
|
||||
|
||||
return (urlparse(server.url).hostname or "").lower()
|
||||
except ValueError: # pragma: no cover - urlparse rarely raises, but fail to ""
|
||||
return ""
|
||||
Reference in New Issue
Block a user