diff --git a/backend/chainlit/mcp_oauth.py b/backend/chainlit/mcp_oauth.py new file mode 100644 index 0000000000..3dd5e8927e --- /dev/null +++ b/backend/chainlit/mcp_oauth.py @@ -0,0 +1,325 @@ +"""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 time +from dataclasses import dataclass +from typing import TYPE_CHECKING, Awaitable, Callable, Dict, Optional, Tuple +from urllib.parse import parse_qs, 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, "", "")) + + +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. + + 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 register( + self, state: str, user_identifier: str, server_url: str + ) -> _PendingAuthorization: + """Record the owner of an in-flight authorization. + + ``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() + pending = _PendingAuthorization( + user_identifier=user_identifier, + server_key=canonical_server_key(server_url), + expires_at=time.monotonic() + self._ttl, + ) + self._pending[state] = pending + return 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 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) + + 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. + + 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) + 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() + + return 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=redirect_handler, + callback_handler=callback_handler, + ) diff --git a/backend/chainlit/server.py b/backend/chainlit/server.py index 4d53390e50..6f0b7c2f07 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 ( @@ -1305,8 +1306,53 @@ 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. + + 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") + + 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, ): @@ -1420,6 +1466,40 @@ 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.", + ) + # 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.", + ) + + async def _emit_authorization_url(auth_url: str) -> None: + await context.emitter.emit( + "mcp_authorization_required", + {"name": payload.name, "url": auth_url}, + ) + + 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, + ) + ready_event: asyncio.Event = asyncio.Event() stop_event: asyncio.Event = asyncio.Event() # Mutable container to pass the ClientSession back from the bg task. @@ -1434,6 +1514,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): @@ -1450,6 +1531,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..6a437c12ee --- /dev/null +++ b/backend/tests/test_mcp_oauth.py @@ -0,0 +1,458 @@ +import pytest +from mcp.client.auth import OAuthClientProvider +from mcp.shared.auth import ( + OAuthClientInformationFull, + OAuthClientMetadata, + OAuthToken, +) + +from chainlit.mcp_oauth import ( + McpOAuthTokenStore, + PendingAuthorizations, + ScopedTokenStorage, + build_oauth_provider, + canonical_server_key, + extract_state, +) + + +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 = "sdk-state-1" + flow = pending.register(state, "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 = "sdk-state-1" + flow = pending.register(state, "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 = "sdk-state-1" + pending.register(state, "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() + 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 = "sdk-state-1" + flow = pending.register(state, "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 = "sdk-state-1" + flow = pending.register(state, "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.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: + 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 = "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}") + + assert res.status_code == 200 + + def test_another_user_is_refused(self, client_and_user): + from chainlit.mcp_oauth import pending_authorizations + + state = "sdk-mallory-state" + flow = pending_authorizations.register( + state, "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): + """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(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" + )