Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(mcp): add browser OAuth for remote servers (#5316)

Browse Source
This commit is contained in:
chengyongru
2026-08-10 23:44:37 +08:00
committed by GitHub
parent b3b0517611
commit 8e77f3f8a4
33 changed files with 4099 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/tools/mcp.py
+78 -13
View File
@@ -38,6 +38,7 @@ if TYPE_CHECKING:
from mcp.types import Prompt, Resource
from mcp.types import Tool as MCPToolDefinition
from nanobot.agent.tools.mcp_oauth import MCPOAuthHandlers
from nanobot.config.schema import MCPServerConfig
# Transient connection errors that warrant a single retry.
@@ -184,6 +185,25 @@ def _is_transient(exc: BaseException) -> bool:
return type(exc).__name__ in _TRANSIENT_EXC_NAMES
def _is_transient_connection_failure(exc: BaseException) -> bool:
if isinstance(exc, BaseExceptionGroup):
group = cast(BaseExceptionGroup[BaseException], exc)
return bool(group.exceptions) and all(
_is_transient_connection_failure(nested) for nested in group.exceptions
)
return isinstance(exc, (httpx.ConnectError, httpx.ConnectTimeout)) or _is_transient(exc)
def _log_mcp_connection_failure(name: str, exc: BaseException, hint: str = "") -> None:
if _is_transient_connection_failure(exc):
logger.warning("MCP server '{}': transient connection failure", name)
logger.opt(exception=exc).debug(
"MCP server '{}' transient connection failure details", name
)
return
logger.opt(exception=exc).error("MCP server '{}': failed to connect: {}", name, hint)
def _is_session_terminated(exc: BaseException) -> bool:
"""Return True when the MCP SDK reports a dead client session."""
if _is_transient(exc):
@@ -961,7 +981,10 @@ class MCPPromptWrapper(_MCPWrapperBase):
async def connect_mcp_servers(
mcp_servers: "dict[str, MCPServerConfig]", registry: ToolRegistry
mcp_servers: "dict[str, MCPServerConfig]",
registry: ToolRegistry,
*,
oauth_handlers: Mapping[str, "MCPOAuthHandlers"] | None = None,
) -> dict[str, MCPConnection]:
"""Connect to configured MCP servers and register their tools, resources, prompts.
@@ -1001,6 +1024,29 @@ async def connect_mcp_servers(
)
return False
oauth_auth: httpx.Auth | None = None
if cfg.auth == "oauth":
if transport_type not in {"sse", "streamableHttp"}:
logger.warning(
"MCP server '{}': OAuth requires an SSE or Streamable HTTP transport",
name,
)
return False
from nanobot.agent.tools.mcp_oauth import (
MCPAuthorizationRequiredError,
create_mcp_oauth_auth,
)
try:
oauth_auth = await create_mcp_oauth_auth(
name,
cfg.url,
(oauth_handlers or {}).get(name),
)
except MCPAuthorizationRequiredError:
logger.info("MCP server '{}': waiting for browser authorization", name)
return False
if transport_type == "stdio":
command, args, env = _normalize_windows_stdio_command(
cfg.command,
@@ -1038,22 +1084,30 @@ async def connect_mcp_servers(
**_pinned_transport_kwargs(),
)
sse_kwargs: dict[str, Any] = {
"httpx_client_factory": httpx_client_factory,
}
if oauth_auth is not None:
sse_kwargs["auth"] = oauth_auth
read, write = await server_stack.enter_async_context(
sse_client(cfg.url, httpx_client_factory=httpx_client_factory)
sse_client(cfg.url, **sse_kwargs)
)
elif transport_type == "streamableHttp":
if not await _probe_http_url(cfg.url):
logger.warning("MCP server '{}': {} unreachable, skipping", name, _redact_url(cfg.url))
return False
http_client_kwargs: dict[str, Any] = {
"headers": cfg.headers or None,
"event_hooks": {"request": [_validate_mcp_request_url]},
"follow_redirects": True,
"timeout": httpx.Timeout(30.0, connect=10.0),
**_pinned_transport_kwargs(),
}
if oauth_auth is not None:
http_client_kwargs["auth"] = oauth_auth
http_client = await server_stack.enter_async_context(
httpx.AsyncClient(
headers=cfg.headers or None,
event_hooks={"request": [_validate_mcp_request_url]},
follow_redirects=True,
timeout=httpx.Timeout(30.0, connect=10.0),
**_pinned_transport_kwargs(),
)
httpx.AsyncClient(**http_client_kwargs)
)
read, write, _ = await server_stack.enter_async_context(
streamable_http_client(cfg.url, http_client=http_client)
@@ -1182,7 +1236,7 @@ async def connect_mcp_servers(
" Hint: this looks like stdio protocol pollution. Make sure the MCP server writes "
"only JSON-RPC to stdout and sends logs/debug output to stderr instead."
)
logger.exception("MCP server '{}': failed to connect: {}", name, hint)
_log_mcp_connection_failure(name, e, hint)
return False
async def connect_single_server(
@@ -1229,7 +1283,7 @@ async def connect_mcp_servers(
try:
result = await connect_single_server(name, cfg)
except Exception as e:
logger.exception("MCP server '{}' connection failed: {}", name, e)
_log_mcp_connection_failure(name, e)
continue
if result[1] is not None:
server_stacks[result[0]] = result[1]
@@ -1302,6 +1356,13 @@ async def reload_servers(state: Any, registry: ToolRegistry) -> dict[str, Any]:
current_servers = dict(state._mcp_servers)
current_names = set(current_servers)
next_names = set(next_servers)
from nanobot.agent.tools.mcp_oauth import mcp_oauth_has_credentials
authorization_pending = {
name
for name, cfg in next_servers.items()
if cfg.auth == "oauth" and not mcp_oauth_has_credentials(name, cfg.url)
}
removed = sorted(current_names - next_names)
added = sorted(next_names - current_names)
changed = sorted(
@@ -1319,9 +1380,13 @@ async def reload_servers(state: Any, registry: ToolRegistry) -> dict[str, Any]:
retry_missing = sorted(
name
for name in next_names
if name not in state._mcp_stacks and name not in set(added) | set(changed)
if name not in state._mcp_stacks
and name not in set(added) | set(changed)
and name not in authorization_pending
)
to_connect_names = sorted(
(set(added) | set(changed) | set(retry_missing)) - authorization_pending
)
to_connect_names = sorted(set(added) | set(changed) | set(retry_missing))
to_connect = {name: next_servers[name] for name in to_connect_names}
connected: dict[str, MCPConnection] = {}
if to_connect:
nanobot/agent/tools/mcp_oauth.py
+401
View File
@@ -0,0 +1,401 @@
"""OAuth support for remote MCP servers.
This module intentionally owns MCP OAuth end to end. Provider OAuth has a
different lifecycle and storage contract, so sharing a higher-level workflow
would couple unrelated extension boundaries.
"""
from __future__ import annotations
import asyncio
import hashlib
import json
import os
import secrets
from collections.abc import Awaitable, Callable
from contextlib import suppress
from dataclasses import dataclass
from pathlib import Path
from typing import Any, TypedDict, cast
from filelock import FileLock
from loguru import logger
from mcp.client.auth import OAuthClientProvider
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken
from pydantic import AnyHttpUrl, AnyUrl
from nanobot.config.paths import get_data_dir
from nanobot.utils.helpers import _write_text_atomic # pyright: ignore[reportPrivateUsage]
MCP_OAUTH_CALLBACK_PATH = "/auth/mcp/callback"
_STORE_VERSION = 1
_STORE_LOCK_TIMEOUT_S = 15
_DEFAULT_REDIRECT_URI = f"http://127.0.0.1{MCP_OAUTH_CALLBACK_PATH}"
_CLIENT_URI = AnyHttpUrl("https://github.com/HKUDS/nanobot")
_LOGO_URI = AnyHttpUrl(
"https://raw.githubusercontent.com/HKUDS/nanobot/main/"
"webui/public/brand/nanobot_apple_touch.png"
)
class _StoredServer(TypedDict, total=False):
server_fingerprint: str
write_lease: str
tokens: dict[str, Any]
client_info: dict[str, Any]
redirect_uri: str
class _CredentialStore(TypedDict):
version: int
servers: dict[str, _StoredServer]
generations: dict[str, str]
class MCPAuthorizationRequiredError(RuntimeError):
"""Raised when a background MCP connection needs interactive authorization."""
@dataclass(frozen=True)
class MCPOAuthHandlers:
"""Browser callbacks supplied only for a user-initiated OAuth attempt."""
redirect_uri: str
redirect_handler: Callable[[str], Awaitable[None]]
callback_handler: Callable[[], Awaitable[tuple[str, str | None]]]
reset_credentials: bool = False
def _store_path() -> Path:
return get_data_dir() / "auth" / "mcp.json"
def _server_fingerprint(server_url: str) -> str:
return hashlib.sha256(server_url.strip().encode("utf-8")).hexdigest()
def _empty_store() -> _CredentialStore:
return {"version": _STORE_VERSION, "servers": {}, "generations": {}}
def _stored_server(value: object) -> _StoredServer | None:
if not isinstance(value, dict):
return None
raw = cast(dict[object, object], value)
entry: _StoredServer = {}
fingerprint = raw.get("server_fingerprint")
if isinstance(fingerprint, str):
entry["server_fingerprint"] = fingerprint
write_lease = raw.get("write_lease")
if isinstance(write_lease, str) and write_lease:
entry["write_lease"] = write_lease
redirect_uri = raw.get("redirect_uri")
if isinstance(redirect_uri, str):
entry["redirect_uri"] = redirect_uri
tokens = raw.get("tokens")
if isinstance(tokens, dict):
token_values = cast(dict[object, object], tokens)
if all(isinstance(key, str) for key in token_values):
entry["tokens"] = cast(dict[str, Any], token_values)
client_info = raw.get("client_info")
if isinstance(client_info, dict):
client_values = cast(dict[object, object], client_info)
if all(isinstance(key, str) for key in client_values):
entry["client_info"] = cast(dict[str, Any], client_values)
return entry
def _read_store_unlocked(path: Path) -> _CredentialStore:
try:
raw = cast(object, json.loads(path.read_text(encoding="utf-8")))
except FileNotFoundError:
return _empty_store()
except (OSError, ValueError, TypeError) as exc:
logger.warning("Could not read MCP OAuth credentials: {}", type(exc).__name__)
return _empty_store()
if not isinstance(raw, dict):
return _empty_store()
payload = cast(dict[object, object], raw)
raw_servers = payload.get("servers")
if not isinstance(raw_servers, dict):
return _empty_store()
servers: dict[str, _StoredServer] = {}
for name, value in cast(dict[object, object], raw_servers).items():
entry = _stored_server(value)
if isinstance(name, str) and entry is not None:
servers[name] = entry
generations: dict[str, str] = {}
raw_generations = payload.get("generations")
if isinstance(raw_generations, dict):
for name, value in cast(dict[object, object], raw_generations).items():
if isinstance(name, str) and isinstance(value, str) and value:
generations[name] = value
return {
"version": _STORE_VERSION,
"servers": servers,
"generations": generations,
}
def _with_store_lock(path: Path) -> FileLock:
path.parent.mkdir(parents=True, exist_ok=True)
return FileLock(str(path.with_suffix(".lock")), timeout=_STORE_LOCK_TIMEOUT_S)
def _write_store_unlocked(path: Path, payload: _CredentialStore) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with suppress(OSError):
os.chmod(path.parent, 0o700)
_write_text_atomic(path, json.dumps(payload, indent=2, ensure_ascii=False))
with suppress(OSError):
os.chmod(path, 0o600)
class MCPOAuthStorage:
"""Persistent MCP SDK token storage, isolated by config name and server URL."""
def __init__(self, server_name: str, server_url: str) -> None:
self.server_name = server_name
self.server_fingerprint = _server_fingerprint(server_url)
self._observed_generation = self._read_generation_sync()
self._write_lease: str | None = None
def _read_generation_sync(self) -> str | None:
path = _store_path()
if not path.exists():
return None
# Writes replace the whole file atomically, so this observes either side
# of a concurrent deletion without blocking the async connection path.
return _read_store_unlocked(path)["generations"].get(self.server_name)
def _generation_is_current(self, payload: _CredentialStore) -> bool:
return payload["generations"].get(self.server_name) == self._observed_generation
def _entry_unlocked(self, payload: _CredentialStore) -> _StoredServer | None:
servers = payload["servers"]
entry = servers.get(self.server_name)
if entry is None or entry.get("server_fingerprint") != self.server_fingerprint:
return None
return entry
def _bind_entry_unlocked(
self,
payload: _CredentialStore,
*,
create: bool,
) -> tuple[_StoredServer | None, bool]:
if not self._generation_is_current(payload):
return None, False
entry = self._entry_unlocked(payload)
if self._write_lease is not None:
if entry is None or entry.get("write_lease") != self._write_lease:
return None, False
return entry, False
if entry is None:
if not create:
return None, False
self._write_lease = secrets.token_urlsafe(24)
entry = _StoredServer(
server_fingerprint=self.server_fingerprint,
write_lease=self._write_lease,
)
payload["servers"][self.server_name] = entry
return entry, True
write_lease = entry.get("write_lease")
changed = not isinstance(write_lease, str) or not write_lease
if changed:
write_lease = secrets.token_urlsafe(24)
entry["write_lease"] = write_lease
self._write_lease = write_lease
return entry, changed
def _read_entry_sync(self) -> _StoredServer | None:
path = _store_path()
with _with_store_lock(path):
payload = _read_store_unlocked(path)
entry, changed = self._bind_entry_unlocked(payload, create=False)
if changed:
_write_store_unlocked(path, payload)
return entry
def _update_entry_sync(
self,
update: Callable[[_StoredServer], None],
*,
create: bool = True,
claim: bool = False,
) -> bool:
path = _store_path()
with _with_store_lock(path):
payload = _read_store_unlocked(path)
if claim:
# A browser flow owns subsequent SDK writes until another flow
# claims the entry or the configured server is removed.
if not self._generation_is_current(payload):
logger.info(
"Ignored stale MCP OAuth credential claim for '{}'",
self.server_name,
)
return False
entry = self._entry_unlocked(payload)
if entry is None:
entry = _StoredServer(server_fingerprint=self.server_fingerprint)
payload["servers"][self.server_name] = entry
self._write_lease = secrets.token_urlsafe(24)
entry["write_lease"] = self._write_lease
else:
entry, _ = self._bind_entry_unlocked(payload, create=create)
if entry is None:
if self._write_lease is not None:
logger.info(
"Ignored stale MCP OAuth credential update for '{}'",
self.server_name,
)
return False
update(entry)
payload["version"] = _STORE_VERSION
_write_store_unlocked(path, payload)
return True
async def get_tokens(self) -> OAuthToken | None:
entry = await asyncio.to_thread(self._read_entry_sync)
raw = entry.get("tokens") if entry is not None else None
if not isinstance(raw, dict):
return None
try:
return OAuthToken.model_validate(raw)
except (ValueError, TypeError):
logger.warning("Ignoring invalid MCP OAuth tokens for '{}'", self.server_name)
return None
async def set_tokens(self, tokens: OAuthToken) -> None:
raw = tokens.model_dump(mode="json", exclude_none=True)
def update(entry: _StoredServer) -> None:
entry["tokens"] = raw
await asyncio.to_thread(self._update_entry_sync, update)
async def clear_tokens(self) -> None:
def update(entry: _StoredServer) -> None:
entry.pop("tokens", None)
await asyncio.to_thread(self._update_entry_sync, update, create=False)
async def get_client_info(self) -> OAuthClientInformationFull | None:
entry = await asyncio.to_thread(self._read_entry_sync)
raw = entry.get("client_info") if entry is not None else None
if not isinstance(raw, dict):
return None
try:
return OAuthClientInformationFull.model_validate(raw)
except (ValueError, TypeError):
logger.warning("Ignoring invalid MCP OAuth client info for '{}'", self.server_name)
return None
async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
raw = client_info.model_dump(mode="json", exclude_none=True)
def update(entry: _StoredServer) -> None:
entry["client_info"] = raw
await asyncio.to_thread(self._update_entry_sync, update)
async def redirect_uri(self) -> str | None:
entry = await asyncio.to_thread(self._read_entry_sync)
value = entry.get("redirect_uri") if entry is not None else None
return value if isinstance(value, str) and value else None
async def prepare_redirect_uri(self, redirect_uri: str, *, reset: bool = False) -> None:
def update(entry: _StoredServer) -> None:
changed = entry.get("redirect_uri") != redirect_uri
if reset:
entry.pop("tokens", None)
entry.pop("client_info", None)
elif changed:
# Dynamic registrations bind a client to its redirect URI.
entry.pop("client_info", None)
entry["redirect_uri"] = redirect_uri
claimed = await asyncio.to_thread(self._update_entry_sync, update, claim=True)
if not claimed:
raise MCPAuthorizationRequiredError("MCP authorization was cancelled")
def has_credentials(self) -> bool:
entry = self._read_entry_sync()
raw_tokens = entry.get("tokens") if entry is not None else None
if not isinstance(raw_tokens, dict):
return False
tokens = cast(dict[str, object], raw_tokens)
access_token = tokens.get("access_token")
return isinstance(access_token, str) and bool(access_token)
async def _missing_callback() -> tuple[str, str | None]:
raise MCPAuthorizationRequiredError("MCP server requires browser authorization")
async def create_mcp_oauth_auth(
server_name: str,
server_url: str,
handlers: MCPOAuthHandlers | None = None,
) -> OAuthClientProvider:
"""Build the official MCP SDK OAuth provider for one configured server."""
storage = MCPOAuthStorage(server_name, server_url)
if handlers is not None:
await storage.prepare_redirect_uri(
handlers.redirect_uri,
reset=handlers.reset_credentials,
)
redirect_uri = handlers.redirect_uri
redirect_handler = handlers.redirect_handler
callback_handler = handlers.callback_handler
else:
if not await asyncio.to_thread(storage.has_credentials):
# Do not perform discovery or dynamic registration from a background
# startup. Interactive OAuth begins only after an explicit user action.
raise MCPAuthorizationRequiredError("MCP server requires browser authorization")
redirect_uri = await storage.redirect_uri() or _DEFAULT_REDIRECT_URI
async def authorization_required(_authorization_url: str) -> None:
await storage.clear_tokens()
raise MCPAuthorizationRequiredError("MCP server requires browser authorization")
redirect_handler = authorization_required
callback_handler = _missing_callback
metadata = OAuthClientMetadata(
redirect_uris=[AnyUrl(redirect_uri)],
token_endpoint_auth_method="none",
client_name="nanobot",
client_uri=_CLIENT_URI,
logo_uri=_LOGO_URI,
software_id="https://github.com/HKUDS/nanobot",
)
return OAuthClientProvider(
server_url,
metadata,
storage,
redirect_handler=redirect_handler,
callback_handler=callback_handler,
timeout=300,
)
def mcp_oauth_has_credentials(server_name: str, server_url: str) -> bool:
"""Return whether this exact configured MCP instance has an access token."""
return MCPOAuthStorage(server_name, server_url).has_credentials()
def delete_mcp_oauth_credentials(server_name: str) -> bool:
"""Delete credentials for one config name without touching other MCP instances."""
path = _store_path()
with _with_store_lock(path):
payload = _read_store_unlocked(path)
servers = payload["servers"]
removed = servers.pop(server_name, None) is not None
# Rotate even when no entry exists so a flow created before removal cannot
# claim the name later and resurrect credentials.
payload["generations"][server_name] = secrets.token_urlsafe(24)
_write_store_unlocked(path, payload)
return removed