langchain-agentx-python 2.3.5__py3-none-any.whl → 2.3.6__py3-none-any.whl

langchain_agentx/__init__.py

@@ -11,7 +11,7 @@ from langchain_agentx import create_loop_agent
 ```
 """
 
-__version__ = "2.3.5"
+__version__ = "2.3.6"
 
 from .mcp.call import McpToolCaller
 from .mcp.client_manager import McpClientManager, McpConnectedClient

langchain_agentx/mcp/auth/__init__.py

@@ -1,5 +1,30 @@
 """mcp.auth — OAuth 协议面（无浏览器）。"""
 
-from langchain_agentx.mcp.auth.oauth_session import McpOAuthSession, McpTokenStore
+from langchain_agentx.mcp.auth.authorization_flow import (
+    McpOAuthAuthorizationFlow,
+    McpOAuthStartResult,
+)
+from langchain_agentx.mcp.auth.errors import McpOAuthError, McpOAuthErrorReason
+from langchain_agentx.mcp.auth.oauth_session import (
+    McpOAuthSession,
+    McpOAuthTokenRecord,
+    McpTokenStore,
+)
+from langchain_agentx.mcp.auth.pkce import McpPkcePair
+from langchain_agentx.mcp.auth.token_store import (
+    McpFileTokenStore,
+    default_token_store_path,
+)
 
-__all__ = ["McpOAuthSession", "McpTokenStore"]
+__all__ = [
+    "McpOAuthAuthorizationFlow",
+    "McpOAuthError",
+    "McpOAuthErrorReason",
+    "McpOAuthSession",
+    "McpOAuthStartResult",
+    "McpOAuthTokenRecord",
+    "McpPkcePair",
+    "McpTokenStore",
+    "McpFileTokenStore",
+    "default_token_store_path",
+]

langchain_agentx/mcp/auth/authorization_flow.py

@@ -0,0 +1,271 @@
+"""OAuth 授权码 + PKCE 主编排。
+
+职责：
+    ``start``：endpoint 发现 + PKCE + state → authorize URL；
+    ``complete``：callback 校验（error/state/code）→ code 交换 → 写 Store。
+
+链路位置：
+    对标 CC ``performMCPOAuthFlow`` 的协议切片（无浏览器 / 无本地端口监听 / 无 XAA）。
+    ``McpOAuthSession`` 门面调用；宿主负责打开浏览器与收 callback。
+
+当前裁剪：
+    无动态 client registration（client_id 由调用方给）。
+    refresh 仅一次、无循环重试（Manager 侧控制）；pending 按 server 名单槽，
+    二次 ``start`` 覆盖（旧 state 失效）。
+"""
+
+from __future__ import annotations
+
+import os
+import secrets
+import time
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Any
+from urllib.parse import parse_qs, urlencode, urlsplit
+
+import httpx
+
+from langchain_agentx.mcp.auth.discovery import McpOAuthDiscovery, McpOAuthEndpoints
+from langchain_agentx.mcp.auth.errors import McpOAuthError
+from langchain_agentx.mcp.auth.pkce import McpPkcePair
+from langchain_agentx.mcp.auth.token_store import (
+    McpOAuthTokenRecord,
+    McpTokenStoreProtocol,
+    build_server_key,
+)
+from langchain_agentx.protocol.mcp.types import McpServerSpec
+
+_ENV_CALLBACK_PORT = "MCP_OAUTH_CALLBACK_PORT"
+_DEFAULT_CALLBACK_PORT = 3118  # CC REDIRECT_PORT_FALLBACK
+_REDIRECT_PATH = "/callback"
+
+
+def default_callback_redirect_uri() -> str:
+    """冻结形态：``http://localhost:{port}/callback``（env 可覆盖端口）。"""
+    port = int(os.environ.get(_ENV_CALLBACK_PORT, "") or 0) or _DEFAULT_CALLBACK_PORT
+    return f"http://localhost:{port}{_REDIRECT_PATH}"
+
+
+@dataclass(frozen=True)
+class McpOAuthStartResult:
+    """``start_authorization`` 产物；宿主拿 URL 去开浏览器。"""
+
+    authorization_url: str
+    state: str
+    redirect_uri: str
+    server_key: str
+
+
+@dataclass
+class _PendingFlow:
+    server_name: str
+    server_key: str
+    endpoints: McpOAuthEndpoints
+    pkce: McpPkcePair
+    state: str
+    redirect_uri: str
+    client_id: str | None
+    scope: str | None
+
+
+class McpOAuthTokenExchanger:
+    """POST token endpoint（form → JSON）。测试注入同形 fake。"""
+
+    def __init__(self, *, timeout: float = 30.0) -> None:
+        self._timeout = timeout
+
+    async def exchange(self, endpoint: str, form: Mapping[str, str]) -> Mapping[str, Any]:
+        async with httpx.AsyncClient(timeout=self._timeout, follow_redirects=True) as client:
+            response = await client.post(
+                endpoint, data=dict(form), headers={"Accept": "application/json"}
+            )
+            response.raise_for_status()
+            payload = response.json()
+        if not isinstance(payload, Mapping):
+            raise ValueError("token endpoint returned non-object JSON")
+        return payload
+
+
+class McpOAuthAuthorizationFlow:
+    """一个 server 至多一个 pending flow；完成/拒绝/交换失败后清槽。"""
+
+    def __init__(
+        self,
+        store: McpTokenStoreProtocol,
+        *,
+        discovery: McpOAuthDiscovery | None = None,
+        exchanger: McpOAuthTokenExchanger | None = None,
+        time_fn: Callable[[], float] = time.time,
+    ) -> None:
+        self._store = store
+        self._discovery = discovery or McpOAuthDiscovery()
+        self._exchanger = exchanger or McpOAuthTokenExchanger()
+        self._time_fn = time_fn
+        self._pending: dict[str, _PendingFlow] = {}
+
+    @property
+    def store(self) -> McpTokenStoreProtocol:
+        return self._store
+
+    async def start(
+        self,
+        spec: McpServerSpec,
+        *,
+        redirect_uri: str | None = None,
+        client_id: str | None = None,
+        scope: str | None = None,
+    ) -> McpOAuthStartResult:
+        _ensure_remote(spec)
+        endpoints = await self._discovery.resolve(spec)
+        pkce = McpPkcePair.generate()
+        state = secrets.token_urlsafe(24)
+        resolved_redirect_uri = redirect_uri or default_callback_redirect_uri()
+        query: dict[str, str] = {
+            "response_type": "code",
+            "redirect_uri": resolved_redirect_uri,
+            "state": state,
+            "code_challenge": pkce.code_challenge,
+            "code_challenge_method": pkce.code_challenge_method,
+        }
+        if client_id:
+            query["client_id"] = client_id
+        if scope:
+            query["scope"] = scope
+        server_key = build_server_key(spec.name, spec.url or "")
+        self._pending[spec.name] = _PendingFlow(
+            server_name=spec.name,
+            server_key=server_key,
+            endpoints=endpoints,
+            pkce=pkce,
+            state=state,
+            redirect_uri=resolved_redirect_uri,
+            client_id=client_id,
+            scope=scope,
+        )
+        return McpOAuthStartResult(
+            authorization_url=f"{endpoints.authorization_endpoint}?{urlencode(query)}",
+            state=state,
+            redirect_uri=resolved_redirect_uri,
+            server_key=server_key,
+        )
+
+    async def complete(self, server_name: str, callback_url: str) -> McpOAuthTokenRecord:
+        pending = self._pending.get(server_name)
+        if pending is None:
+            raise McpOAuthError(
+                "no_pending_flow", f"server {server_name!r} has no pending authorization flow"
+            )
+        params = {
+            key: values[-1] for key, values in parse_qs(urlsplit(callback_url).query).items()
+        }
+        error = params.get("error")
+        if error:
+            self._pending.pop(server_name, None)
+            raise McpOAuthError("provider_denied", f"authorization server returned error={error}")
+        if params.get("state") != pending.state:
+            # state 不匹配保留 pending：正确的 callback 仍可到达。
+            raise McpOAuthError(
+                "state_mismatch", f"callback state does not match pending flow for {server_name!r}"
+            )
+        code = params.get("code")
+        if not code:
+            # 无 error 也无 code 的畸形回跳：AS 未授予，按拒绝收口。
+            raise McpOAuthError("provider_denied", "callback carries neither code nor error")
+        try:
+            record = await self._exchange_code(pending, code)
+        except McpOAuthError:
+            self._pending.pop(server_name, None)
+            raise
+        self._store.set_record(pending.server_key, record, server_name=pending.server_name)
+        self._pending.pop(server_name, None)
+        return record
+
+    async def _exchange_code(self, pending: _PendingFlow, code: str) -> McpOAuthTokenRecord:
+        form: dict[str, str] = {
+            "grant_type": "authorization_code",
+            "code": code,
+            "redirect_uri": pending.redirect_uri,
+            "code_verifier": pending.pkce.code_verifier,
+        }
+        if pending.client_id:
+            form["client_id"] = pending.client_id
+        try:
+            payload = await self._exchanger.exchange(pending.endpoints.token_endpoint, form)
+        except Exception as exc:
+            raise McpOAuthError("token_exchange_failed", f"token exchange failed: {exc}") from exc
+        return _record_from_payload(payload, self._time_fn(), client_id=pending.client_id)
+
+    async def refresh(self, spec: McpServerSpec) -> McpOAuthTokenRecord:
+        """401 后的一次静默 refresh；成功写回 Store。
+
+        响应缺 ``refresh_token`` 时沿用旧值（对齐 CC 换票语义）；
+        无记录 / 无 refresh_token / 交换失败统一 ``token_exchange_failed``，
+        由 ``McpOAuthSession.refresh_authorization`` 收口清凭据。
+        """
+        server_key = build_server_key(spec.name, spec.url or "")
+        previous = self._store.get_record(server_key)
+        if previous is None or not previous.refresh_token:
+            raise McpOAuthError(
+                "token_exchange_failed", f"server {spec.name!r} has no refresh_token"
+            )
+        endpoints = await self._discovery.resolve(spec)
+        form: dict[str, str] = {
+            "grant_type": "refresh_token",
+            "refresh_token": previous.refresh_token,
+        }
+        if previous.client_id:
+            form["client_id"] = previous.client_id
+        try:
+            payload = await self._exchanger.exchange(endpoints.token_endpoint, form)
+        except Exception as exc:
+            raise McpOAuthError("token_exchange_failed", f"refresh failed: {exc}") from exc
+        record = _record_from_payload(
+            payload, self._time_fn(), client_id=previous.client_id, previous=previous
+        )
+        self._store.set_record(server_key, record, server_name=spec.name)
+        return record
+
+
+def _ensure_remote(spec: McpServerSpec) -> None:
+    if spec.transport not in ("sse", "http") or not spec.url:
+        raise McpOAuthError(
+            "not_remote", f"{spec.name}: OAuth 仅适用于带 url 的 sse/http 远程 server"
+        )
+
+
+def _record_from_payload(
+    payload: Mapping[str, Any],
+    now: float,
+    *,
+    client_id: str | None,
+    previous: McpOAuthTokenRecord | None = None,
+) -> McpOAuthTokenRecord:
+    """token 响应 → 记录；``previous`` 提供缺省（refresh 沿用旧 refresh_token/scope）。"""
+    access_token = payload.get("access_token")
+    if not isinstance(access_token, str) or not access_token:
+        raise McpOAuthError("token_exchange_failed", "token response missing access_token")
+    refresh_token = _optional_str(payload.get("refresh_token"))
+    if refresh_token is None and previous is not None:
+        refresh_token = previous.refresh_token
+    scope = _optional_str(payload.get("scope"))
+    if scope is None and previous is not None:
+        scope = previous.scope
+    token_type = _optional_str(payload.get("token_type"))
+    if token_type is None:
+        token_type = previous.token_type if previous is not None else "Bearer"
+    expires_in = payload.get("expires_in")
+    return McpOAuthTokenRecord(
+        access_token=access_token,
+        refresh_token=refresh_token,
+        expires_at=(
+            now + float(expires_in) if isinstance(expires_in, (int, float)) else None
+        ),
+        token_type=token_type,
+        scope=scope,
+        client_id=client_id,
+    )
+
+
+def _optional_str(value: object) -> str | None:
+    return value if isinstance(value, str) and value else None

langchain_agentx/mcp/auth/discovery.py

@@ -0,0 +1,134 @@
+"""OAuth endpoint 发现。
+
+职责：
+    解析 authorization / token endpoint：显式配置 → RFC 9728 → RFC 8414（path-aware）
+    → 约定路径兜底。
+
+链路位置：
+    ``McpOAuthAuthorizationFlow.start`` 调用；对齐 CC ``fetchAuthServerMetadata``
+    的裁剪版（无 configuredMetadataUrl / 无严格 schema 校验）。
+
+当前裁剪：
+    取 ``authorization_endpoint`` / ``token_endpoint`` 两字段即用；
+    网络失败不抛（落兜底）；已取得的 metadata 残缺才 ``discovery_failed``。
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, Protocol
+from urllib.parse import urlsplit, urlunsplit
+
+import httpx
+
+from langchain_agentx.mcp.auth.errors import McpOAuthError
+from langchain_agentx.protocol.mcp.types import McpServerSpec
+
+
+@dataclass(frozen=True)
+class McpOAuthEndpoints:
+    """授权码流两端点。"""
+
+    authorization_endpoint: str
+    token_endpoint: str
+
+
+class McpOAuthJsonFetcher(Protocol):
+    """GET JSON；404/无文档返回 None，网络/HTTP 错误抛异常。"""
+
+    async def get_json(self, url: str) -> Mapping[str, Any] | None: ...
+
+
+class McpOAuthHttpJsonFetcher:
+    """httpx 默认实现。"""
+
+    def __init__(self, *, timeout: float = 10.0) -> None:
+        self._timeout = timeout
+
+    async def get_json(self, url: str) -> Mapping[str, Any] | None:
+        async with httpx.AsyncClient(timeout=self._timeout, follow_redirects=True) as client:
+            response = await client.get(url, headers={"Accept": "application/json"})
+            if response.status_code == 404:
+                return None
+            response.raise_for_status()
+            payload = response.json()
+        if not isinstance(payload, Mapping):
+            raise ValueError(f"{url} returned non-object JSON")
+        return payload
+
+
+class McpOAuthDiscovery:
+    """endpoint 解析；注入 endpoints 可整体跳过网络（测试/已知配置）。"""
+
+    def __init__(
+        self,
+        fetch_json: McpOAuthJsonFetcher | None = None,
+        *,
+        endpoints: McpOAuthEndpoints | None = None,
+    ) -> None:
+        self._fetch_json = fetch_json or McpOAuthHttpJsonFetcher()
+        self._endpoints = endpoints
+
+    async def resolve(self, spec: McpServerSpec) -> McpOAuthEndpoints:
+        if self._endpoints is not None:
+            return self._endpoints
+        base = _base_url(spec.url or "")
+        try:
+            endpoints = await self._resolve_from_metadata(base)
+            if endpoints is not None:
+                return endpoints
+        except McpOAuthError:
+            raise
+        except Exception:
+            pass
+        return _conventional_endpoints(base)
+
+    async def _resolve_from_metadata(self, base: str) -> McpOAuthEndpoints | None:
+        protected = await self._fetch_json.get_json(_well_known(base, "oauth-protected-resource"))
+        if protected is not None:
+            servers = protected.get("authorization_servers")
+            if isinstance(servers, list) and servers and isinstance(servers[0], str):
+                as_base = _base_url(servers[0])
+                metadata = await self._fetch_json.get_json(
+                    _well_known(as_base, "oauth-authorization-server")
+                )
+                if metadata is not None:
+                    return _endpoints_from_metadata(metadata)
+        metadata = await self._fetch_json.get_json(_well_known(base, "oauth-authorization-server"))
+        if metadata is not None:
+            return _endpoints_from_metadata(metadata)
+        return None
+
+
+def _base_url(url: str) -> str:
+    return url.strip().rstrip("/")
+
+
+def _well_known(base: str, suffix: str) -> str:
+    parts = urlsplit(base)
+    path = parts.path.strip("/")
+    well_known = f"/.well-known/{suffix}" + (f"/{path}" if path else "")
+    return urlunsplit((parts.scheme, parts.netloc, well_known, "", ""))
+
+
+def _endpoints_from_metadata(metadata: Mapping[str, Any]) -> McpOAuthEndpoints:
+    authorization = metadata.get("authorization_endpoint")
+    token = metadata.get("token_endpoint")
+    if not isinstance(authorization, str) or not authorization:
+        raise McpOAuthError(
+            "discovery_failed", "authorization server metadata lacks authorization_endpoint"
+        )
+    if not isinstance(token, str) or not token:
+        raise McpOAuthError(
+            "discovery_failed", "authorization server metadata lacks token_endpoint"
+        )
+    return McpOAuthEndpoints(authorization_endpoint=authorization, token_endpoint=token)
+
+
+def _conventional_endpoints(base: str) -> McpOAuthEndpoints:
+    """legacy 兜底：沿 P2a stub 的约定路径。"""
+    return McpOAuthEndpoints(
+        authorization_endpoint=f"{base}/oauth/authorize",
+        token_endpoint=f"{base}/oauth/token",
+    )

langchain_agentx/mcp/auth/errors.py

@@ -0,0 +1,34 @@
+"""MCP OAuth 错误类型。
+
+职责：
+    授权流失败原因枚举，禁止调用方靠裸 Exception 字符串分叉。
+
+链路位置：
+    authorization_flow / oauth_session 抛出；宿主映射 UI。
+
+当前裁剪：
+    Phase 1 不含 refresh 专用 reason（Phase 2 可扩）。
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+McpOAuthErrorReason = Literal[
+    "not_remote",
+    "state_mismatch",
+    "provider_denied",
+    "token_exchange_failed",
+    "discovery_failed",
+    "no_pending_flow",
+    "cancelled",
+]
+
+
+class McpOAuthError(Exception):
+    """第三方 MCP OAuth 协议失败。"""
+
+    def __init__(self, reason: McpOAuthErrorReason, message: str) -> None:
+        self.reason: McpOAuthErrorReason = reason
+        self.message = message
+        super().__init__(f"{reason}: {message}")

langchain_agentx/mcp/auth/oauth_session.py

@@ -1,57 +1,107 @@
-"""OAuth 协议状态（无浏览器）。
+"""OAuth 协议门面（无浏览器）。
 
 职责：
-    授权 URL、callback/token 登记。宿主打开浏览器；SDK 只持 token。
+    ``needs_auth`` 后宿主经 ``start_authorization`` / ``complete_authorization`` /
+    ``logout`` 驱动授权码 + PKCE；兼容 P2a ``submit_token`` / ``authorized_headers``。
 
 链路位置：
-    对标 CC ``performMCPOAuthFlow`` 的协议切片（忽略 xaa / WebView）。
-    Manager connect 无 token 时 ``needs_auth``。
+    Manager connect 无 token 时 ``needs_auth``（经 ``get_token`` /
+    ``authorized_headers`` 消费）；打开浏览器 / 收 callback 在宿主。
 
 当前裁剪：
-    内存 token；无 PKCE/IdP 实现。
+    Memory store 默认（生产接 ``McpFileTokenStore``）；refresh 单次、由 Manager
+    在 401 时触发。
 """
 
 from __future__ import annotations
 
+from langchain_agentx.mcp.auth.authorization_flow import (
+    McpOAuthAuthorizationFlow,
+    McpOAuthStartResult,
+)
+from langchain_agentx.mcp.auth.errors import McpOAuthError
+from langchain_agentx.mcp.auth.token_store import (
+    McpOAuthTokenRecord,
+    McpTokenStore,
+    McpTokenStoreProtocol,
+    build_server_key,
+)
 from langchain_agentx.protocol.mcp.types import McpServerSpec
 
-
-class McpTokenStore:
-    """进程内 token。测试注入 mock。"""
-
-    def __init__(self) -> None:
-        self._tokens: dict[str, str] = {}
-
-    def get_token(self, server_name: str) -> str | None:
-        return self._tokens.get(server_name)
-
-    def set_token(self, server_name: str, access_token: str) -> None:
-        self._tokens[server_name] = access_token
+__all__ = [
+    "McpOAuthSession",
+    "McpOAuthStartResult",
+    "McpOAuthTokenRecord",
+    "McpTokenStore",
+]
 
 
 class McpOAuthSession:
-    """needs_auth → 宿主填 token → 再 connect。"""
-
-    def __init__(self, store: McpTokenStore | None = None) -> None:
-        self._store = store or McpTokenStore()
+    """needs_auth → 宿主 PKCE（或 submit_token）→ 再 connect。"""
+
+    def __init__(
+        self,
+        store: McpTokenStoreProtocol | None = None,
+        *,
+        flow: McpOAuthAuthorizationFlow | None = None,
+    ) -> None:
+        if flow is not None and store is None:
+            store = flow.store
+        self._store: McpTokenStoreProtocol = store or McpTokenStore()
+        self._flow = flow or McpOAuthAuthorizationFlow(self._store)
 
     @property
-    def store(self) -> McpTokenStore:
+    def store(self) -> McpTokenStoreProtocol:
         return self._store
 
+    @property
+    def flow(self) -> McpOAuthAuthorizationFlow:
+        return self._flow
+
     def get_token(self, server_name: str) -> str | None:
         return self._store.get_token(server_name)
 
-    def build_authorization_url(self, spec: McpServerSpec) -> str:
-        base = (spec.url or "").rstrip("/")
-        return f"{base}/oauth/authorize"
-
     def submit_token(self, server_name: str, access_token: str) -> None:
         self._store.set_token(server_name, access_token)
 
     def authorized_headers(self, spec: McpServerSpec) -> dict[str, str]:
         headers = dict(spec.headers or {})
-        token = self.get_token(spec.name)
+        token = self._token_for_spec(spec)
         if token:
             headers["Authorization"] = f"Bearer {token}"
         return headers
+
+    def _token_for_spec(self, spec: McpServerSpec) -> str | None:
+        """同 name 多 url 时按 server_key 取（不串票）；无记录回退 by-name（P2a）。"""
+        if spec.url:
+            record = self._store.get_record(build_server_key(spec.name, spec.url))
+            if record is not None:
+                return record.access_token
+        return self.get_token(spec.name)
+
+    async def refresh_authorization(self, spec: McpServerSpec) -> bool:
+        """401 后的一次静默 refresh；任何失败清该 server 凭据并返回 False。"""
+        try:
+            await self._flow.refresh(spec)
+            return True
+        except McpOAuthError:
+            self._store.clear(spec.name)
+            return False
+
+    async def start_authorization(
+        self,
+        spec: McpServerSpec,
+        *,
+        redirect_uri: str | None = None,
+        client_id: str | None = None,
+        scope: str | None = None,
+    ) -> McpOAuthStartResult:
+        return await self._flow.start(
+            spec, redirect_uri=redirect_uri, client_id=client_id, scope=scope
+        )
+
+    async def complete_authorization(self, server_name: str, callback_url: str) -> None:
+        await self._flow.complete(server_name, callback_url)
+
+    def logout(self, server_name: str) -> None:
+        self._store.clear(server_name)

langchain_agentx/mcp/auth/pkce.py

@@ -0,0 +1,38 @@
+"""PKCE (S256) 值对象。
+
+职责：
+    生成 code_verifier / code_challenge，供授权码流使用。
+
+链路位置：
+    McpOAuthAuthorizationFlow.start 调用。
+
+当前裁剪：
+    仅 S256；无 plain。
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import secrets
+from dataclasses import dataclass
+
+
+def _b64url_no_pad(raw: bytes) -> str:
+    return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
+
+
+@dataclass(frozen=True)
+class McpPkcePair:
+    """RFC 7636 S256 PKCE 对。"""
+
+    code_verifier: str
+    code_challenge: str
+    code_challenge_method: str = "S256"
+
+    @classmethod
+    def generate(cls) -> McpPkcePair:
+        verifier = _b64url_no_pad(secrets.token_bytes(32))
+        digest = hashlib.sha256(verifier.encode("ascii")).digest()
+        challenge = _b64url_no_pad(digest)
+        return cls(code_verifier=verifier, code_challenge=challenge)