From 94ff2f559460f60bb3e70407b68f4fe12010ea0d Mon Sep 17 00:00:00 2001 From: Rohit Behera <126186063+r0h1tb@users.noreply.github.com> Date: Wed, 19 Aug 2026 12:28:58 +0530 Subject: [PATCH 1/2] feat(mcp): per-user OAuth for HTTP MCP connections MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Static credential injection (#2292) covers servers where the user already holds a token. This adds the authorization flow itself, opt-in per connection via `useOAuth` so existing connections are unaffected. Discovery, dynamic client registration and PKCE come from the MCP SDK's OAuthClientProvider. What the SDK cannot supply is the part that only matters on a multi-user server: its TokenStorage takes no arguments, so a single storage shared across requests would hand one user's access token to the next caller. McpOAuthTokenStore keys tokens by (user identifier, server) and hands the SDK a view fixed to one pair. Servers are compared on scheme, host and path, so a token issued for one server mounted on a host is never sent to another mounted beside it. The redirect returns on a route shared by every user, so PendingAuthorizations resolves a callback only for the user who started it; a state belonging to someone else is refused rather than completed against the wrong account. States are single-use, expire, and are URL-safe — random_secret's alphabet contains %, /, = and ?, which do not survive a query string. Co-Authored-By: Claude Opus 5 --- backend/chainlit/mcp_oauth.py | 282 +++++++++++++++++++++++++++ backend/chainlit/server.py | 72 +++++++ backend/chainlit/types.py | 4 + backend/tests/test_mcp_oauth.py | 329 ++++++++++++++++++++++++++++++++ 4 files changed, 687 insertions(+) create mode 100644 backend/chainlit/mcp_oauth.py create mode 100644 backend/tests/test_mcp_oauth.py diff --git a/backend/chainlit/mcp_oauth.py b/backend/chainlit/mcp_oauth.py new file mode 100644 index 0000000000..9c709e52d9 --- /dev/null +++ b/backend/chainlit/mcp_oauth.py @@ -0,0 +1,282 @@ +"""Per-user OAuth token storage for MCP servers. + +The MCP SDK's ``OAuthClientProvider`` handles discovery, dynamic client +registration and PKCE. It delegates persistence to a ``TokenStorage``, whose +methods take no arguments — the SDK expects one storage instance per +authorization context and has no notion of who the end user is. + +That assumption holds for the single-user clients the SDK targets, where the +process belongs to one person. Chainlit is a multi-user server, so a storage +shared across requests would hand one user's access token to another. This +module supplies the missing scope: tokens are keyed by +``(user identifier, server)`` and a scoped view is handed to the SDK. + +Two properties are load-bearing: + +- A token is only ever readable by the user it was issued for, so a reconnect + from another user cannot inherit it. +- A token is bound to the server that issued it, so it is never presented to a + different server (the confused-deputy case). +""" + +from __future__ import annotations + +import asyncio +import secrets +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING, Awaitable, Callable, Dict, Optional, Tuple +from urllib.parse import urlsplit, urlunsplit + +if TYPE_CHECKING: + from mcp.shared.auth import OAuthClientInformationFull, OAuthToken + +__all__ = [ + "McpOAuthTokenStore", + "PendingAuthorizations", + "ScopedTokenStorage", + "build_oauth_provider", + "canonical_server_key", + "pending_authorizations", + "token_store", +] + +_DEFAULT_PORTS = {"http": "80", "https": "443"} + + +def canonical_server_key(server_url: str) -> str: + """Return a stable key identifying the MCP server behind ``server_url``. + + Scheme and host are compared case-insensitively and a default port is + dropped, so the same server reached by equivalent URLs shares one token. + The path is kept: two MCP servers are routinely mounted on one host, and + collapsing them would let a token issued for one be sent to the other. + """ + parts = urlsplit(server_url) + scheme = parts.scheme.lower() + hostname = (parts.hostname or "").lower() + + netloc = hostname + if parts.port is not None and _DEFAULT_PORTS.get(scheme) != str(parts.port): + netloc = f"{hostname}:{parts.port}" + + path = parts.path.rstrip("/") + + # Query and fragment never identify the server, and credentials in the + # netloc must not leak into a dictionary key. + return urlunsplit((scheme, netloc, path, "", "")) + + +class McpOAuthTokenStore: + """Holds MCP OAuth tokens for every user, keyed by user and server. + + The store is deliberately not a cache of one user's tokens: ``scoped()`` + is the only way to reach an entry, and it fixes both halves of the key up + front so a caller cannot read across users by accident. + """ + + def __init__(self) -> None: + self._tokens: Dict[Tuple[str, str], OAuthToken] = {} + self._clients: Dict[Tuple[str, str], OAuthClientInformationFull] = {} + + def scoped(self, user_identifier: str, server_url: str) -> ScopedTokenStorage: + """Return a ``TokenStorage`` bound to one user and one server.""" + if not user_identifier: + raise ValueError( + "An MCP OAuth token cannot be stored without a user identifier." + ) + return ScopedTokenStorage( + self, user_identifier, canonical_server_key(server_url) + ) + + def forget_user(self, user_identifier: str) -> None: + """Drop every token held for a user, e.g. on logout.""" + for mapping in (self._tokens, self._clients): + for key in [k for k in mapping if k[0] == user_identifier]: + del mapping[key] + + def forget_server(self, user_identifier: str, server_url: str) -> None: + """Drop the token a user holds for one server.""" + key = (user_identifier, canonical_server_key(server_url)) + self._tokens.pop(key, None) + self._clients.pop(key, None) + + +class ScopedTokenStorage: + """A ``mcp.client.auth.TokenStorage`` fixed to one user and one server.""" + + def __init__( + self, store: McpOAuthTokenStore, user_identifier: str, server_key: str + ) -> None: + self._store = store + self._key = (user_identifier, server_key) + + @property + def user_identifier(self) -> str: + return self._key[0] + + @property + def server_key(self) -> str: + return self._key[1] + + async def get_tokens(self) -> Optional[OAuthToken]: + return self._store._tokens.get(self._key) + + async def set_tokens(self, tokens: OAuthToken) -> None: + self._store._tokens[self._key] = tokens + + async def get_client_info(self) -> Optional[OAuthClientInformationFull]: + return self._store._clients.get(self._key) + + async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + self._store._clients[self._key] = client_info + + +@dataclass +class _PendingAuthorization: + user_identifier: str + server_key: str + expires_at: float + # The callback can land before anyone awaits it, so the result is kept + # here and the future is only created once a waiter appears — which also + # keeps ``start()`` usable outside a running event loop. + result: Optional[Tuple[str, Optional[str]]] = None + future: Optional[asyncio.Future[Tuple[str, Optional[str]]]] = None + cancelled: bool = False + + def wait(self) -> asyncio.Future[Tuple[str, Optional[str]]]: + if self.future is None: + self.future = asyncio.get_running_loop().create_future() + if self.result is not None: + self.future.set_result(self.result) + elif self.cancelled: + self.future.cancel() + return self.future + + def _complete(self, value: Tuple[str, Optional[str]]) -> None: + self.result = value + if self.future is not None and not self.future.done(): + self.future.set_result(value) + + def _cancel(self) -> None: + self.cancelled = True + if self.future is not None and not self.future.done(): + self.future.cancel() + + +class PendingAuthorizations: + """Correlates an OAuth redirect back to the user who started it. + + The redirect leaves Chainlit and comes back on a shared callback route, so + the ``state`` parameter is the only link to the originating user. Resolving + a callback therefore requires the state to be known *and* to belong to the + user presenting it — otherwise one user could complete another's flow and + have the resulting token stored against them. + + States are single-use and expire, so a leaked redirect URL cannot be + replayed later. + """ + + def __init__(self, ttl_seconds: float = 300.0) -> None: + self._ttl = ttl_seconds + self._pending: Dict[str, _PendingAuthorization] = {} + + def start( + self, user_identifier: str, server_url: str + ) -> Tuple[str, _PendingAuthorization]: + """Register a flow and return its ``state`` and the pending record. + + Await ``pending.wait()`` for the authorization code. + """ + if not user_identifier: + raise ValueError( + "An MCP OAuth flow cannot be started without a user identifier." + ) + self._drop_expired() + # The state travels in the authorization URL, so it must be URL-safe. + state = secrets.token_urlsafe(32) + pending = _PendingAuthorization( + user_identifier=user_identifier, + server_key=canonical_server_key(server_url), + expires_at=time.monotonic() + self._ttl, + ) + self._pending[state] = pending + return state, pending + + def resolve( + self, state: str, user_identifier: str, code: str + ) -> _PendingAuthorization: + """Complete the flow named by ``state`` on behalf of ``user_identifier``. + + Raises KeyError when the state is unknown, already used or expired, and + PermissionError when it belongs to a different user. + """ + self._drop_expired() + pending = self._pending.get(state) + if pending is None: + raise KeyError("Unknown or expired OAuth state.") + if pending.user_identifier != user_identifier: + # Do not consume it: the rightful owner may still complete it. + raise PermissionError("OAuth state belongs to a different user.") + + del self._pending[state] + pending._complete((code, state)) + return pending + + def cancel(self, state: str) -> None: + pending = self._pending.pop(state, None) + if pending: + pending._cancel() + + def __len__(self) -> int: + self._drop_expired() + return len(self._pending) + + def _drop_expired(self) -> None: + now = time.monotonic() + for state in [s for s, p in self._pending.items() if p.expires_at <= now]: + self._pending.pop(state)._cancel() + + +# One store per Chainlit process. Tokens are held in memory: a restart forces a +# re-authorization, which is the safe default for credentials this sensitive. +token_store = McpOAuthTokenStore() +pending_authorizations = PendingAuthorizations() + + +def build_oauth_provider( + user_identifier: str, + server_url: str, + redirect_uri: str, + on_redirect: Callable[[str], Awaitable[None]], + client_name: str = "Chainlit", +): + """Build an SDK OAuth provider scoped to one user and one MCP server. + + The SDK owns discovery, dynamic client registration and PKCE. Chainlit + supplies the two things it cannot know: which user this is, and how to get + the authorization URL in front of them. + """ + from mcp.client.auth import OAuthClientProvider + from mcp.shared.auth import OAuthClientMetadata + from pydantic import AnyUrl + + storage = token_store.scoped(user_identifier, server_url) + state, pending = pending_authorizations.start(user_identifier, server_url) + + async def callback_handler() -> Tuple[str, Optional[str]]: + return await pending.wait() + + provider = OAuthClientProvider( + server_url=server_url, + client_metadata=OAuthClientMetadata( + client_name=client_name, + redirect_uris=[AnyUrl(redirect_uri)], + grant_types=["authorization_code", "refresh_token"], + response_types=["code"], + ), + storage=storage, + redirect_handler=on_redirect, + callback_handler=callback_handler, + ) + return provider, state diff --git a/backend/chainlit/server.py b/backend/chainlit/server.py index cab8fef9e6..517aa2acd1 100644 --- a/backend/chainlit/server.py +++ b/backend/chainlit/server.py @@ -58,6 +58,7 @@ from chainlit.data.acl import is_thread_author from chainlit.logger import logger from chainlit.markdown import get_markdown_str +from chainlit.mcp_oauth import build_oauth_provider, pending_authorizations from chainlit.oauth_providers import get_oauth_provider from chainlit.secret import random_secret from chainlit.types import ( @@ -1289,8 +1290,45 @@ async def call_action( return JSONResponse(content={"success": True, "response": response}) +@router.get("/mcp/oauth/callback") +async def mcp_oauth_callback( + current_user: UserParam, + code: Optional[str] = None, + state: Optional[str] = None, + error: Optional[str] = None, +): + """Complete an MCP authorization started by this user. + + The redirect comes back on a route shared by every user, so the pending + flow is resolved against the caller's identity: a state issued to someone + else is refused rather than completed on their behalf. + """ + if error: + raise HTTPException(status_code=400, detail=error) + + if not code or not state: + raise HTTPException(status_code=400, detail="Missing code or state") + + if not current_user: + raise HTTPException( + status_code=401, detail="MCP OAuth requires an authenticated user." + ) + + try: + pending_authorizations.resolve(state, current_user.identifier, code) + except PermissionError: + raise HTTPException(status_code=403, detail="Unauthorized") + except KeyError: + raise HTTPException( + status_code=400, detail="Unknown or expired authorization request." + ) + + return HTMLResponse("

Authorization complete. You can close this window.

") + + @router.post("/mcp") async def connect_mcp( + request: Request, payload: ConnectMCPRequest, current_user: UserParam, ): @@ -1404,6 +1442,38 @@ async def connect_mcp( # opened it — avoiding the cross-task cancel-scope corruption from # https://github.com/Chainlit/chainlit/issues/2182. + # ── Optional per-user OAuth ── + # + # The SDK provider performs discovery, dynamic client registration and + # PKCE. Chainlit supplies the user scope, so a token obtained here is only + # ever replayed for the same user and the same server. + oauth_provider = None + if getattr(payload, "useOAuth", False): + if not isinstance(mcp_connection, (SseMcpConnection, HttpMcpConnection)): + raise HTTPException( + status_code=400, + detail="OAuth is only supported for HTTP MCP transports.", + ) + user = context.session.user + if not user: + raise HTTPException( + status_code=401, + detail="MCP OAuth requires an authenticated user.", + ) + + async def _emit_authorization_url(auth_url: str) -> None: + await context.emitter.emit( + "mcp_authorization_required", + {"name": payload.name, "url": auth_url}, + ) + + oauth_provider, _oauth_state = build_oauth_provider( + user_identifier=user.identifier, + server_url=mcp_connection.url, + redirect_uri=f"{get_user_facing_url(request.url)}/oauth/callback", + on_redirect=_emit_authorization_url, + ) + ready_event: asyncio.Event = asyncio.Event() stop_event: asyncio.Event = asyncio.Event() # Mutable container to pass the ClientSession back from the bg task. @@ -1418,6 +1488,7 @@ async def _mcp_session_runner() -> None: sse_client( url=mcp_connection.url, headers=mcp_connection.headers, + auth=oauth_provider, ) ) elif isinstance(mcp_connection, StdioMcpConnection): @@ -1434,6 +1505,7 @@ async def _mcp_session_runner() -> None: streamablehttp_client( url=mcp_connection.url, headers=mcp_connection.headers, + auth=oauth_provider, ) ) else: diff --git a/backend/chainlit/types.py b/backend/chainlit/types.py index 2adcc8ad4e..1922bff3a8 100644 --- a/backend/chainlit/types.py +++ b/backend/chainlit/types.py @@ -259,6 +259,8 @@ class ConnectSseMCPRequest(BaseModel): url: str # Optional HTTP headers to forward to the MCP transport (e.g. Authorization) headers: Optional[Dict[str, str]] = None + # Obtain a token via OAuth instead of relying on a static header. + useOAuth: bool = False class ConnectStreamableHttpMCPRequest(BaseModel): @@ -268,6 +270,8 @@ class ConnectStreamableHttpMCPRequest(BaseModel): url: str # Optional HTTP headers to forward to the MCP transport (e.g. Authorization) headers: Dict[str, str] | None = None + # Obtain a token via OAuth instead of relying on a static header. + useOAuth: bool = False ConnectMCPRequest = Union[ diff --git a/backend/tests/test_mcp_oauth.py b/backend/tests/test_mcp_oauth.py new file mode 100644 index 0000000000..d6e439fdb1 --- /dev/null +++ b/backend/tests/test_mcp_oauth.py @@ -0,0 +1,329 @@ +import pytest +from mcp.client.auth import OAuthClientProvider +from mcp.shared.auth import ( + OAuthClientInformationFull, + OAuthClientMetadata, + OAuthToken, +) + +from chainlit.mcp_oauth import ( + McpOAuthTokenStore, + PendingAuthorizations, + ScopedTokenStorage, + canonical_server_key, +) + + +def make_token(access_token: str) -> OAuthToken: + return OAuthToken(access_token=access_token, token_type="Bearer") + + +def make_client_info(client_id: str) -> OAuthClientInformationFull: + return OAuthClientInformationFull( + client_id=client_id, + redirect_uris=["https://app.example.com/mcp/oauth/callback"], + ) + + +class TestCanonicalServerKey: + """The key decides which token is reused, so equivalence must be exact.""" + + @pytest.mark.parametrize( + ("a", "b"), + [ + ("https://mcp.example.com/sse", "https://MCP.Example.com/sse"), + ("https://mcp.example.com/sse", "https://mcp.example.com/sse/"), + ("https://mcp.example.com/sse", "https://mcp.example.com:443/sse"), + ("http://mcp.example.com/sse", "http://mcp.example.com:80/sse"), + ("https://mcp.example.com/sse", "https://mcp.example.com/sse?x=1"), + ("https://mcp.example.com/sse", "https://mcp.example.com/sse#frag"), + ], + ) + def test_equivalent_urls_share_a_key(self, a, b): + assert canonical_server_key(a) == canonical_server_key(b) + + @pytest.mark.parametrize( + ("a", "b"), + [ + # Two MCP servers on one host must not share a token. + ("https://example.com/jira", "https://example.com/confluence"), + # A non-default port is a different endpoint. + ("https://example.com/sse", "https://example.com:8443/sse"), + # http and https are different origins. + ("http://example.com/sse", "https://example.com/sse"), + ("https://a.example.com/sse", "https://b.example.com/sse"), + ], + ) + def test_distinct_servers_get_distinct_keys(self, a, b): + assert canonical_server_key(a) != canonical_server_key(b) + + def test_credentials_are_not_kept_in_the_key(self): + key = canonical_server_key("https://user:pw@example.com/sse") + assert "user" not in key + assert "pw" not in key + assert key == canonical_server_key("https://example.com/sse") + + +class TestScopedTokenStorageProtocol: + def test_is_accepted_by_the_sdk_oauth_provider(self): + """The SDK drives the flow, so the storage has to satisfy it for real. + + ``TokenStorage`` is not ``@runtime_checkable``, so construct the provider + that consumes it instead of asserting ``isinstance``. + """ + scoped = McpOAuthTokenStore().scoped("alice", "https://example.com/sse") + + provider = OAuthClientProvider( + server_url="https://example.com/sse", + client_metadata=OAuthClientMetadata( + client_name="Chainlit", + redirect_uris=["https://app.example.com/mcp/oauth/callback"], + grant_types=["authorization_code", "refresh_token"], + response_types=["code"], + ), + storage=scoped, + ) + + assert provider.context.storage is scoped + + async def test_round_trips_tokens_and_client_info(self): + scoped = McpOAuthTokenStore().scoped("alice", "https://example.com/sse") + + assert await scoped.get_tokens() is None + assert await scoped.get_client_info() is None + + await scoped.set_tokens(make_token("tok-1")) + await scoped.set_client_info(make_client_info("client-1")) + + stored = await scoped.get_tokens() + assert stored is not None + assert stored.access_token == "tok-1" + client = await scoped.get_client_info() + assert client is not None + assert client.client_id == "client-1" + + +class TestIsolation: + """The properties that stop one user's access leaking to another.""" + + async def test_a_token_is_not_readable_by_another_user(self): + store = McpOAuthTokenStore() + server = "https://jira.example.com/mcp" + + await store.scoped("alice", server).set_tokens(make_token("alice-token")) + + # Bob reconnects to the same server and must start unauthenticated. + assert await store.scoped("bob", server).get_tokens() is None + + async def test_a_token_is_not_sent_to_a_different_server(self): + store = McpOAuthTokenStore() + + await store.scoped("alice", "https://jira.example.com/mcp").set_tokens( + make_token("jira-token") + ) + + # Same user, different server: the Jira token must not be offered. + other = store.scoped("alice", "https://evil.example.com/mcp") + assert await other.get_tokens() is None + + async def test_client_registration_is_also_scoped(self): + store = McpOAuthTokenStore() + server = "https://jira.example.com/mcp" + + await store.scoped("alice", server).set_client_info(make_client_info("alice-c")) + + assert await store.scoped("bob", server).get_client_info() is None + + async def test_equivalent_urls_reuse_the_same_token(self): + store = McpOAuthTokenStore() + await store.scoped("alice", "https://Jira.Example.com/mcp/").set_tokens( + make_token("jira-token") + ) + + reconnect = store.scoped("alice", "https://jira.example.com:443/mcp") + stored = await reconnect.get_tokens() + assert stored is not None + assert stored.access_token == "jira-token" + + def test_a_scope_requires_a_user_identifier(self): + store = McpOAuthTokenStore() + with pytest.raises(ValueError, match="user identifier"): + store.scoped("", "https://example.com/mcp") + + +class TestForgetting: + async def test_forget_user_drops_every_server_for_that_user_only(self): + store = McpOAuthTokenStore() + jira = "https://jira.example.com/mcp" + confluence = "https://confluence.example.com/mcp" + + await store.scoped("alice", jira).set_tokens(make_token("a-jira")) + await store.scoped("alice", confluence).set_tokens(make_token("a-conf")) + await store.scoped("bob", jira).set_tokens(make_token("b-jira")) + + store.forget_user("alice") + + assert await store.scoped("alice", jira).get_tokens() is None + assert await store.scoped("alice", confluence).get_tokens() is None + bob = await store.scoped("bob", jira).get_tokens() + assert bob is not None + assert bob.access_token == "b-jira" + + async def test_forget_server_drops_one_server_only(self): + store = McpOAuthTokenStore() + jira = "https://jira.example.com/mcp" + confluence = "https://confluence.example.com/mcp" + + await store.scoped("alice", jira).set_tokens(make_token("a-jira")) + await store.scoped("alice", confluence).set_tokens(make_token("a-conf")) + + store.forget_server("alice", "https://JIRA.example.com/mcp/") + + assert await store.scoped("alice", jira).get_tokens() is None + kept = await store.scoped("alice", confluence).get_tokens() + assert kept is not None + assert kept.access_token == "a-conf" + + async def test_forgetting_an_unknown_entry_is_a_no_op(self): + store = McpOAuthTokenStore() + store.forget_user("nobody") + store.forget_server("nobody", "https://example.com/mcp") + + +class TestPendingAuthorizations: + """The callback route is shared, so `state` must carry the user with it.""" + + async def test_the_originating_user_completes_the_flow(self): + pending = PendingAuthorizations() + state, flow = pending.start("alice", "https://jira.example.com/mcp") + + resolved = pending.resolve(state, "alice", "auth-code") + + assert resolved.user_identifier == "alice" + assert resolved.server_key == "https://jira.example.com/mcp" + assert await flow.wait() == ("auth-code", state) + + async def test_another_user_cannot_complete_someone_elses_flow(self): + pending = PendingAuthorizations() + state, flow = pending.start("alice", "https://jira.example.com/mcp") + + with pytest.raises(PermissionError): + pending.resolve(state, "mallory", "stolen-code") + + # Alice's flow survives so she can still finish it. + assert flow.result is None + pending.resolve(state, "alice", "auth-code") + assert await flow.wait() == ("auth-code", state) + + async def test_a_state_cannot_be_replayed(self): + pending = PendingAuthorizations() + state, _flow = pending.start("alice", "https://jira.example.com/mcp") + pending.resolve(state, "alice", "auth-code") + + with pytest.raises(KeyError): + pending.resolve(state, "alice", "auth-code") + + def test_an_unknown_state_is_rejected(self): + pending = PendingAuthorizations() + with pytest.raises(KeyError): + pending.resolve("never-issued", "alice", "code") + + async def test_states_are_unique_per_flow(self): + pending = PendingAuthorizations() + first, _f1 = pending.start("alice", "https://jira.example.com/mcp") + second, _f2 = pending.start("alice", "https://jira.example.com/mcp") + assert first != second + + async def test_expired_flows_are_dropped(self): + pending = PendingAuthorizations(ttl_seconds=0) + state, flow = pending.start("alice", "https://jira.example.com/mcp") + + with pytest.raises(KeyError): + pending.resolve(state, "alice", "auth-code") + assert len(pending) == 0 + assert flow.cancelled + + async def test_cancel_releases_the_waiter(self): + pending = PendingAuthorizations() + state, flow = pending.start("alice", "https://jira.example.com/mcp") + + pending.cancel(state) + + assert flow.cancelled + assert len(pending) == 0 + + def test_a_flow_requires_a_user_identifier(self): + pending = PendingAuthorizations() + with pytest.raises(ValueError, match="user identifier"): + pending.start("", "https://example.com/mcp") + + +class TestScopeAccessors: + def test_exposes_the_scope_it_is_bound_to(self): + scoped = McpOAuthTokenStore().scoped("alice", "https://Example.com/mcp/") + assert isinstance(scoped, ScopedTokenStorage) + assert scoped.user_identifier == "alice" + assert scoped.server_key == "https://example.com/mcp" + + +class TestMcpOAuthCallbackRoute: + """The callback is reachable by any logged-in user, so it must check scope.""" + + @pytest.fixture + def client_and_user(self): + + from fastapi.testclient import TestClient + + from chainlit.auth import get_current_user + from chainlit.server import app + from chainlit.user import User + + def _as(identifier): + app.dependency_overrides[get_current_user] = lambda: User( + identifier=identifier + ) + return TestClient(app) + + yield _as + app.dependency_overrides.pop(get_current_user, None) + + def test_the_owner_completes_the_flow(self, client_and_user): + from chainlit.mcp_oauth import pending_authorizations + + state, _flow = pending_authorizations.start( + "alice", "https://jira.example.com/mcp" + ) + client = client_and_user("alice") + + res = client.get(f"/mcp/oauth/callback?code=abc&state={state}") + + assert res.status_code == 200 + + def test_another_user_is_refused(self, client_and_user): + from chainlit.mcp_oauth import pending_authorizations + + state, flow = pending_authorizations.start( + "alice", "https://jira.example.com/mcp" + ) + client = client_and_user("mallory") + + res = client.get(f"/mcp/oauth/callback?code=stolen&state={state}") + + assert res.status_code == 403 + # Alice's flow is untouched, so she can still complete it. + assert flow.result is None + + def test_an_unknown_state_is_rejected(self, client_and_user): + client = client_and_user("alice") + res = client.get("/mcp/oauth/callback?code=abc&state=never-issued") + assert res.status_code == 400 + + def test_a_provider_error_is_surfaced(self, client_and_user): + client = client_and_user("alice") + res = client.get("/mcp/oauth/callback?error=access_denied") + assert res.status_code == 400 + + def test_missing_code_is_rejected(self, client_and_user): + client = client_and_user("alice") + res = client.get("/mcp/oauth/callback?state=abc") + assert res.status_code == 400 From 0fc62eb8433ef5b68d506065fc574d282542c435 Mon Sep 17 00:00:00 2001 From: Rohit Behera <126186063+r0h1tb@users.noreply.github.com> Date: Thu, 20 Aug 2026 00:19:28 +0530 Subject: [PATCH 2/2] fix(mcp): correlate the OAuth callback on the SDK's state The pending map was keyed on a state Chainlit generated, but the SDK mints its own, embeds it in the authorization URL and compares the returned value with compare_digest. The browser therefore echoed the SDK's state, resolve() raised KeyError, and no flow could ever complete. Register the owner when the redirect is handed over, keyed by the state already in that URL, and return that same state from the callback handler so the SDK's comparison passes. That also removes the reason for generating a URL-safe state here: the SDK owns state generation now, so random_secret's alphabet is no longer involved. A failed or malformed callback now abandons the caller's own flow, so the waiting connection fails fast instead of hanging until the state expires. Ownership is checked there for the same reason resolve() checks it: otherwise anyone holding a state could cancel someone else's connection. Scope the flow to the HTTP-authenticated caller rather than the session user; the callback route sees the former, so the two must agree or every callback is refused as belonging to someone else. Co-Authored-By: Claude Opus 5 --- backend/chainlit/mcp_oauth.py | 71 +++++++++++--- backend/chainlit/server.py | 28 ++++-- backend/tests/test_mcp_oauth.py | 159 +++++++++++++++++++++++++++++--- 3 files changed, 220 insertions(+), 38 deletions(-) diff --git a/backend/chainlit/mcp_oauth.py b/backend/chainlit/mcp_oauth.py index 9c709e52d9..3dd5e8927e 100644 --- a/backend/chainlit/mcp_oauth.py +++ b/backend/chainlit/mcp_oauth.py @@ -22,11 +22,10 @@ from __future__ import annotations import asyncio -import secrets import time from dataclasses import dataclass from typing import TYPE_CHECKING, Awaitable, Callable, Dict, Optional, Tuple -from urllib.parse import urlsplit, urlunsplit +from urllib.parse import parse_qs, urlsplit, urlunsplit if TYPE_CHECKING: from mcp.shared.auth import OAuthClientInformationFull, OAuthToken @@ -67,6 +66,12 @@ def canonical_server_key(server_url: str) -> str: return urlunsplit((scheme, netloc, path, "", "")) +def extract_state(authorization_url: str) -> Optional[str]: + """Return the ``state`` the SDK embedded in an authorization URL.""" + values = parse_qs(urlsplit(authorization_url).query).get("state") + return values[0] if values else None + + class McpOAuthTokenStore: """Holds MCP OAuth tokens for every user, keyed by user and server. @@ -181,27 +186,30 @@ def __init__(self, ttl_seconds: float = 300.0) -> None: self._ttl = ttl_seconds self._pending: Dict[str, _PendingAuthorization] = {} - def start( - self, user_identifier: str, server_url: str - ) -> Tuple[str, _PendingAuthorization]: - """Register a flow and return its ``state`` and the pending record. + def register( + self, state: str, user_identifier: str, server_url: str + ) -> _PendingAuthorization: + """Record the owner of an in-flight authorization. - Await ``pending.wait()`` for the authorization code. + ``state`` is generated by the SDK and already embedded in the + authorization URL, so it is passed in rather than minted here: the SDK + compares the returned state against its own and would reject anything + else. Await ``pending.wait()`` for the authorization code. """ if not user_identifier: raise ValueError( "An MCP OAuth flow cannot be started without a user identifier." ) + if not state: + raise ValueError("An MCP OAuth flow cannot be registered without a state.") self._drop_expired() - # The state travels in the authorization URL, so it must be URL-safe. - state = secrets.token_urlsafe(32) pending = _PendingAuthorization( user_identifier=user_identifier, server_key=canonical_server_key(server_url), expires_at=time.monotonic() + self._ttl, ) self._pending[state] = pending - return state, pending + return pending def resolve( self, state: str, user_identifier: str, code: str @@ -228,6 +236,21 @@ def cancel(self, state: str) -> None: if pending: pending._cancel() + def abandon(self, state: str, user_identifier: str) -> bool: + """Drop a flow its owner has given up on, releasing the waiting caller. + + Ownership is checked for the same reason ``resolve`` checks it: without + it, anyone holding a state could cancel someone else's connection. + Returns whether anything was dropped. + """ + self._drop_expired() + pending = self._pending.get(state) + if pending is None or pending.user_identifier != user_identifier: + return False + del self._pending[state] + pending._cancel() + return True + def __len__(self) -> int: self._drop_expired() return len(self._pending) @@ -256,18 +279,39 @@ def build_oauth_provider( The SDK owns discovery, dynamic client registration and PKCE. Chainlit supplies the two things it cannot know: which user this is, and how to get the authorization URL in front of them. + + The SDK also owns the ``state``: it mints one, embeds it in the + authorization URL and compares the returned value against it. So the owner + is recorded when the redirect is handed over, keyed by the state already in + that URL — registering a Chainlit-generated state instead would key the map + on a value the browser never echoes back. """ from mcp.client.auth import OAuthClientProvider from mcp.shared.auth import OAuthClientMetadata from pydantic import AnyUrl storage = token_store.scoped(user_identifier, server_url) - state, pending = pending_authorizations.start(user_identifier, server_url) + flow: Dict[str, _PendingAuthorization] = {} + + async def redirect_handler(authorization_url: str) -> None: + state = extract_state(authorization_url) + if not state: + raise ValueError( + "The MCP authorization URL carries no state parameter, so the " + "callback could not be matched to a user." + ) + flow["pending"] = pending_authorizations.register( + state, user_identifier, server_url + ) + await on_redirect(authorization_url) async def callback_handler() -> Tuple[str, Optional[str]]: + pending = flow.get("pending") + if pending is None: # pragma: no cover - the SDK always redirects first + raise RuntimeError("MCP authorization was awaited before it started.") return await pending.wait() - provider = OAuthClientProvider( + return OAuthClientProvider( server_url=server_url, client_metadata=OAuthClientMetadata( client_name=client_name, @@ -276,7 +320,6 @@ async def callback_handler() -> Tuple[str, Optional[str]]: response_types=["code"], ), storage=storage, - redirect_handler=on_redirect, + redirect_handler=redirect_handler, callback_handler=callback_handler, ) - return provider, state diff --git a/backend/chainlit/server.py b/backend/chainlit/server.py index 517aa2acd1..fddc273172 100644 --- a/backend/chainlit/server.py +++ b/backend/chainlit/server.py @@ -1302,18 +1302,26 @@ async def mcp_oauth_callback( The redirect comes back on a route shared by every user, so the pending flow is resolved against the caller's identity: a state issued to someone else is refused rather than completed on their behalf. + + A failed or malformed callback abandons the caller's own flow, so the + connection waiting on it fails fast instead of hanging until the state + expires. """ + if not current_user: + raise HTTPException( + status_code=401, detail="MCP OAuth requires an authenticated user." + ) + if error: + if state: + pending_authorizations.abandon(state, current_user.identifier) raise HTTPException(status_code=400, detail=error) if not code or not state: + if state: + pending_authorizations.abandon(state, current_user.identifier) raise HTTPException(status_code=400, detail="Missing code or state") - if not current_user: - raise HTTPException( - status_code=401, detail="MCP OAuth requires an authenticated user." - ) - try: pending_authorizations.resolve(state, current_user.identifier, code) except PermissionError: @@ -1454,8 +1462,10 @@ async def connect_mcp( status_code=400, detail="OAuth is only supported for HTTP MCP transports.", ) - user = context.session.user - if not user: + # Scope to the HTTP-authenticated caller, which is the identity the + # callback route sees. Falling back to the session user would start a + # flow the callback then refuses as belonging to someone else. + if not current_user: raise HTTPException( status_code=401, detail="MCP OAuth requires an authenticated user.", @@ -1467,8 +1477,8 @@ async def _emit_authorization_url(auth_url: str) -> None: {"name": payload.name, "url": auth_url}, ) - oauth_provider, _oauth_state = build_oauth_provider( - user_identifier=user.identifier, + oauth_provider = build_oauth_provider( + user_identifier=current_user.identifier, server_url=mcp_connection.url, redirect_uri=f"{get_user_facing_url(request.url)}/oauth/callback", on_redirect=_emit_authorization_url, diff --git a/backend/tests/test_mcp_oauth.py b/backend/tests/test_mcp_oauth.py index d6e439fdb1..6a437c12ee 100644 --- a/backend/tests/test_mcp_oauth.py +++ b/backend/tests/test_mcp_oauth.py @@ -10,7 +10,9 @@ McpOAuthTokenStore, PendingAuthorizations, ScopedTokenStorage, + build_oauth_provider, canonical_server_key, + extract_state, ) @@ -195,7 +197,8 @@ class TestPendingAuthorizations: async def test_the_originating_user_completes_the_flow(self): pending = PendingAuthorizations() - state, flow = pending.start("alice", "https://jira.example.com/mcp") + state = "sdk-state-1" + flow = pending.register(state, "alice", "https://jira.example.com/mcp") resolved = pending.resolve(state, "alice", "auth-code") @@ -205,7 +208,8 @@ async def test_the_originating_user_completes_the_flow(self): async def test_another_user_cannot_complete_someone_elses_flow(self): pending = PendingAuthorizations() - state, flow = pending.start("alice", "https://jira.example.com/mcp") + state = "sdk-state-1" + flow = pending.register(state, "alice", "https://jira.example.com/mcp") with pytest.raises(PermissionError): pending.resolve(state, "mallory", "stolen-code") @@ -217,7 +221,8 @@ async def test_another_user_cannot_complete_someone_elses_flow(self): async def test_a_state_cannot_be_replayed(self): pending = PendingAuthorizations() - state, _flow = pending.start("alice", "https://jira.example.com/mcp") + state = "sdk-state-1" + pending.register(state, "alice", "https://jira.example.com/mcp") pending.resolve(state, "alice", "auth-code") with pytest.raises(KeyError): @@ -230,13 +235,14 @@ def test_an_unknown_state_is_rejected(self): async def test_states_are_unique_per_flow(self): pending = PendingAuthorizations() - first, _f1 = pending.start("alice", "https://jira.example.com/mcp") - second, _f2 = pending.start("alice", "https://jira.example.com/mcp") - assert first != second + pending.register("state-a", "alice", "https://jira.example.com/mcp") + pending.register("state-b", "alice", "https://jira.example.com/mcp") + assert len(pending) == 2 async def test_expired_flows_are_dropped(self): pending = PendingAuthorizations(ttl_seconds=0) - state, flow = pending.start("alice", "https://jira.example.com/mcp") + state = "sdk-state-1" + flow = pending.register(state, "alice", "https://jira.example.com/mcp") with pytest.raises(KeyError): pending.resolve(state, "alice", "auth-code") @@ -245,7 +251,8 @@ async def test_expired_flows_are_dropped(self): async def test_cancel_releases_the_waiter(self): pending = PendingAuthorizations() - state, flow = pending.start("alice", "https://jira.example.com/mcp") + state = "sdk-state-1" + flow = pending.register(state, "alice", "https://jira.example.com/mcp") pending.cancel(state) @@ -255,7 +262,12 @@ async def test_cancel_releases_the_waiter(self): def test_a_flow_requires_a_user_identifier(self): pending = PendingAuthorizations() with pytest.raises(ValueError, match="user identifier"): - pending.start("", "https://example.com/mcp") + pending.register("s", "", "https://example.com/mcp") + + def test_a_flow_requires_a_state(self): + pending = PendingAuthorizations() + with pytest.raises(ValueError, match="state"): + pending.register("", "alice", "https://example.com/mcp") class TestScopeAccessors: @@ -290,9 +302,8 @@ def _as(identifier): def test_the_owner_completes_the_flow(self, client_and_user): from chainlit.mcp_oauth import pending_authorizations - state, _flow = pending_authorizations.start( - "alice", "https://jira.example.com/mcp" - ) + state = "sdk-owner-state" + pending_authorizations.register(state, "alice", "https://jira.example.com/mcp") client = client_and_user("alice") res = client.get(f"/mcp/oauth/callback?code=abc&state={state}") @@ -302,8 +313,9 @@ def test_the_owner_completes_the_flow(self, client_and_user): def test_another_user_is_refused(self, client_and_user): from chainlit.mcp_oauth import pending_authorizations - state, flow = pending_authorizations.start( - "alice", "https://jira.example.com/mcp" + state = "sdk-mallory-state" + flow = pending_authorizations.register( + state, "alice", "https://jira.example.com/mcp" ) client = client_and_user("mallory") @@ -324,6 +336,123 @@ def test_a_provider_error_is_surfaced(self, client_and_user): assert res.status_code == 400 def test_missing_code_is_rejected(self, client_and_user): + """Use a real state so the 400 can only come from the missing code.""" + from chainlit.mcp_oauth import pending_authorizations + + state = "sdk-missing-code-state" + flow = pending_authorizations.register( + state, "alice", "https://jira.example.com/mcp" + ) client = client_and_user("alice") - res = client.get("/mcp/oauth/callback?state=abc") + + res = client.get(f"/mcp/oauth/callback?state={state}") + assert res.status_code == 400 + assert res.json()["detail"] == "Missing code or state" + # The waiting connection is released rather than left hanging. + assert flow.cancelled + + def test_a_provider_error_abandons_the_flow(self, client_and_user): + from chainlit.mcp_oauth import pending_authorizations + + state = "sdk-error-state" + flow = pending_authorizations.register( + state, "alice", "https://jira.example.com/mcp" + ) + client = client_and_user("alice") + + res = client.get(f"/mcp/oauth/callback?error=access_denied&state={state}") + + assert res.status_code == 400 + assert flow.cancelled + + def test_an_error_cannot_abandon_someone_elses_flow(self, client_and_user): + from chainlit.mcp_oauth import pending_authorizations + + state = "sdk-victim-state" + flow = pending_authorizations.register( + state, "alice", "https://jira.example.com/mcp" + ) + client = client_and_user("mallory") + + res = client.get(f"/mcp/oauth/callback?error=access_denied&state={state}") + + assert res.status_code == 400 + # Alice's connection is untouched. + assert not flow.cancelled + + +class TestProviderStateCorrelation: + """The SDK mints the state, so the pending map must be keyed on that one. + + Registering a Chainlit-generated state would key the map on a value the + browser never echoes back, and every callback would 404 the flow. + """ + + def test_extract_state_reads_the_url(self): + assert ( + extract_state("https://as.example.com/authorize?client_id=x&state=s1") + == "s1" + ) + assert extract_state("https://as.example.com/authorize?client_id=x") is None + + async def test_the_flow_is_registered_under_the_sdk_state(self): + """Drive the provider's redirect handler and check the key it used.""" + from chainlit.mcp_oauth import pending_authorizations + + seen: list[str] = [] + + async def on_redirect(url: str) -> None: + seen.append(url) + + provider = build_oauth_provider( + user_identifier="alice", + server_url="https://jira.example.com/mcp", + redirect_uri="https://app.example.com/mcp/oauth/callback", + on_redirect=on_redirect, + ) + + sdk_url = "https://as.example.com/authorize?client_id=c&state=sdk-generated" + await provider.context.redirect_handler(sdk_url) + + assert seen == [sdk_url] + # The callback route looks the flow up by the state in that URL. + resolved = pending_authorizations.resolve("sdk-generated", "alice", "code-1") + assert resolved.user_identifier == "alice" + + async def test_the_callback_handler_returns_the_sdk_state(self): + """The SDK compares the returned state with compare_digest.""" + from chainlit.mcp_oauth import pending_authorizations + + async def on_redirect(url: str) -> None: + return None + + provider = build_oauth_provider( + user_identifier="alice", + server_url="https://jira.example.com/mcp", + redirect_uri="https://app.example.com/mcp/oauth/callback", + on_redirect=on_redirect, + ) + + await provider.context.redirect_handler( + "https://as.example.com/authorize?state=sdk-echo" + ) + pending_authorizations.resolve("sdk-echo", "alice", "the-code") + + assert await provider.context.callback_handler() == ("the-code", "sdk-echo") + + async def test_a_url_without_state_is_refused(self): + async def on_redirect(url: str) -> None: + return None + + provider = build_oauth_provider( + user_identifier="alice", + server_url="https://jira.example.com/mcp", + redirect_uri="https://app.example.com/mcp/oauth/callback", + on_redirect=on_redirect, + ) + + with pytest.raises(ValueError, match="state"): + await provider.context.redirect_handler( + "https://as.example.com/authorize?client_id=c" + )