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,
+ )