From 5c689ebbf457885458b25eff9f42cf4490e71907 Mon Sep 17 00:00:00 2001 From: julia-timbal Date: Thu, 23 Jul 2026 15:42:54 +0200 Subject: [PATCH] feat(mcp): add OAuth 2.0 browser login for HTTP MCP servers Wire MCP SDK OAuthClientProvider into MCPServer with local callback server, browser redirect, and pluggable token storage so users can authenticate via the server's OAuth portal --- python/tests/core/test_mcp_oauth.py | 206 ++++++++++++++++++ python/timbal/core/__init__.py | 3 + python/timbal/core/mcp.py | 34 ++- python/timbal/core/mcp_oauth.py | 311 ++++++++++++++++++++++++++++ 4 files changed, 553 insertions(+), 1 deletion(-) create mode 100644 python/tests/core/test_mcp_oauth.py create mode 100644 python/timbal/core/mcp_oauth.py diff --git a/python/tests/core/test_mcp_oauth.py b/python/tests/core/test_mcp_oauth.py new file mode 100644 index 00000000..f8e5a307 --- /dev/null +++ b/python/tests/core/test_mcp_oauth.py @@ -0,0 +1,206 @@ +import asyncio +import socket + +import httpx +import pytest +from pydantic import ValidationError +from timbal.core.mcp import MCPServer +from timbal.core.mcp_oauth import ( + FileTokenStorage, + InMemoryTokenStorage, + OAuth, + _CallbackServer, +) + +try: + from mcp.client.auth import OAuthClientProvider + from mcp.shared.auth import OAuthClientInformationFull, OAuthToken +except ImportError: # pragma: no cover + OAuthClientProvider = None + +pytestmark = pytest.mark.skipif(OAuthClientProvider is None, reason="mcp package not installed") + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("localhost", 0)) + return s.getsockname()[1] + + +def _sample_tokens() -> "OAuthToken": + return OAuthToken(access_token="abc123", token_type="Bearer", expires_in=3600, refresh_token="ref456") + + +def _sample_client_info() -> "OAuthClientInformationFull": + return OAuthClientInformationFull( + client_id="client-1", + redirect_uris=["http://localhost:51820/callback"], + ) + + +class TestInMemoryTokenStorage: + @pytest.mark.asyncio + async def test_roundtrip(self): + storage = InMemoryTokenStorage() + assert await storage.get_tokens() is None + assert await storage.get_client_info() is None + + await storage.set_tokens(_sample_tokens()) + await storage.set_client_info(_sample_client_info()) + + assert (await storage.get_tokens()).access_token == "abc123" + assert (await storage.get_client_info()).client_id == "client-1" + + +class TestFileTokenStorage: + @pytest.mark.asyncio + async def test_roundtrip_and_persistence(self, tmp_path): + url = "https://mcp.example.com/mcp" + storage = FileTokenStorage(url, storage_dir=tmp_path) + assert await storage.get_tokens() is None + + await storage.set_tokens(_sample_tokens()) + await storage.set_client_info(_sample_client_info()) + + # A fresh instance pointing at the same URL sees the same credentials. + storage2 = FileTokenStorage(url, storage_dir=tmp_path) + assert (await storage2.get_tokens()).refresh_token == "ref456" + assert (await storage2.get_client_info()).client_id == "client-1" + + @pytest.mark.asyncio + async def test_isolated_per_server_url(self, tmp_path): + a = FileTokenStorage("https://a.example.com/mcp", storage_dir=tmp_path) + b = FileTokenStorage("https://b.example.com/mcp", storage_dir=tmp_path) + await a.set_tokens(_sample_tokens()) + assert await b.get_tokens() is None + + @pytest.mark.asyncio + async def test_clear(self, tmp_path): + storage = FileTokenStorage("https://mcp.example.com/mcp", storage_dir=tmp_path) + await storage.set_tokens(_sample_tokens()) + storage.clear() + assert await storage.get_tokens() is None + + +class TestCallbackServer: + @pytest.mark.asyncio + async def test_captures_code_and_state(self): + server = _CallbackServer("localhost", 0) + await server.start() + port = server._server.sockets[0].getsockname()[1] + + async def hit_callback(): + async with httpx.AsyncClient() as client: + return await client.get(f"http://localhost:{port}/callback?code=the-code&state=the-state") + + response, (code, state) = await asyncio.gather(hit_callback(), server.wait_for_code(timeout=5)) + assert response.status_code == 200 + assert b"Authentication complete" in response.content + assert code == "the-code" + assert state == "the-state" + + @pytest.mark.asyncio + async def test_error_callback_raises(self): + server = _CallbackServer("localhost", 0) + await server.start() + port = server._server.sockets[0].getsockname()[1] + + async def hit_callback(): + async with httpx.AsyncClient() as client: + return await client.get( + f"http://localhost:{port}/callback?error=access_denied&error_description=nope" + ) + + hit_task = asyncio.create_task(hit_callback()) + with pytest.raises(RuntimeError, match="access_denied"): + await server.wait_for_code(timeout=5) + response = await hit_task + assert response.status_code == 400 + + @pytest.mark.asyncio + async def test_unknown_path_is_404_and_flow_still_completes(self): + server = _CallbackServer("localhost", 0) + await server.start() + port = server._server.sockets[0].getsockname()[1] + + async def hit(): + async with httpx.AsyncClient() as client: + r404 = await client.get(f"http://localhost:{port}/favicon.ico") + r200 = await client.get(f"http://localhost:{port}/callback?code=c&state=s") + return r404, r200 + + (r404, r200), (code, state) = await asyncio.gather(hit(), server.wait_for_code(timeout=5)) + assert r404.status_code == 404 + assert r200.status_code == 200 + assert (code, state) == ("c", "s") + + @pytest.mark.asyncio + async def test_timeout(self): + server = _CallbackServer("localhost", 0) + await server.start() + with pytest.raises(TimeoutError): + await server.wait_for_code(timeout=0.05) + + +class TestOAuthConfig: + def test_build_provider(self, tmp_path): + oauth = OAuth(scopes="read write", callback_port=51999, storage_dir=tmp_path) + provider = oauth.build_provider("https://mcp.example.com/mcp") + assert isinstance(provider, OAuthClientProvider) + + metadata = provider.context.client_metadata + assert str(metadata.redirect_uris[0]) == "http://localhost:51999/callback" + assert metadata.scope == "read write" + assert metadata.token_endpoint_auth_method == "none" + assert isinstance(provider.context.storage, FileTokenStorage) + + def test_custom_storage_is_used(self): + storage = InMemoryTokenStorage() + provider = OAuth(storage=storage).build_provider("https://mcp.example.com/mcp") + assert provider.context.storage is storage + + @pytest.mark.asyncio + async def test_redirect_and_callback_handlers_complete_flow(self, tmp_path, monkeypatch): + """Simulate the browser: redirect handler opens the portal, then the portal + redirects the user back to the local callback URL with the auth code.""" + opened_urls: list[str] = [] + monkeypatch.setattr("webbrowser.open", lambda url: opened_urls.append(url) or True) + + port = _free_port() + oauth = OAuth(callback_port=port, storage_dir=tmp_path, flow_timeout=5) + provider = oauth.build_provider("https://mcp.example.com/mcp") + + # SDK calls redirect_handler first (opens browser)... + await provider.context.redirect_handler("https://auth.example.com/authorize?state=xyz") + assert opened_urls == ["https://auth.example.com/authorize?state=xyz"] + + # ...then awaits callback_handler while the user logs in. + async def simulate_browser_redirect(): + async with httpx.AsyncClient() as client: + return await client.get(f"http://localhost:{port}/callback?code=ok&state=xyz") + + response, (code, state) = await asyncio.gather( + simulate_browser_redirect(), + provider.context.callback_handler(), + ) + assert response.status_code == 200 + assert (code, state) == ("ok", "xyz") + + +class TestMCPServerAuthValidation: + def test_auth_string_coerces_to_oauth(self): + server = MCPServer(transport="http", url="https://mcp.example.com/mcp", auth="oauth") + assert isinstance(server.auth, OAuth) + + def test_auth_instance_preserved(self): + oauth = OAuth(scopes="read") + server = MCPServer(transport="http", url="https://mcp.example.com/mcp", auth=oauth) + assert server.auth is oauth + + def test_auth_rejected_for_stdio(self): + with pytest.raises(ValidationError, match="only supported for http"): + MCPServer(transport="stdio", command="echo", auth="oauth") + + def test_no_auth_by_default(self): + server = MCPServer(transport="http", url="https://mcp.example.com/mcp") + assert server.auth is None diff --git a/python/timbal/core/__init__.py b/python/timbal/core/__init__.py index ee10ee51..7c4549f6 100644 --- a/python/timbal/core/__init__.py +++ b/python/timbal/core/__init__.py @@ -8,6 +8,7 @@ from .agent import Agent from .fallback_model import FallbackModel, ModelEntry from .mcp import MCPServer + from .mcp_oauth import OAuth from .memory_compaction import MemoryCompactor # noqa: F401 - type alias from .skill import Skill from .test_model import TestModel @@ -20,6 +21,7 @@ "FallbackModel", "MCPServer", "ModelEntry", + "OAuth", "Skill", "TestModel", "Tool", @@ -32,6 +34,7 @@ "FallbackModel": ".fallback_model", "MCPServer": ".mcp", "ModelEntry": ".fallback_model", + "OAuth": ".mcp_oauth", "Skill": ".skill", "TestModel": ".test_model", "Tool": ".tool", diff --git a/python/timbal/core/mcp.py b/python/timbal/core/mcp.py index 8f5bc6fa..feb4cb36 100644 --- a/python/timbal/core/mcp.py +++ b/python/timbal/core/mcp.py @@ -20,6 +20,7 @@ from ..types.content import FileContent, TextContent from ..types.file import File from ..types.message import Message +from .mcp_oauth import OAuth from .runnable import Runnable from .tool import Tool from .tool_set import ToolSet @@ -99,6 +100,15 @@ class MCPServer(ToolSet): MCPServer(transport="stdio", command="npx", args=["-y", "@modelcontextprotocol/server-filesystem", "."]) MCPServer(transport="http", url="https://api.timbal.ai/mcp") + + HTTP servers that require user login support OAuth 2.0 out of the box: + + MCPServer(transport="http", url="https://mcp.example.com/mcp", auth="oauth") + MCPServer(transport="http", url="https://mcp.example.com/mcp", auth=OAuth(scopes="read write")) + + On the first 401 response the user's browser opens at the server's + authorization portal; tokens are persisted under ``~/.timbal/mcp_oauth/`` + and refreshed automatically on subsequent runs. """ name: str | None = None @@ -119,6 +129,15 @@ class MCPServer(ToolSet): url: str | None = None headers: dict[str, str] = Field(default_factory=dict) + auth: OAuth | Literal["oauth"] | None = None + """OAuth 2.0 configuration for HTTP servers that require user login. + + Pass ``"oauth"`` for the default browser-based flow, or an :class:`OAuth` + instance to customize scopes, callback port, token storage, etc. When the + server responds 401, the OAuth portal opens in the user's browser and the + resulting tokens are stored and refreshed automatically. + """ + _session: Any | None = PrivateAttr(default=None) _tools_cache: list[Runnable] | None = PrivateAttr(default=None) _context: Any | None = PrivateAttr(default=None) @@ -135,9 +154,13 @@ def _validate_transport_fields(self) -> "MCPServer": if self.transport == "stdio": if not self.command: raise ValueError("'command' is required for stdio transport") + if self.auth is not None: + raise ValueError("'auth' is only supported for http transport") elif self.transport == "http": if not self.url: raise ValueError("'url' is required for http transport") + if self.auth == "oauth": + self.auth = OAuth() return self @asynccontextmanager @@ -157,7 +180,16 @@ async def _connect_stdio(self): @asynccontextmanager async def _connect_http(self): assert self.url is not None - async with httpx.AsyncClient(headers=self.headers if self.headers else None) as http_client: + auth = self.auth.build_provider(self.url) if isinstance(self.auth, OAuth) else None + # Mirror the MCP SDK's recommended client defaults (30s timeout, + # follow_redirects) — OAuth discovery/authorization endpoints often redirect. + http_client = httpx.AsyncClient( + headers=self.headers if self.headers else None, + auth=auth, + timeout=httpx.Timeout(30.0), + follow_redirects=True, + ) + async with http_client: async with streamable_http_client(self.url, http_client=http_client) as (read, write, _): async with ClientSession(read, write) as session: await session.initialize() diff --git a/python/timbal/core/mcp_oauth.py b/python/timbal/core/mcp_oauth.py new file mode 100644 index 00000000..7b024981 --- /dev/null +++ b/python/timbal/core/mcp_oauth.py @@ -0,0 +1,311 @@ +"""OAuth 2.0 authorization for MCP servers. + +Implements the client side of the MCP authorization spec on top of the +official SDK's ``OAuthClientProvider`` (authorization code flow with PKCE, +metadata discovery, dynamic client registration, automatic token refresh). + +This module supplies the three pieces the SDK leaves to the application: + +- **Browser redirect**: when the server responds 401, the user's default + browser is opened at the authorization portal so they can log in. +- **Callback capture**: a one-shot local HTTP server listens on + ``http://localhost:{port}/callback`` and captures the authorization code + the portal redirects back with. +- **Token persistence**: ``FileTokenStorage`` keeps tokens and the registered + client info under ``~/.timbal/mcp_oauth/`` (keyed by server URL) so users + log in once, not on every run. ``InMemoryTokenStorage`` is available for + ephemeral sessions. + +Usage: + + from timbal.core import MCPServer + from timbal.core.mcp_oauth import OAuth + + server = MCPServer( + transport="http", + url="https://mcp.example.com/mcp", + auth=OAuth(), # or auth="oauth" for all defaults + ) +""" + +import asyncio +import hashlib +import json +import webbrowser +from pathlib import Path +from typing import Any +from urllib.parse import parse_qs, urlparse + +import structlog +from mcp.client.auth import OAuthClientProvider +from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken +from pydantic import BaseModel, ConfigDict + +logger = structlog.get_logger("timbal.core.mcp_oauth") + +DEFAULT_STORAGE_DIR = Path.home() / ".timbal" / "mcp_oauth" + +_CALLBACK_SUCCESS_HTML = b"""\ + + +Authentication complete + +

Authentication complete

+

You can close this window and return to your application.

+ + +""" + +_CALLBACK_ERROR_HTML = b"""\ + + +Authentication failed + +

Authentication failed

+

You can close this window and retry from your application.

+ + +""" + + +class InMemoryTokenStorage: + """Ephemeral token storage. Tokens are lost when the process exits.""" + + def __init__(self) -> None: + self._tokens: OAuthToken | None = None + self._client_info: OAuthClientInformationFull | None = None + + async def get_tokens(self) -> OAuthToken | None: + return self._tokens + + async def set_tokens(self, tokens: OAuthToken) -> None: + self._tokens = tokens + + async def get_client_info(self) -> OAuthClientInformationFull | None: + return self._client_info + + async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + self._client_info = client_info + + +class FileTokenStorage: + """Persist tokens and client registration to disk, keyed by server URL. + + Each server gets its own JSON file under ``storage_dir`` named by a hash + of the server URL, created with owner-only permissions (0600). + """ + + def __init__(self, server_url: str, storage_dir: Path | None = None) -> None: + self.server_url = server_url + self.storage_dir = storage_dir or DEFAULT_STORAGE_DIR + digest = hashlib.sha256(server_url.encode()).hexdigest()[:16] + host = urlparse(server_url).hostname or "server" + self.path = self.storage_dir / f"{host}_{digest}.json" + + def _read(self) -> dict: + try: + return json.loads(self.path.read_text()) + except (FileNotFoundError, json.JSONDecodeError): + return {} + + def _write(self, data: dict) -> None: + self.storage_dir.mkdir(parents=True, exist_ok=True) + self.path.touch(mode=0o600, exist_ok=True) + self.path.write_text(json.dumps(data, indent=2)) + + async def get_tokens(self) -> OAuthToken | None: + tokens = self._read().get("tokens") + return OAuthToken.model_validate(tokens) if tokens else None + + async def set_tokens(self, tokens: OAuthToken) -> None: + data = self._read() + data["tokens"] = tokens.model_dump(mode="json", exclude_none=True) + self._write(data) + + async def get_client_info(self) -> OAuthClientInformationFull | None: + client_info = self._read().get("client_info") + return OAuthClientInformationFull.model_validate(client_info) if client_info else None + + async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + data = self._read() + data["client_info"] = client_info.model_dump(mode="json", exclude_none=True) + self._write(data) + + def clear(self) -> None: + """Delete stored credentials for this server (forces re-authentication).""" + self.path.unlink(missing_ok=True) + + +class _CallbackServer: + """One-shot local HTTP server that captures the OAuth authorization callback. + + Listens on ``http://{host}:{port}{path}`` for a single redirect carrying + ``code`` and ``state`` query parameters, serves a small confirmation page, + and hands the values back to the OAuth flow. + """ + + def __init__(self, host: str, port: int, path: str = "/callback") -> None: + self.host = host + self.port = port + self.path = path + self._server: asyncio.AbstractServer | None = None + self._result: asyncio.Future[tuple[str, str | None]] | None = None + + @property + def redirect_uri(self) -> str: + return f"http://{self.host}:{self.port}{self.path}" + + async def start(self) -> None: + if self._server is not None: + return + self._result = asyncio.get_running_loop().create_future() + self._server = await asyncio.start_server(self._handle_connection, self.host, self.port) + logger.debug("OAuth callback server listening", redirect_uri=self.redirect_uri) + + async def _handle_connection(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + request_line = await reader.readline() + # Drain headers so the browser doesn't see a reset connection. + while await reader.readline() not in (b"\r\n", b"\n", b""): + pass + + parts = request_line.decode("latin-1").split(" ") + target = parts[1] if len(parts) >= 2 else "/" + parsed = urlparse(target) + + if parsed.path != self.path: + writer.write(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + await writer.drain() + return + + params = parse_qs(parsed.query) + code = params.get("code", [None])[0] + state = params.get("state", [None])[0] + error = params.get("error", [None])[0] + + if error or not code: + body = _CALLBACK_ERROR_HTML + status = b"HTTP/1.1 400 Bad Request" + if self._result is not None and not self._result.done(): + description = params.get("error_description", [None])[0] + message = f"OAuth authorization failed: {error or 'no authorization code returned'}" + if description: + message = f"{message} ({description})" + self._result.set_exception(RuntimeError(message)) + else: + body = _CALLBACK_SUCCESS_HTML + status = b"HTTP/1.1 200 OK" + if self._result is not None and not self._result.done(): + self._result.set_result((code, state)) + + writer.write( + status + + b"\r\nContent-Type: text/html; charset=utf-8" + + f"\r\nContent-Length: {len(body)}".encode() + + b"\r\nConnection: close\r\n\r\n" + + body + ) + await writer.drain() + finally: + writer.close() + + async def wait_for_code(self, timeout: float) -> tuple[str, str | None]: + assert self._result is not None, "start() must be called before wait_for_code()" + try: + return await asyncio.wait_for(asyncio.shield(self._result), timeout) + except TimeoutError: + raise TimeoutError(f"Timed out after {timeout:.0f}s waiting for the OAuth callback. ") from None + finally: + await self.stop() + + async def stop(self) -> None: + if self._server is not None: + self._server.close() + await self._server.wait_closed() + self._server = None + self._result = None + + +class OAuth(BaseModel): + """OAuth 2.0 configuration for an HTTP :class:`~timbal.core.mcp.MCPServer`. + + When the MCP server responds 401, the flow kicks in automatically: the + user's browser opens at the server's authorization portal, they log in, + and the resulting tokens are stored (and refreshed) transparently. All + fields have sensible defaults, so ``OAuth()`` works out of the box. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True) + + client_name: str = "Timbal Agent" + """Client name sent during dynamic client registration (shown on consent screens).""" + + scopes: str | None = None + """Space-separated OAuth scopes to request. If None, the server's advertised scopes are used.""" + + callback_host: str = "localhost" + """Host the local callback server binds to. Must match what the browser can reach.""" + + callback_port: int = 51820 + """Port for the local callback server. The redirect URI becomes http://{host}:{port}/callback.""" + + callback_path: str = "/callback" + """Path component of the redirect URI.""" + + flow_timeout: float = 300.0 + """Seconds to wait for the user to complete the login in the browser.""" + + client_metadata_url: str | None = None + """Optional URL-based client ID (CIMD). Skips dynamic registration on servers that support it.""" + + storage: Any = None + """Custom TokenStorage. Defaults to FileTokenStorage under ~/.timbal/mcp_oauth/.""" + + storage_dir: Path | None = None + """Directory for the default FileTokenStorage. Ignored when a custom storage is given.""" + + open_browser: bool = True + """Open the system browser automatically. The URL is always logged for headless environments.""" + + def build_provider(self, server_url: str) -> OAuthClientProvider: + """Build the httpx auth provider that drives the OAuth flow for ``server_url``.""" + callback_server = _CallbackServer(self.callback_host, self.callback_port, self.callback_path) + storage = self.storage or FileTokenStorage(server_url, storage_dir=self.storage_dir) + + client_metadata = OAuthClientMetadata( + client_name=self.client_name, + redirect_uris=[callback_server.redirect_uri], + grant_types=["authorization_code", "refresh_token"], + response_types=["code"], + token_endpoint_auth_method="none", # public client with PKCE + scope=self.scopes, + ) + + async def redirect_handler(authorization_url: str) -> None: + # Start listening *before* opening the browser so the redirect + # can't race the server startup. + await callback_server.start() + logger.info( + "MCP server requires authentication. Complete the login in your browser.", + url=authorization_url, + ) + if self.open_browser: + opened = webbrowser.open(authorization_url) + if not opened: + logger.warning( + "Could not open a browser automatically. Open the authorization URL manually.", + url=authorization_url, + ) + + async def callback_handler() -> tuple[str, str | None]: + return await callback_server.wait_for_code(self.flow_timeout) + + return OAuthClientProvider( + server_url=server_url, + client_metadata=client_metadata, + storage=storage, + redirect_handler=redirect_handler, + callback_handler=callback_handler, + timeout=self.flow_timeout, + client_metadata_url=self.client_metadata_url, + )