kaiyuan/chatgpt-on-wechat
1
0
Fork 0
mirror of https://github.com/zhayujie/chatgpt-on-wechat.git synced 2026-07-17 11:07:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(mcp): support OAuth authorization for remote MCP servers

Browse Source
This commit is contained in:
zhayujie
2026-07-13 12:00:14 +08:00
parent 42a5cf9538
commit 8c7cda89dc
8 changed files with 792 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

165
agent/tools/mcp/mcp_client.py
View File

@@ -21,6 +21,48 @@ from common.log import logger
_STREAMABLE_HTTP_ALIASES = {"streamable-http", "streamable_http", "streamablehttp", "http"}
# Optional callback invoked after an OAuth authorization completes, so the
# tool manager can bring the newly-authorized server online. Signature:
# reload_fn(server_name: str) -> None. Installed by the tool manager.
_reload_callback = None
def set_reload_callback(fn) -> None:
"""Register a callback fired after a server's OAuth flow succeeds."""
global _reload_callback
_reload_callback = fn
def notify_server_authorized(server_name: str) -> None:
"""Called by the web callback once tokens are stored for a server."""
fn = _reload_callback
if fn is None:
logger.debug(f"[MCP:{server_name}] Authorized but no reload callback registered")
return
try:
fn(server_name)
except Exception as e:
logger.warning(f"[MCP:{server_name}] reload callback failed: {e}")
def _oauth_redirect_uri() -> str:
"""Build the OAuth redirect URI served by the web console callback.
Priority: explicit mcp_oauth_redirect_base config, otherwise the local
web console address (127.0.0.1:<web_port>). Both point at the shared
/mcp/oauth/callback route.
"""
try:
from config import conf
base = (conf().get("mcp_oauth_redirect_base") or "").strip().rstrip("/")
if not base:
port = int(os.environ.get("COW_WEB_PORT") or conf().get("web_port", 9899))
base = f"http://127.0.0.1:{port}"
except Exception:
base = "http://127.0.0.1:9899"
return f"{base}/mcp/oauth/callback"
class McpClient:
"""Single MCP Server client supporting stdio, SSE and Streamable HTTP transports."""
@@ -56,6 +98,13 @@ class McpClient:
self._http_headers: dict = {} # extra headers from user config (e.g. Authorization)
self._http_session_id: Optional[str] = None # Mcp-Session-Id assigned by the server
# OAuth state (streamable-http only). Lazily created when the server
# responds with 401 and the user has not supplied a static token.
self._oauth = None # OAuthHandler instance
# Set to True once a 401 could not be satisfied and the user must
# complete the browser authorization. Callers can surface this state.
self.needs_auth: bool = False
# Shared state
self._next_id = 1
self._id_lock = threading.Lock()
@@ -325,13 +374,118 @@ class McpClient:
if isinstance(extra_headers, dict):
self._http_headers = {str(k): str(v) for k, v in extra_headers.items()}
# Restore any previously stored OAuth credentials for this server so a
# restart reuses the token instead of forcing re-authorization.
self._maybe_load_oauth()
return self._handshake()
# ------------------------------------------------------------------
# OAuth helpers (streamable-http only)
# ------------------------------------------------------------------
def _has_static_auth(self) -> bool:
"""True when the user supplied their own Authorization header."""
return any(k.lower() == "authorization" for k in self._http_headers)
def _maybe_load_oauth(self) -> None:
"""Attach an OAuthHandler when stored credentials exist for this server."""
if self._has_static_auth():
return
try:
from agent.tools.mcp.mcp_oauth import OAuthHandler, load_server_record
except Exception:
return
rec = load_server_record(self.name)
# Only create a handler when we have something to reuse; otherwise it
# is created lazily on the first 401.
if rec.get("access_token") or rec.get("client_id"):
self._oauth = OAuthHandler(
server_name=self.name,
resource_url=self._http_url,
redirect_uri=_oauth_redirect_uri(),
scope=self.config.get("scope", ""),
)
def _current_bearer(self) -> Optional[str]:
"""Return a valid access token, refreshing if needed."""
if self._oauth is None:
return None
return self._oauth.get_valid_access_token()
def _begin_oauth(self, www_authenticate: str = "") -> None:
"""Kick off the OAuth flow after a 401: discover, register, prompt user."""
if self._has_static_auth():
return
try:
from agent.tools.mcp.mcp_oauth import OAuthHandler
except Exception as e:
logger.warning(f"[MCP:{self.name}] OAuth module unavailable: {e}")
return
if self._oauth is None:
self._oauth = OAuthHandler(
server_name=self.name,
resource_url=self._http_url,
redirect_uri=_oauth_redirect_uri(),
scope=self.config.get("scope", ""),
)
if not self._oauth.ensure_registered(www_authenticate):
logger.warning(
f"[MCP:{self.name}] OAuth discovery/registration failed; "
f"cannot authorize automatically"
)
return
auth_url = self._oauth.build_authorization_url()
if not auth_url:
logger.warning(f"[MCP:{self.name}] Failed to build authorization URL")
return
self.needs_auth = True
logger.warning(
f"[MCP:{self.name}] ⚠️ Authorization required. Open this URL in a "
f"browser to authorize, then this server will come online automatically:\n"
f" {auth_url}"
)
# On a machine with a local browser (desktop/dev), open it directly.
if os.environ.get("COW_DESKTOP") == "1" or not os.environ.get("COW_HEADLESS"):
try:
import webbrowser
webbrowser.open(auth_url)
except Exception:
pass
def _streamable_http_send(self, message: dict) -> dict:
"""POST a JSON-RPC request and return the response (JSON or SSE-wrapped)."""
return self._streamable_http_post(message, expect_response=True)
def _streamable_http_post(self, message: dict, expect_response: bool) -> dict:
def _handle_401(self, err, message: dict, expect_response: bool, retried: bool) -> dict:
"""Handle a 401: refresh the token and retry once, else begin OAuth."""
www_auth = ""
try:
www_auth = err.headers.get("WWW-Authenticate", "") or ""
except Exception:
pass
try:
err.read()
except Exception:
pass
# First try a silent refresh with the stored refresh token.
if not retried and self._oauth is not None and self._oauth.refresh():
logger.info(f"[MCP:{self.name}] Token refreshed after 401, retrying")
return self._streamable_http_post(message, expect_response, _retried=True)
# No usable token — start (or restart) the interactive OAuth flow.
self._begin_oauth(www_auth)
raise IOError(
f"[MCP:{self.name}] streamable-http HTTP 401: authorization required "
f"(complete the OAuth flow to enable this server)"
)
def _streamable_http_post(self, message: dict, expect_response: bool, _retried: bool = False) -> dict:
"""
POST a JSON-RPC message over Streamable HTTP.
@@ -351,6 +505,12 @@ class McpClient:
if sid:
headers["Mcp-Session-Id"] = sid
headers.update(self._http_headers)
# Inject OAuth bearer token when we have one (unless the user set a
# static Authorization header, which takes precedence).
if not self._has_static_auth():
token = self._current_bearer()
if token:
headers["Authorization"] = f"Bearer {token}"
req = urllib.request.Request(
self._http_url,
@@ -362,6 +522,9 @@ class McpClient:
try:
resp = urllib.request.urlopen(req, timeout=30)
except urllib.error.HTTPError as e:
# 401 is the spec-compliant "needs authorization" signal.
if e.code == 401 and not self._has_static_auth():
return self._handle_401(e, message, expect_response, _retried)
# Surface the server-provided error body for easier debugging
detail = ""
try:

466
agent/tools/mcp/mcp_oauth.py Normal file
View File

@@ -0,0 +1,466 @@
"""
MCP OAuth 2.1 client (authorization code + PKCE) with zero external deps.
Implements the subset of the MCP authorization spec needed to connect to
remote MCP servers that guard their endpoint behind OAuth (e.g. Xmind):
1. Metadata discovery via RFC 9728 (protected-resource) + RFC 8414
(authorization-server) .well-known documents.
2. Dynamic Client Registration (RFC 7591) to obtain a client_id.
3. PKCE (RFC 7636, S256) authorization-code flow.
4. Token exchange + refresh, persisted to ~/.cow/mcp_oauth.json.
The actual browser round-trip is completed out-of-band: McpClient generates
an authorization URL, the user opens it, and the web console callback
(/mcp/oauth/callback) feeds the returned code back into finish_authorization().
"""
import base64
import hashlib
import json
import os
import secrets
import threading
import time
import urllib.parse
import urllib.request
import urllib.error
from typing import Optional
from common.log import logger
# ------------------------------------------------------------------
# Token store: ~/.cow/mcp_oauth.json {server_name: {...credentials...}}
# ------------------------------------------------------------------
_STORE_LOCK = threading.Lock()
def _store_path() -> str:
base = os.path.expanduser("~/.cow")
try:
os.makedirs(base, exist_ok=True)
except OSError:
pass
return os.path.join(base, "mcp_oauth.json")
def _load_store() -> dict:
path = _store_path()
if not os.path.exists(path):
return {}
try:
with open(path, "r", encoding="utf-8") as f:
data = json.load(f)
return data if isinstance(data, dict) else {}
except Exception as e:
logger.warning(f"[MCP-OAuth] Failed to read token store: {e}")
return {}
def _save_store(store: dict) -> None:
path = _store_path()
tmp = f"{path}.tmp"
try:
with open(tmp, "w", encoding="utf-8") as f:
json.dump(store, f, ensure_ascii=False, indent=2)
os.replace(tmp, path)
# Credentials file: restrict to owner read/write when possible.
try:
os.chmod(path, 0o600)
except OSError:
pass
except Exception as e:
logger.warning(f"[MCP-OAuth] Failed to persist token store: {e}")
def load_server_record(server_name: str) -> dict:
with _STORE_LOCK:
return dict(_load_store().get(server_name, {}))
def save_server_record(server_name: str, record: dict) -> None:
with _STORE_LOCK:
store = _load_store()
store[server_name] = record
_save_store(store)
def clear_server_record(server_name: str) -> None:
with _STORE_LOCK:
store = _load_store()
if server_name in store:
store.pop(server_name, None)
_save_store(store)
# ------------------------------------------------------------------
# Pending authorizations, keyed by the OAuth `state` param.
# Populated when an authorization URL is generated; consumed by the
# web callback when the browser redirects back with ?code&state.
# ------------------------------------------------------------------
_PENDING_LOCK = threading.Lock()
_PENDING: dict = {} # state -> {"handler": OAuthHandler, "created": ts}
_PENDING_TTL = 600 # seconds
def _register_pending(state: str, handler: "OAuthHandler") -> None:
with _PENDING_LOCK:
_prune_pending_locked()
_PENDING[state] = {"handler": handler, "created": time.time()}
def _prune_pending_locked() -> None:
now = time.time()
stale = [s for s, v in _PENDING.items() if now - v["created"] > _PENDING_TTL]
for s in stale:
_PENDING.pop(s, None)
def pop_pending(state: str) -> Optional["OAuthHandler"]:
with _PENDING_LOCK:
_prune_pending_locked()
entry = _PENDING.pop(state, None)
return entry["handler"] if entry else None
def has_pending() -> bool:
with _PENDING_LOCK:
_prune_pending_locked()
return bool(_PENDING)
# ------------------------------------------------------------------
# HTTP helpers (stdlib only)
# ------------------------------------------------------------------
_UA = "CowAgent-MCP-OAuth/1.0"
def _http_get_json(url: str, timeout: int = 15) -> Optional[dict]:
req = urllib.request.Request(url, headers={"Accept": "application/json", "User-Agent": _UA})
try:
with urllib.request.urlopen(req, timeout=timeout) as resp:
raw = resp.read().decode("utf-8")
return json.loads(raw)
except urllib.error.HTTPError as e:
logger.debug(f"[MCP-OAuth] GET {url} -> HTTP {e.code}")
return None
except Exception as e:
logger.debug(f"[MCP-OAuth] GET {url} failed: {e}")
return None
def _http_post_form(url: str, fields: dict, timeout: int = 20) -> dict:
body = urllib.parse.urlencode(fields).encode("utf-8")
req = urllib.request.Request(
url,
data=body,
method="POST",
headers={
"Content-Type": "application/x-www-form-urlencoded",
"Accept": "application/json",
"User-Agent": _UA,
},
)
with urllib.request.urlopen(req, timeout=timeout) as resp:
raw = resp.read().decode("utf-8")
return json.loads(raw) if raw else {}
def _http_post_json(url: str, payload: dict, timeout: int = 20) -> dict:
body = json.dumps(payload).encode("utf-8")
req = urllib.request.Request(
url,
data=body,
method="POST",
headers={
"Content-Type": "application/json",
"Accept": "application/json",
"User-Agent": _UA,
},
)
with urllib.request.urlopen(req, timeout=timeout) as resp:
raw = resp.read().decode("utf-8")
return json.loads(raw) if raw else {}
# ------------------------------------------------------------------
# Discovery (RFC 9728 + RFC 8414)
# ------------------------------------------------------------------
def _origin(url: str) -> str:
p = urllib.parse.urlparse(url)
return f"{p.scheme}://{p.netloc}"
def discover_metadata(resource_url: str, www_authenticate: str = "") -> Optional[dict]:
"""
Resolve the authorization server metadata for a protected MCP resource.
Returns a dict with at least authorization_endpoint + token_endpoint,
plus registration_endpoint when the server supports DCR. Returns None
when discovery fails.
"""
as_metadata_url = _parse_resource_metadata_url(www_authenticate)
# 1) Protected-resource metadata (RFC 9728) to locate the auth server.
auth_server = None
prm = None
if as_metadata_url:
prm = _http_get_json(as_metadata_url)
if prm is None:
origin = _origin(resource_url)
prm = _http_get_json(f"{origin}/.well-known/oauth-protected-resource")
if prm and isinstance(prm.get("authorization_servers"), list) and prm["authorization_servers"]:
auth_server = prm["authorization_servers"][0]
# 2) Authorization-server metadata (RFC 8414). Fall back to the resource
# origin when the resource did not advertise a separate auth server.
base = auth_server or _origin(resource_url)
asm = _fetch_as_metadata(base)
if not asm:
return None
if not asm.get("authorization_endpoint") or not asm.get("token_endpoint"):
logger.warning("[MCP-OAuth] Authorization server metadata missing required endpoints")
return None
# Derive the scope to request. Prefer the resource's required_scopes
# (RFC 9728), then its scopes_supported, then the auth server's
# scopes_supported. Stored so callers don't have to configure it.
discovered_scope = ""
if prm:
scopes = prm.get("required_scopes") or prm.get("scopes_supported")
if isinstance(scopes, list) and scopes:
discovered_scope = " ".join(str(s) for s in scopes)
if not discovered_scope and isinstance(asm.get("scopes_supported"), list) and asm["scopes_supported"]:
discovered_scope = " ".join(str(s) for s in asm["scopes_supported"])
if discovered_scope:
asm["_discovered_scope"] = discovered_scope
return asm
def _parse_resource_metadata_url(www_authenticate: str) -> Optional[str]:
"""Extract resource_metadata="..." from a WWW-Authenticate: Bearer header."""
if not www_authenticate:
return None
# naive but sufficient parse for `resource_metadata="URL"`
marker = "resource_metadata="
idx = www_authenticate.find(marker)
if idx < 0:
return None
rest = www_authenticate[idx + len(marker):].strip()
if rest.startswith('"'):
end = rest.find('"', 1)
return rest[1:end] if end > 0 else None
# unquoted, up to comma/space
for sep in (",", " "):
if sep in rest:
rest = rest.split(sep, 1)[0]
return rest or None
def _fetch_as_metadata(base: str) -> Optional[dict]:
"""Try both RFC 8414 and OIDC well-known locations."""
base = base.rstrip("/")
candidates = [
f"{base}/.well-known/oauth-authorization-server",
f"{base}/.well-known/openid-configuration",
]
for url in candidates:
data = _http_get_json(url)
if data and data.get("authorization_endpoint"):
return data
return None
# ------------------------------------------------------------------
# PKCE
# ------------------------------------------------------------------
def _b64url(data: bytes) -> str:
return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")
def _make_pkce() -> tuple:
verifier = _b64url(secrets.token_bytes(32))
challenge = _b64url(hashlib.sha256(verifier.encode("ascii")).digest())
return verifier, challenge
# ------------------------------------------------------------------
# OAuthHandler: per-server OAuth state machine
# ------------------------------------------------------------------
class OAuthHandler:
"""Drives the OAuth flow and token lifecycle for a single MCP server."""
def __init__(self, server_name: str, resource_url: str, redirect_uri: str,
scope: str = "", client_name: str = "CowAgent"):
self.server_name = server_name
self.resource_url = resource_url
self.redirect_uri = redirect_uri
self.scope = scope
self.client_name = client_name
rec = load_server_record(server_name)
self.metadata: dict = rec.get("metadata", {})
self.client_id: Optional[str] = rec.get("client_id")
self.client_secret: Optional[str] = rec.get("client_secret")
self.access_token: Optional[str] = rec.get("access_token")
self.refresh_token: Optional[str] = rec.get("refresh_token")
self.expires_at: float = float(rec.get("expires_at", 0) or 0)
self._verifier: Optional[str] = None
# --- persistence -------------------------------------------------
def _persist(self) -> None:
save_server_record(self.server_name, {
"resource_url": self.resource_url,
"metadata": self.metadata,
"client_id": self.client_id,
"client_secret": self.client_secret,
"access_token": self.access_token,
"refresh_token": self.refresh_token,
"expires_at": self.expires_at,
})
# --- token access ------------------------------------------------
def get_valid_access_token(self, leeway: int = 60) -> Optional[str]:
"""Return a usable access token, refreshing proactively when near expiry."""
if not self.access_token:
return None
if self.expires_at and time.time() >= self.expires_at - leeway:
if not self.refresh():
return None
return self.access_token
def refresh(self) -> bool:
"""Refresh the access token using the stored refresh token."""
if not self.refresh_token or not self.metadata.get("token_endpoint"):
return False
fields = {
"grant_type": "refresh_token",
"refresh_token": self.refresh_token,
"client_id": self.client_id or "",
}
if self.client_secret:
fields["client_secret"] = self.client_secret
try:
resp = _http_post_form(self.metadata["token_endpoint"], fields)
except Exception as e:
logger.warning(f"[MCP-OAuth:{self.server_name}] refresh failed: {e}")
return False
return self._absorb_token_response(resp)
# --- authorization-code flow ------------------------------------
def ensure_registered(self, www_authenticate: str = "") -> bool:
"""Discover metadata + register a client if not already done."""
if not self.metadata.get("authorization_endpoint"):
meta = discover_metadata(self.resource_url, www_authenticate)
if not meta:
return False
self.metadata = meta
# Adopt the scope discovered from metadata when the user didn't set one.
if not self.scope and self.metadata.get("_discovered_scope"):
self.scope = self.metadata["_discovered_scope"]
logger.info(f"[MCP-OAuth:{self.server_name}] Using discovered scope: {self.scope}")
if not self.client_id:
if not self._register_client():
return False
self._persist()
return True
def _register_client(self) -> bool:
reg_endpoint = self.metadata.get("registration_endpoint")
if not reg_endpoint:
logger.warning(
f"[MCP-OAuth:{self.server_name}] No registration_endpoint; "
f"DCR unavailable. Provide client_id manually."
)
return False
payload = {
"client_name": self.client_name,
"redirect_uris": [self.redirect_uri],
"grant_types": ["authorization_code", "refresh_token"],
"response_types": ["code"],
"token_endpoint_auth_method": "none",
}
if self.scope:
payload["scope"] = self.scope
try:
resp = _http_post_json(reg_endpoint, payload)
except Exception as e:
logger.warning(f"[MCP-OAuth:{self.server_name}] DCR failed: {e}")
return False
client_id = resp.get("client_id")
if not client_id:
logger.warning(f"[MCP-OAuth:{self.server_name}] DCR returned no client_id")
return False
self.client_id = client_id
self.client_secret = resp.get("client_secret")
logger.info(f"[MCP-OAuth:{self.server_name}] Registered client_id={client_id}")
return True
def build_authorization_url(self) -> Optional[str]:
"""Create an authorization URL and register this handler as pending."""
if not self.metadata.get("authorization_endpoint") or not self.client_id:
return None
self._verifier, challenge = _make_pkce()
state = secrets.token_urlsafe(24)
params = {
"response_type": "code",
"client_id": self.client_id,
"redirect_uri": self.redirect_uri,
"code_challenge": challenge,
"code_challenge_method": "S256",
"state": state,
}
if self.scope:
params["scope"] = self.scope
# Advertise the resource we intend to access (RFC 8707).
params["resource"] = self.resource_url
_register_pending(state, self)
return f"{self.metadata['authorization_endpoint']}?{urllib.parse.urlencode(params)}"
def finish_authorization(self, code: str) -> bool:
"""Exchange an authorization code for tokens."""
if not self.metadata.get("token_endpoint") or not self._verifier:
return False
fields = {
"grant_type": "authorization_code",
"code": code,
"redirect_uri": self.redirect_uri,
"client_id": self.client_id or "",
"code_verifier": self._verifier,
"resource": self.resource_url,
}
if self.client_secret:
fields["client_secret"] = self.client_secret
try:
resp = _http_post_form(self.metadata["token_endpoint"], fields)
except Exception as e:
logger.warning(f"[MCP-OAuth:{self.server_name}] token exchange failed: {e}")
return False
ok = self._absorb_token_response(resp)
self._verifier = None
return ok
def _absorb_token_response(self, resp: dict) -> bool:
access = resp.get("access_token")
if not access:
logger.warning(f"[MCP-OAuth:{self.server_name}] token response missing access_token: {resp}")
return False
self.access_token = access
if resp.get("refresh_token"):
self.refresh_token = resp["refresh_token"]
expires_in = resp.get("expires_in")
self.expires_at = time.time() + int(expires_in) if expires_in else 0
self._persist()
logger.info(f"[MCP-OAuth:{self.server_name}] Access token stored")
return True