"""OAuth login flow and token storage for MCP servers. Note: `mcp.shared.auth.OAuthToken` is a pydantic model whose default `repr` includes the access and refresh token strings verbatim. Never log one via `%r`, `str()`, f-string interpolation, or `logger.exception`/`exc_info` on an exception that wraps one — the tokens will land in stdout, log files, and error-reporting pipelines. Pass only structural facts ("refreshed token for server X") rather than the token itself. """ from __future__ import annotations import concurrent.futures import contextlib import hashlib import html import json import logging import os import re import secrets import stat import threading import time from collections.abc import AsyncGenerator, Awaitable, Callable from typing import TYPE_CHECKING, Literal, TypedDict from urllib.parse import parse_qs, urlparse import httpx from mcp.client.auth import OAuthClientProvider, TokenStorage from mcp.client.auth.utils import ( build_oauth_authorization_server_metadata_discovery_urls, build_protected_resource_metadata_discovery_urls, create_oauth_metadata_request, handle_auth_metadata_response, handle_protected_resource_response, ) from mcp.client.streamable_http import MCP_PROTOCOL_VERSION from mcp.shared.auth import ( OAuthClientInformationFull, OAuthMetadata, OAuthToken, ) from pydantic import BaseModel, ConfigDict, ValidationError if TYPE_CHECKING: from pathlib import Path from deepagents_code.mcp_oauth_ui import OAuthInteraction class _DeviceCodeResponse(BaseModel): """RFC 8628 §3.2 device-authorization response payload.""" model_config = ConfigDict(extra="ignore") device_code: str """Opaque device code the client polls with at the token endpoint.""" user_code: str """Short code the user enters in the browser to approve the device.""" verification_uri: str """Provider URL the user visits to complete device authorization.""" expires_in: int """Lifetime of the device code in seconds.""" interval: int = 5 """Recommended polling interval in seconds when the provider omits one.""" class McpServerSpec(TypedDict, total=False): """Parsed MCP server config entry. All keys are optional at the type level because `mcpServers` entries are validated shape-first by `_validate_server_config` rather than by the type system. This TypedDict documents the accepted shape for readers and static checkers — validate the fields at use sites before relying on them. """ auth: Literal["oauth"] """Authentication mode for remote MCP servers that require OAuth login.""" type: Literal["stdio", "http", "sse"] """Transport type when the config uses the `type` key.""" transport: Literal["stdio", "http", "sse"] """Transport type when the config uses the `transport` key.""" url: str """Remote endpoint URL for HTTP or SSE MCP servers.""" headers: dict[str, str] """Optional request headers sent when connecting to the remote server.""" command: str """Executable for stdio MCP servers.""" args: list[str] """Command-line arguments passed to the stdio server executable.""" env: dict[str, str] """Environment overrides for launching a stdio MCP server.""" logger = logging.getLogger(__name__) _REF_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}") """Matches `${VAR}` placeholders inside config strings for env-var substitution.""" _SAFE_SERVER_NAME_RE = re.compile(r"^[A-Za-z0-9_-]+$") """Matches server names that are safe to embed in token-file basenames. Mirrors `_SERVER_NAME_RE` in `mcp_tools` — duplicated here because `mcp_auth` cannot import from `mcp_tools` at module top-level without risking a circular import. Keep both regexes in sync.""" _STORAGE_VERSION = 1 """Schema version stamped into persisted credential files; bump on incompatible shape changes so `_load_*` can reject or migrate older payloads.""" _REFRESH_SAFETY_MARGIN_SECONDS = 30.0 """Refresh access tokens this many seconds before their advertised expiry. Absorbs clock skew and request latency so a token deemed valid locally isn't rejected as expired by the server — without the margin, a 401 sends the SDK into the full re-auth (browser) flow instead of the cheaper refresh grant. """ def resolve_headers( headers: dict[str, str], *, server_name: str | None = None, ) -> dict[str, str]: """Resolve `${VAR}` env-var references in header values. Args: headers: Raw header mapping from MCP config. server_name: Optional server name for error messages. Returns: A new dict with env-var references resolved to current values. Raises: TypeError: If a header value is not a string. RuntimeError: If a `${VAR}` reference points to an unset env var. """ # noqa: DOC502 - RuntimeError is raised via `_interpolate` resolved: dict[str, str] = {} for name, value in headers.items(): if not isinstance(value, str): where = f"mcpServers.{server_name}.headers.{name}" if server_name else name msg = f"{where} must be a string, got {type(value).__name__}" raise TypeError(msg) resolved[name] = _interpolate(value, header=name, server_name=server_name) return resolved def _interpolate(s: str, *, header: str, server_name: str | None) -> str: """Expand `${VAR}` references in `s` against the current environment. Args: s: Raw header value. header: Header name, used in error messages. server_name: Owning server name for error messages. Returns: Interpolated string. Raises: RuntimeError: If a referenced env var is unset. """ # noqa: DOC502 - raised inside the inner `replace` substitution callback def replace(match: re.Match[str]) -> str: var_name = match.group(1) val = os.environ.get(var_name) if val is None: where = ( f"mcpServers.{server_name}.headers.{header}" if server_name else header ) msg = ( f"{where} references unset env var {var_name}. " f"Set {var_name} in the environment or remove the reference." ) raise RuntimeError(msg) return val return _REF_RE.sub(replace, s) def _tokens_dir() -> Path: """Return `~/.deepagents/.state/mcp-tokens/`. The deferred import lets tests redirect token storage into a temp directory by patching `deepagents_code.model_config.DEFAULT_STATE_DIR`. """ from deepagents_code.model_config import DEFAULT_STATE_DIR return DEFAULT_STATE_DIR / "mcp-tokens" def _token_file_stem(server_name: str, server_url: str | None) -> str: """Return a path-safe storage stem for this server identity. Safety of the stem depends on `server_name` already having passed `_SERVER_NAME_RE` in `_validate_server_config` — the URL is hashed to a hex digest, so only the server name can carry path separators. """ if server_url is None: return server_name digest = hashlib.sha256(server_url.encode("utf-8")).hexdigest()[:16] return f"{server_name}-{digest}" class FileTokenStorage(TokenStorage): """File-backed `TokenStorage` under `~/.deepagents/.state/mcp-tokens/`.""" def __init__(self, server_name: str, *, server_url: str | None = None) -> None: """Bind this storage to a configured MCP server identity. Raises: ValueError: If `server_name` contains characters that would let it escape the `~/.deepagents/.state/mcp-tokens/` directory when used as the token-file basename. """ if not _SAFE_SERVER_NAME_RE.fullmatch(server_name): msg = ( f"Invalid MCP server name {server_name!r}: token storage " "names must match [A-Za-z0-9_-]+ to keep the on-disk path " "inside ~/.deepagents/.state/mcp-tokens/." ) raise ValueError(msg) self._server_name = server_name self._server_url = server_url @property def path(self) -> Path: """Return the on-disk token file path for this server.""" stem = _token_file_stem(self._server_name, self._server_url) return _tokens_dir() / f"{stem}.json" async def get_tokens(self) -> OAuthToken | None: """Return the stored `OAuthToken`, or `None` if none is persisted.""" data = self._read() if data is None: return None raw = data.get("tokens") if raw is None: return None return OAuthToken.model_validate(raw) async def set_tokens(self, tokens: OAuthToken) -> None: """Persist `tokens` to disk, preserving any stored client info. Also records the absolute Unix-epoch expiry as a sidecar so a cold-started provider can detect a stale access token and trigger the SDK's `refresh_token` grant instead of a full browser re-auth. Cleared when `expires_in` is absent so the sidecar can't go stale. """ data = self._read() or {} data["version"] = _STORAGE_VERSION data["tokens"] = json.loads(tokens.model_dump_json(exclude_none=True)) if tokens.expires_in is not None: data["expires_at"] = time.time() + tokens.expires_in else: data.pop("expires_at", None) self._write(data) async def get_client_info(self) -> OAuthClientInformationFull | None: """Return the stored client registration, or `None` if none is persisted.""" data = self._read() if data is None: return None raw = data.get("client_info") if raw is None: return None return OAuthClientInformationFull.model_validate(raw) async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: """Persist `client_info` to disk, preserving any stored tokens.""" data = self._read() or {} data["version"] = _STORAGE_VERSION data["client_info"] = json.loads(client_info.model_dump_json(exclude_none=True)) self._write(data) async def get_oauth_metadata(self) -> OAuthMetadata | None: """Return stored public OAuth authorization metadata, if available.""" data = self._read() if data is None: return None raw = data.get("oauth_metadata") if raw is None: return None return OAuthMetadata.model_validate(raw) async def set_oauth_metadata(self, metadata: OAuthMetadata) -> None: """Persist public OAuth authorization metadata beside the token state.""" data = self._read() or {} data["version"] = _STORAGE_VERSION data["oauth_metadata"] = json.loads(metadata.model_dump_json(exclude_none=True)) self._write(data) async def set_tokens_and_client_info( self, tokens: OAuthToken, client_info: OAuthClientInformationFull, ) -> None: """Persist tokens and client info in a single atomic write. Prevents the state where one call succeeds and the other fails, leaving an orphan on disk. """ data = self._read() or {} data["version"] = _STORAGE_VERSION data["tokens"] = json.loads(tokens.model_dump_json(exclude_none=True)) data["client_info"] = json.loads(client_info.model_dump_json(exclude_none=True)) if tokens.expires_in is not None: data["expires_at"] = time.time() + tokens.expires_in else: data.pop("expires_at", None) self._write(data) async def get_expires_at(self) -> float | None: """Return the stored absolute token expiry (Unix epoch), or `None`. Returns `None` for token files written before this field existed, for tokens whose provider omitted `expires_in`, or when the sidecar value fails to coerce to `float`. Callers should treat `None` as "expiry unknown" and decide policy (skip, assume-expired, etc.). """ data = self._read() if data is None: return None raw = data.get("expires_at") if raw is None: return None try: return float(raw) except (TypeError, ValueError): # Log the value's type, never the value itself — the sidecar lives # next to the bearer token in the same JSON envelope, so a # misplaced token string would land here. logger.warning( "MCP token sidecar 'expires_at' for %s is not numeric (%s); " "treating as unknown. The next request will trigger a refresh " "or browser re-auth.", self._server_name, type(raw).__name__, ) return None def stored_loopback_port(self) -> int | None: """Return the stored loopback redirect URI port, if one is reusable. DCR registers `client_id` against a specific `redirect_uri`. If the callback server binds a fresh random port on a later launch, the authorize request will carry a `redirect_uri` that no longer matches the one registered with the persisted `client_id`, and the authorization server will reject it ("invalid or missing redirect_uri"). Reusing the persisted port keeps the registration valid across runs. Returns: The integer port parsed from a stored `http://localhost:/callback` redirect URI, or `None` if no usable port is on disk. """ try: data = self._read() except RuntimeError as exc: logger.warning( "MCP token file for %s is unreadable during loopback port " "lookup; falling back to a fresh random port. Delete the file " "and log in again if OAuth authorization fails: %s", self.path, exc, ) return None if data is None: return None client_info = data.get("client_info") or {} redirect_uris = client_info.get("redirect_uris") or [] if not redirect_uris: return None uri = str(redirect_uris[0]) port = self._loopback_callback_port(uri) if port is None: logger.warning( "Stored MCP OAuth redirect URI for %s is not a reusable " "loopback callback URI; falling back to a fresh random port. " "OAuth authorization may fail if the server requires the " "persisted client registration redirect URI: %s", self.path, uri, ) return None return port @staticmethod def _loopback_callback_port(uri: str) -> int | None: """Return `uri`'s port if it is a reusable loopback callback URI. A reusable URI is `http://localhost:/callback` with an explicit port. Anything else (a different scheme/host/path, or a portless `http://localhost/callback`) returns `None` because it cannot be paired with the loopback callback server a CLI login binds. """ parsed = urlparse(uri) try: port = parsed.port except ValueError: return None if ( parsed.scheme != "http" or parsed.hostname != _LOOPBACK_URI_HOST or parsed.path != _LOOPBACK_CALLBACK_PATH or port is None ): return None return port def discard_client_info_if_loopback_unusable(self) -> bool: """Drop a persisted client registration that can't serve loopback login. `stored_loopback_port` explains why the authorize request's `redirect_uri` must match the one registered against the persisted `client_id`. When that registered URI is not a reusable loopback callback (e.g. a portless `http://localhost/callback` left by an earlier non-loopback login), no port can be reused, so a CLI login binds a fresh random port and the server rejects the mismatched `redirect_uri` ("invalid or missing redirect_uri"). Removing the stale `client_info` makes the SDK perform a fresh DCR with the loopback redirect URI it will actually use. Only discards when no token is persisted at all, so a session with a usable (or refreshable) token is never downgraded to a full re-auth. The presence check is deliberately conservative: it errs toward keeping the registration, never toward deleting one that might still be needed. Returns: `True` if a stale client registration was removed. `False` covers both "nothing to discard" and "could not discard" (an unreadable or unwritable token file, logged where it happens). """ try: data = self._read() except RuntimeError as exc: # Mirror `stored_loopback_port`: a corrupt or unsupported-version # token file carries actionable "delete the file" guidance in the # `RuntimeError` message, so surface it rather than dropping it. logger.warning( "MCP token file for %s is unreadable while checking for a " "stale client registration; skipping self-heal. Delete the " "file and log in again if OAuth authorization fails: %s", self.path, exc, ) return False if data is None or "client_info" not in data: return False # A persisted access/refresh token can still authenticate (or refresh) # without re-running the authorization-code grant, so leave the # registration intact rather than forcing an avoidable re-auth. if data.get("tokens") is not None: return False redirect_uris = (data.get("client_info") or {}).get("redirect_uris") or [] if ( redirect_uris and self._loopback_callback_port(str(redirect_uris[0])) is not None ): return False del data["client_info"] try: self._write(data) except OSError as exc: # Surface the failure but don't crash login: the stale registration # simply remains, and the login attempt fails the same way it would # have without this self-heal. logger.warning( "Could not remove stale MCP client registration in %s: %s", self.path, exc, ) return False return True def _read(self) -> dict | None: path = self.path if not path.exists(): return None try: raw = path.read_text(encoding="utf-8") data = json.loads(raw) except (OSError, json.JSONDecodeError) as exc: msg = ( f"Failed to read MCP token file {path}: {exc}. " f"Delete the file and run `/mcp login {self._server_name}` " f"in the TUI (or `dcode mcp login {self._server_name}`)." ) raise RuntimeError(msg) from exc if data.get("version") != _STORAGE_VERSION: msg = ( f"MCP token file {path} has unsupported version " f"{data.get('version')!r} (expected {_STORAGE_VERSION}). " f"Delete it and run `/mcp login {self._server_name}` in the " f"TUI (or `dcode mcp login {self._server_name}`)." ) raise RuntimeError(msg) return data def _write(self, data: dict) -> None: path = self.path path.parent.mkdir(parents=True, exist_ok=True) if hasattr(os, "chmod"): try: path.parent.chmod(stat.S_IRWXU) except OSError as exc: # A failing chmod on the parent dir leaves the tokens # directory at the default umask. Warn so operators on # shared hosts notice. logger.warning( "Could not lock down MCP tokens dir %s (mode 0700): %s. " "Tokens may be readable by other local users.", path.parent, exc, ) tmp = path.with_suffix(path.suffix + ".tmp") payload = json.dumps(data, separators=(",", ":")).encode("utf-8") # O_EXCL + mode 0600 means the token file is never visible at the # default umask between open() and chmod(). On Windows, os.open() # ignores the mode bits, so the explicit chmod below is the # cross-platform guarantee. flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL with contextlib.suppress(FileNotFoundError): tmp.unlink() fd = os.open(str(tmp), flags, 0o600) try: with os.fdopen(fd, "wb") as fh: fh.write(payload) except Exception: with contextlib.suppress(OSError): tmp.unlink() raise try: tmp.replace(path) except Exception: with contextlib.suppress(OSError): tmp.unlink() raise if hasattr(os, "chmod"): # Already 0600 from os.open on POSIX; a second chmod covers # filesystems that ignore the create-mode argument. try: path.chmod(stat.S_IRUSR | stat.S_IWUSR) except OSError as exc: logger.warning( "Could not set mode 0600 on MCP token file %s: %s. " "Stored refresh/access tokens may be world-readable.", path, exc, ) RedirectHandler = Callable[[str], Awaitable[None]] CallbackHandler = Callable[[], Awaitable[tuple[str, str | None]]] _LOOPBACK_BIND_HOST = "127.0.0.1" _LOOPBACK_URI_HOST = "localhost" _LOOPBACK_CALLBACK_PATH = "/callback" _LOOPBACK_CALLBACK_TIMEOUT = 300.0 class _LoopbackCallbackTimeoutError(RuntimeError): """Raised when the browser never reaches the local callback server.""" class _LoopbackCallbackUnavailableError(RuntimeError): """Raised when the local callback server cannot be started.""" def _choose_loopback_port() -> int: """Return a high local TCP port candidate without opening a socket. The OAuth redirect URI must be known before the provider starts the handshake, but the actual callback server should not keep a socket open unless a browser redirect is needed. Returns: A port number from the dynamic/private port range. """ return 49152 + secrets.randbelow(65535 - 49152 + 1) class _LoopbackOAuthCallbackServer: """Single-use loopback HTTP server for CLI OAuth callbacks. Port selection and socket binding are intentionally separated: the redirect URI is fixed at construction time so it can be registered with the OAuth provider, while the socket is not opened until `start()` is called from the redirect handler. """ def __init__(self, *, port: int) -> None: """Prepare a callback server for a previously selected loopback port. Args: port: TCP port to bind when `start()` is called. """ self._port = port self.redirect_uri = ( f"http://{_LOOPBACK_URI_HOST}:{port}{_LOOPBACK_CALLBACK_PATH}" ) self._future: concurrent.futures.Future[tuple[str, str | None]] = ( concurrent.futures.Future() ) self._server: object | None = None self._started = False self._closed = False self._thread: threading.Thread | None = None @property def port(self) -> int: """Loopback TCP port this server will bind on `start()`.""" return self._port def start(self) -> None: """Bind and start serving callback requests in a background thread.""" if self._started: return from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer parent = self class Handler(BaseHTTPRequestHandler): def do_GET(self) -> None: parent._handle_get(self) def log_message( # noqa: PLR6301 # stdlib override self, format: str, # noqa: A002 *args: object, ) -> None: del format, args self._server = ThreadingHTTPServer((_LOOPBACK_BIND_HOST, self._port), Handler) self._started = True self._thread = threading.Thread( target=self._serve_forever, name="deepagents-mcp-oauth-callback", daemon=True, ) self._thread.start() def _serve_forever(self) -> None: from http.server import ThreadingHTTPServer if isinstance(self._server, ThreadingHTTPServer): self._server.serve_forever() async def wait(self) -> tuple[str, str | None]: """Wait for the authorization callback and return `(code, state)`. Returns: The OAuth authorization code and optional state. Raises: _LoopbackCallbackTimeoutError: If no callback arrives before the timeout. _LoopbackCallbackUnavailableError: If the server could not be started. RuntimeError: If the provider returned an OAuth error or the callback URL lacked a `code` parameter. """ # noqa: DOC502 - _LoopbackCallbackUnavailableError/RuntimeError set on future import asyncio try: return await asyncio.wait_for( asyncio.wrap_future(self._future), timeout=_LOOPBACK_CALLBACK_TIMEOUT, ) except TimeoutError as exc: msg = "Browser callback was not received before the timeout." raise _LoopbackCallbackTimeoutError(msg) from exc def fail(self, exc: Exception) -> None: """Poison the future so `wait()` raises `exc` immediately. Args: exc: Exception to surface to the awaiting coroutine. """ if not self._future.done(): self._future.set_exception(exc) def close(self) -> None: """Stop the local callback server and release its socket.""" if self._closed: return self._closed = True from http.server import ThreadingHTTPServer if isinstance(self._server, ThreadingHTTPServer): if self._started: self._server.shutdown() self._server.server_close() def _handle_get(self, request: object) -> None: from http.server import BaseHTTPRequestHandler handler = request if not isinstance(handler, BaseHTTPRequestHandler): return if self._future.done(): # Duplicate browser request (retry, prefetch, favicon) after the # flow already completed. Respond and return without touching the # future — avoids InvalidStateError from a concurrent set_result. # Branch on the future's terminal state: a previous error must # not be papered over with a success page. if self._future.exception() is None: self._send_html( handler, 200, _oauth_success_html( "MCP authorization complete. " "This tab will close automatically.", ), ) else: self._send_html( handler, 400, _oauth_error_html( "Authorization did not complete. " "Return to your terminal for details.", ), ) return parsed = urlparse(handler.path) if parsed.path != _LOOPBACK_CALLBACK_PATH: self._send_html( handler, 404, _oauth_error_html("Callback route not found."), ) return params = parse_qs(parsed.query) if "error" in params: err_code = params["error"][0] err_desc = (params.get("error_description") or [""])[0] detail = f": {err_desc}" if err_desc else "" msg = f"Authorization denied by provider: {err_code}{detail}" self._future.set_exception(RuntimeError(msg)) self._send_html(handler, 400, _oauth_error_html(msg)) return if "code" not in params or not params["code"]: msg = "Callback URL is missing the 'code' parameter." self._future.set_exception(RuntimeError(msg)) self._send_html(handler, 400, _oauth_error_html(msg)) return self._future.set_result((params["code"][0], (params.get("state") or [None])[0])) self._send_html( handler, 200, _oauth_success_html( "MCP authorization complete. This tab will close automatically.", ), ) @staticmethod def _send_html(handler: object, status: int, body: str) -> None: from http.server import BaseHTTPRequestHandler if not isinstance(handler, BaseHTTPRequestHandler): return payload = body.encode("utf-8") handler.send_response(status) handler.send_header("Content-Type", "text/html; charset=utf-8") handler.send_header("Content-Length", str(len(payload))) handler.end_headers() handler.wfile.write(payload) def _oauth_success_html(message: str) -> str: return _oauth_result_html( title="Authorization complete", heading="You're signed in", message=message, status="success", ) def _oauth_error_html(message: str) -> str: return _oauth_result_html( title="Authorization failed", heading="Authorization failed", message=message, status="error", ) def _oauth_result_html( *, title: str, heading: str, message: str, status: Literal["success", "error"], ) -> str: accent = "#137333" if status == "success" else "#b3261e" background = "#eef7f0" if status == "success" else "#fceeee" mark = "✓" if status == "success" else "!" escaped_title = html.escape(title) escaped_heading = html.escape(heading) escaped = html.escape(message) # `window.close()` is only honored for tabs the script itself opened # (browser policy), but for the common case where the auth flow was # launched via `window.open` / `target=_blank` from another page, the # tab closes cleanly. When the browser refuses, the user still sees the # static success page and the message text remains accurate. auto_close = ( "" if status == "success" else "" ) return ( '' '' f"{escaped_title}" "" '
' f'
{mark}
' f"

{escaped_heading}

{escaped}

" "
" f"{auto_close}" "" ) class MCPReauthRequiredError(RuntimeError): """Raised when an MCP server needs interactive re-authentication.""" def __init__(self, server_name: str) -> None: """Build with `server_name` so the message tells the user what to fix.""" self.server_name = server_name super().__init__( f"MCP server {server_name!r} needs re-authentication. " f"Run `/mcp login {server_name}` in the TUI, or " f"`dcode mcp login {server_name}` from the shell.", ) def _make_reauth_required_handlers( server_name: str, ) -> tuple[RedirectHandler, CallbackHandler]: """Return OAuth handlers that refuse to prompt and raise instead. Used in non-interactive server mode so that a missing or expired token surfaces as `MCPReauthRequiredError` rather than hanging on `input()`. """ async def redirect(_auth_url: str) -> None: # noqa: RUF029 raise MCPReauthRequiredError(server_name) async def callback() -> tuple[str, str | None]: # noqa: RUF029 raise MCPReauthRequiredError(server_name) return redirect, callback def _make_paste_back_handlers( *, extra_auth_params: dict[str, str] | None = None, ui: OAuthInteraction | None = None, ) -> tuple[RedirectHandler, CallbackHandler]: """Create paste-back redirect and callback handlers for OAuth. Args: extra_auth_params: Extra query params to append to the auth URL. ui: Interaction surface for the auth URL display and the pasted-back callback URL prompt. Returns: A tuple of `(redirect_handler, callback_handler)`. """ extras = dict(extra_auth_params or {}) interaction = ui if ui is not None else _default_ui() async def redirect(auth_url: str) -> None: final_url = _append_query_params(auth_url, extras) if extras else auth_url await interaction.show_authorize_url(final_url, opened_in_browser=False) async def callback() -> tuple[str, str | None]: url = await interaction.request_callback_url() return _parse_callback_url(url) return redirect, callback def _parse_callback_url(url: str) -> tuple[str, str | None]: """Parse a provider callback URL into `(code, state)`. Args: url: Raw callback URL pasted by the user. Returns: The `code` and optional `state` query parameters. Raises: RuntimeError: If the URL contains `error=` or lacks `code`. """ params = parse_qs(urlparse(url).query) if "error" in params: err_code = params["error"][0] err_desc = (params.get("error_description") or [""])[0] detail = f": {err_desc}" if err_desc else "" msg = f"Authorization denied by provider: {err_code}{detail}" raise RuntimeError(msg) if "code" not in params or not params["code"]: msg = "Callback URL is missing the 'code' parameter." raise RuntimeError(msg) return params["code"][0], (params.get("state") or [None])[0] def _default_ui() -> OAuthInteraction: """Return the default `OAuthInteraction` implementation (CLI stdio).""" from deepagents_code.mcp_oauth_ui import CliOAuthInteraction return CliOAuthInteraction() def _make_loopback_handlers( *, callback_server: _LoopbackOAuthCallbackServer, extra_auth_params: dict[str, str] | None = None, ui: OAuthInteraction | None = None, ) -> tuple[RedirectHandler, CallbackHandler]: """Create browser loopback redirect and callback handlers for OAuth. Args: callback_server: Prepared local callback server for this login attempt. The socket is bound when the returned redirect handler is first called. extra_auth_params: Extra query params to append to the auth URL. ui: Interaction surface for the browser-opened or fallback prompts. Returns: A tuple of `(redirect_handler, callback_handler)`. """ extras = dict(extra_auth_params or {}) interaction = ui if ui is not None else _default_ui() last_authorize_url: str | None = None _paste_redirect, paste_callback = _make_paste_back_handlers( extra_auth_params=extra_auth_params, ui=interaction, ) async def redirect(auth_url: str) -> None: import asyncio import webbrowser nonlocal last_authorize_url final_url = _append_query_params(auth_url, extras) if extras else auth_url last_authorize_url = final_url # Resolve a browser explicitly before opening so headless / SSH # environments fall through to paste-back without burning the # 300s loopback timeout. `webbrowser.open` can return `True` in # those environments even when nothing launches. try: await asyncio.to_thread(webbrowser.get) has_browser = True except webbrowser.Error: has_browser = False if has_browser: opened = await asyncio.to_thread(webbrowser.open, final_url) else: opened = False if not opened: callback_server.fail( _LoopbackCallbackUnavailableError( "No browser is available to complete the OAuth flow.", ), ) await interaction.show_authorize_url(final_url, opened_in_browser=False) return try: callback_server.start() except OSError as exc: logger.warning( "Could not start loopback OAuth callback server on port %s: %s", callback_server.port, exc, ) msg = "Local OAuth callback server could not be started." callback_server.fail(_LoopbackCallbackUnavailableError(msg)) await interaction.show_notice( "Could not start the local OAuth callback server.", ) await interaction.show_authorize_url(final_url, opened_in_browser=False) return await interaction.show_authorize_url(final_url, opened_in_browser=True) async def callback() -> tuple[str, str | None]: try: return await callback_server.wait() except ( _LoopbackCallbackTimeoutError, _LoopbackCallbackUnavailableError, ) as exc: if last_authorize_url is not None: await interaction.show_authorize_url( last_authorize_url, opened_in_browser=False, ) await interaction.show_notice( f"{exc}\nPaste the full callback URL instead.", ) return await paste_callback() finally: callback_server.close() return redirect, callback def _append_query_params(url: str, params: dict[str, str]) -> str: """Return `url` with `params` replacing any same-named query keys.""" from urllib.parse import urlencode, urlunparse parsed = urlparse(url) existing = dict(parse_qs(parsed.query, keep_blank_values=True)) for key, value in params.items(): existing[key] = [value] return urlunparse(parsed._replace(query=urlencode(existing, doseq=True))) class _ExpiryAwareOAuthClientProvider(OAuthClientProvider): """`OAuthClientProvider` that restores `token_expiry_time` from storage. Upstream `_initialize` loads stored tokens but leaves `context.token_expiry_time` at `None`, which makes `is_token_valid` report any stored access token — even one that expired hours ago — as valid. The SDK then sends a stale `Bearer`, gets a 401, and falls into a full re-auth (browser) instead of the `refresh_token` grant. Restoring the persisted absolute expiry to the context after load lets the SDK's refresh-when-invalid-and-refreshable branch fire on the first request after a cold start. When the sidecar is absent (older token files written before this field existed), assume the token is expired so the refresh path still gets a chance before falling back to 401. """ async def _initialize(self) -> None: # Overrides a leading-underscore SDK method; behavior depends on # `super()._initialize()` populating `context.current_tokens` from # storage. If an upstream rename or refactor breaks that contract, # the test suite's TestExpiryAwareOAuthClientProvider cases will # fail loudly rather than silently regress to the 401-on-restart # bug this class exists to prevent. await super()._initialize() if self.context.oauth_metadata is None: get_oauth_metadata = getattr( self.context.storage, "get_oauth_metadata", None, ) if get_oauth_metadata is not None: self.context.oauth_metadata = await get_oauth_metadata() get_expires_at = getattr(self.context.storage, "get_expires_at", None) if get_expires_at is None: return expires_at = await get_expires_at() tokens = self.context.current_tokens if expires_at is None: # Use 1.0 (one second after the Unix epoch) rather than 0.0 so the # SDK's `not self.token_expiry_time` falsy-zero check doesn't treat # the sentinel as "no expiry known" and mark the token valid again. if tokens is not None and tokens.refresh_token: self.context.token_expiry_time = 1.0 elif tokens is not None and tokens.access_token: # Legacy file with no refresh_token: nothing we can do to # pre-empt expiry. Surface a structural breadcrumb so the # 401-then-browser-reauth flow isn't completely silent. logger.info( "Legacy MCP token file for %s has no refresh_token; " "cannot pre-empt expiry. The next 401 will trigger " "browser re-auth.", self.context.server_url, ) return if expires_at - time.time() < _REFRESH_SAFETY_MARGIN_SECONDS: # Token already inside its safety margin (or past it) — likely a # cold start after a long pause, or a misconfigured server issuing # sub-margin lifetimes. Log only the duration, never any token # material. logger.debug( "MCP token for %s is within %.0fs of expiry on load; " "scheduling refresh on next request.", self.context.server_url, _REFRESH_SAFETY_MARGIN_SECONDS, ) self.context.token_expiry_time = expires_at - _REFRESH_SAFETY_MARGIN_SECONDS async def _persist_oauth_metadata(self) -> None: """Persist discovered public OAuth metadata when storage supports it.""" if self.context.oauth_metadata is None: return set_oauth_metadata = getattr(self.context.storage, "set_oauth_metadata", None) if set_oauth_metadata is not None: await set_oauth_metadata(self.context.oauth_metadata) async def _handle_token_response(self, response: httpx.Response) -> None: """Persist tokens and any metadata discovered during full OAuth login.""" await super()._handle_token_response(response) await self._persist_oauth_metadata() async def async_auth_flow( self, request: httpx.Request, ) -> AsyncGenerator[httpx.Request, httpx.Response]: """Discover and cache OAuth metadata before the SDK refresh branch. Yields: HTTP requests for OAuth metadata discovery and the delegated SDK auth flow. """ async with self.context.lock: if not self._initialized: await self._initialize() self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION) if ( not self.context.is_token_valid() and self.context.can_refresh_token() and self.context.oauth_metadata is None ): # Pre-empt the SDK's 401-path discovery so its refresh branch # finds populated `oauth_metadata` and uses the advertised token # endpoint instead of guessing `/token`. The resource-metadata # URL is `None`: no 401 yet, so no `WWW-Authenticate` to read. try: prm_urls = build_protected_resource_metadata_discovery_urls( None, self.context.server_url, ) for url in prm_urls: # ASYNC119: yielding the request to receive its response is # this auth generator's handshake protocol, not a value # escaping a context manager. response = yield create_oauth_metadata_request(url) # noqa: ASYNC119 prm = await handle_protected_resource_response(response) if prm is None: logger.debug( "Protected resource metadata discovery failed: %s", url, ) continue self.context.protected_resource_metadata = prm self.context.auth_server_url = str(prm.authorization_servers[0]) break asm_urls = build_oauth_authorization_server_metadata_discovery_urls( self.context.auth_server_url, self.context.server_url, ) for url in asm_urls: # ASYNC119: yielding the request to receive its response is # this auth generator's handshake protocol, not a value # escaping a context manager. response = yield create_oauth_metadata_request(url) # noqa: ASYNC119 ok, metadata = await handle_auth_metadata_response(response) if not ok: break if metadata is None: logger.debug("OAuth metadata discovery failed: %s", url) continue self.context.oauth_metadata = metadata await self._persist_oauth_metadata() break except httpx.HTTPError as exc: # Log only the exception type, never its payload — discovery # responses travel the same channel as bearer tokens. logger.debug( "Pre-emptive OAuth metadata discovery for %s raised %s; " "deferring to the SDK auth flow.", self.context.server_url, type(exc).__name__, ) # Delegate to the SDK flow by manually pumping the inner generator so # the HTTP responses httpx feeds back via `auth_flow.asend(response)` # are forwarded into it. A plain `async for` would advance the inner # generator with `__anext__()` (i.e. `asend(None)`), discarding every # response — the SDK's `response = yield request` and refresh-path # `yield refresh_request` would then see `None` and raise # `AttributeError: 'NoneType' object has no attribute 'status_code'`, # surfacing as the `ExceptionGroup` users hit on MCP OAuth login. # httpx primes the flow with `__anext__()`, then drives it with # `asend`/`aclose` (never `athrow`), so forwarding sent values and # closing the inner generator on `GeneratorExit` is sufficient — no # `athrow` forwarding needed. inner = super().async_auth_flow(request) try: # Prime with `anext()` (no response to send yet); thereafter every # resume carries httpx's response back in via `asend`. flow_request = await anext(inner) while True: response = yield flow_request flow_request = await inner.asend(response) except StopAsyncIteration: return finally: await inner.aclose() def build_oauth_provider( *, server_name: str, server_url: str, storage: TokenStorage, extra_auth_params: dict[str, str] | None = None, interactive: bool = True, ui: OAuthInteraction | None = None, ) -> OAuthClientProvider: """Construct an `OAuthClientProvider` for an MCP server. Args: server_name: MCP server name used in re-auth messages. server_url: Remote MCP server URL. storage: Token storage implementation for this server. extra_auth_params: Optional query params for the interactive auth URL. interactive: Whether the provider may prompt on stdin. ui: Interaction surface used for URL display and paste-back input in interactive mode. Returns: A configured `OAuthClientProvider`. """ from deepagents_code.mcp_providers import resolve_provider policy = resolve_provider(server_url) redirect_uri: str | None = None if interactive: if policy.supports_loopback_callback(): fixed = policy.loopback_port() if fixed is not None: port = fixed else: # Reuse the port from a prior DCR registration when available, # so the authorize request's redirect_uri matches what was # registered against the persisted client_id. A fresh random # port on every launch would otherwise invalidate the URI on # the second run and force the server to reject the request. stored = ( storage.stored_loopback_port() if isinstance(storage, FileTokenStorage) else None ) # No reusable port means any persisted registration can't be # paired with the random loopback port we're about to bind. Drop # a stale registration so the handshake re-runs DCR with a # matching redirect URI instead of failing with "invalid or # missing redirect_uri". if ( stored is None and isinstance(storage, FileTokenStorage) and storage.discard_client_info_if_loopback_unusable() ): logger.info( "Discarded a stale MCP client registration for %s " "whose redirect URI can't serve loopback login; the " "handshake will register a fresh client.", server_name, ) port = stored if stored is not None else _choose_loopback_port() callback_server = _LoopbackOAuthCallbackServer(port=port) redirect_uri = callback_server.redirect_uri redirect, callback = _make_loopback_handlers( callback_server=callback_server, extra_auth_params=extra_auth_params, ui=ui, ) else: redirect, callback = _make_paste_back_handlers( extra_auth_params=extra_auth_params, ui=ui, ) else: redirect, callback = _make_reauth_required_handlers(server_name=server_name) metadata = ( policy.client_metadata(redirect_uri=redirect_uri) if redirect_uri is not None else policy.client_metadata() ) return _ExpiryAwareOAuthClientProvider( server_url=server_url, client_metadata=metadata, storage=storage, redirect_handler=redirect, callback_handler=callback, ) async def _run_device_flow( *, device_code_url: str, token_url: str, client_id: str, scope: str | None = None, ui: OAuthInteraction | None = None, ) -> OAuthToken: """Run OAuth 2.0 Device Authorization Grant and return the token. Args: device_code_url: Provider endpoint that issues a device + user code. token_url: Provider endpoint to poll for the access token. client_id: Registered OAuth client ID. scope: Optional space-delimited scope string. ui: Interaction surface used to display the device code. Returns: The issued OAuth access token payload. Raises: RuntimeError: If the device flow fails, times out, or the provider returns an unexpected HTTP status on the device-code request. """ import asyncio import httpx interaction = ui if ui is not None else _default_ui() init_data = {"client_id": client_id} if scope is not None: init_data["scope"] = scope async with httpx.AsyncClient(timeout=30.0) as client: response = await client.post( device_code_url, data=init_data, headers={"Accept": "application/json"}, ) try: response.raise_for_status() except httpx.HTTPStatusError as exc: msg = ( f"Device code request failed: HTTP {response.status_code} " f"from {device_code_url}." ) raise RuntimeError(msg) from exc try: device = _DeviceCodeResponse.model_validate(response.json()) except (ValueError, ValidationError) as exc: msg = ( f"Device code response from {device_code_url} is missing " f"required fields: {exc}" ) raise RuntimeError(msg) from exc await interaction.show_device_code( verification_uri=device.verification_uri, user_code=device.user_code, expires_in=device.expires_in, ) interval = max(device.interval, 1) loop = asyncio.get_running_loop() deadline = loop.time() + device.expires_in while loop.time() < deadline: await asyncio.sleep(interval) token_response = await client.post( token_url, data={ "client_id": client_id, "device_code": device.device_code, "grant_type": "urn:ietf:params:oauth:grant-type:device_code", }, headers={"Accept": "application/json"}, ) # RFC 8628 §3.5 lets providers return `authorization_pending` / # `slow_down` with either a 200 or 400 response. Check the body # before raise_for_status so 400-returning providers work. try: body = token_response.json() except ValueError as exc: # Malformed JSON would otherwise cascade into a confusing # OAuthToken.model_validate({}) error below; log the cause # explicitly so debugging is possible. logger.warning( "Token endpoint %s returned non-JSON body: %s", token_url, exc, ) body = {} err = body.get("error") if err == "authorization_pending": continue if err == "slow_down": interval += 5 continue if err: msg = f"Device flow failed: {err}: {body.get('error_description', '')}" raise RuntimeError(msg) try: token_response.raise_for_status() except httpx.HTTPStatusError as exc: msg = ( f"Token request failed: HTTP {token_response.status_code} " f"from {token_url}." ) raise RuntimeError(msg) from exc try: return OAuthToken.model_validate(body) except ValidationError as exc: msg = ( f"Token response from {token_url} is not a valid " f"OAuth token payload: {exc}" ) raise RuntimeError(msg) from exc msg = "Device flow timed out. Try logging in again." raise RuntimeError(msg) def format_login_failure(exc: BaseException) -> str: """Return a token-safe single-line summary of an OAuth-login exception. OAuth handshakes commonly surface as `ExceptionGroup` (anyio task groups) or as MCP-SDK errors whose `args`/`repr` may include an `OAuthToken`. Never call `str()`/`repr()` on the raw exception for display or logging — instead, prefer a known-safe nested `MCPReauthRequiredError` message, fall back to the messages of our own loopback-related exception types, and degrade to a class-name chain for anything else. Args: exc: Root exception caught from the login worker. Returns: A user-displayable string that is safe to log and to render. """ reauth = find_reauth_required(exc) if reauth is not None: return str(reauth) safe_types = ( _LoopbackCallbackTimeoutError, _LoopbackCallbackUnavailableError, ) if isinstance(exc, safe_types): return f"{type(exc).__name__}: {exc}" parts: list[str] = [] current: BaseException | None = exc visited: set[int] = set() while current is not None and id(current) not in visited: visited.add(id(current)) parts.append(type(current).__name__) if isinstance(current, BaseExceptionGroup): parts.append( "[" + ", ".join(type(e).__name__ for e in current.exceptions[:5]) + "]" ) break current = current.__cause__ or current.__context__ return " -> ".join(parts) if parts else type(exc).__name__ def find_reauth_required(exc: BaseException) -> MCPReauthRequiredError | None: """Find an `MCPReauthRequiredError` anywhere inside `exc`'s tree. Walks `exceptions` (for `ExceptionGroup`), then `__cause__` and `__context__`, tracking visited nodes to terminate on cyclic chains. Args: exc: Root exception to inspect. Returns: The nested `MCPReauthRequiredError`, or `None` if not present. """ visited: set[int] = set() stack: list[BaseException] = [exc] while stack: current = stack.pop() if id(current) in visited: continue visited.add(id(current)) if isinstance(current, MCPReauthRequiredError): return current if isinstance(current, BaseExceptionGroup): stack.extend(current.exceptions) cause = current.__cause__ or current.__context__ if cause is not None: stack.append(cause) return None async def _drive_handshake(connections: dict) -> None: """Open a one-shot MCP session for `connections` to trigger OAuth handshake.""" from langchain_mcp_adapters.client import MultiServerMCPClient client = MultiServerMCPClient(connections=connections) server_name = next(iter(connections)) async with client.session(server_name): pass async def login( *, server_name: str, server_config: McpServerSpec, ui: OAuthInteraction, ) -> None: """Drive OAuth login for `server_name`, persisting tokens on success. Args: server_name: Name of the configured MCP server. server_config: Parsed server config for that entry. ui: Interaction surface for all user prompts and progress messages during the flow. Raises: ValueError: If `server_config` isn't an OAuth http/sse server. RuntimeError: If header env-var interpolation fails, the device flow fails or times out, or the OAuth handshake aborts. """ # noqa: DOC502 - `RuntimeError` surfaces via `resolve_headers` / `_run_device_flow` from langchain_mcp_adapters.sessions import ( SSEConnection, StreamableHttpConnection, ) if server_config.get("auth") != "oauth": msg = ( f"Server '{server_name}' does not use OAuth " '(set "auth": "oauth" in mcpServers).' ) raise ValueError(msg) from deepagents_code.mcp_tools import _resolve_server_type transport = _resolve_server_type(server_config) if transport not in {"http", "sse"}: msg = ( f"Server '{server_name}' uses {transport!r} transport; " "OAuth login is only valid for http/sse." ) raise ValueError(msg) from deepagents_code.mcp_providers import resolve_provider storage = FileTokenStorage(server_name, server_url=server_config["url"]) policy = resolve_provider(server_config["url"]) result = await policy.run_login( server_name=server_name, server_url=server_config["url"], storage=storage, ui=ui, ) success_message = ( f"Logged in to MCP server '{server_name}'. Tokens saved to {storage.path}." ) if result.completed: await ui.show_success(success_message) return provider = build_oauth_provider( server_name=server_name, server_url=server_config["url"], storage=storage, extra_auth_params=result.extra_auth_params or None, ui=ui, ) conn: StreamableHttpConnection | SSEConnection if transport == "http": conn = StreamableHttpConnection( transport="streamable_http", url=server_config["url"], auth=provider, ) else: conn = SSEConnection( transport="sse", url=server_config["url"], auth=provider, ) if "headers" in server_config: conn["headers"] = resolve_headers( server_config["headers"], server_name=server_name, ) await _drive_handshake({server_name: conn}) await ui.show_success(success_message)