feat(mcp): add OAuth 2.0 browser login for HTTP MCP servers

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

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",

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()