"""Configured authentication sources retain stable principal identities.""" import json from collections.abc import AsyncGenerator from types import SimpleNamespace from typing import Any, cast from urllib.parse import parse_qs import httpx2 import pytest from fastmcp_tasks.context import ( TaskContextSnapshot, _apply_snapshot_to_context, get_task_scope, ) from mcp.server.auth.middleware.auth_context import auth_context_var from mcp.server.auth.middleware.bearer_auth import ( AuthenticatedUser, authorization_context, ) from mcp.server.auth.provider import AccessToken as SDKAccessToken from mcp.server.auth.provider import principal_components from mcp.server.context import ServerRequestContext from mcp.server.request_state import authenticated_principal from mcp_types import InputRequiredResult from starlette.testclient import TestClient from fastmcp import Context, FastMCP from fastmcp.server.auth import ( AccessToken, MultiAuth, RemoteAuthProvider, TokenVerifier, ) from fastmcp.server.auth.oauth_proxy import OAuthProxy from fastmcp.server.auth.providers.introspection import IntrospectionTokenVerifier from fastmcp.server.auth.providers.jwt import JWTVerifier, RSAKeyPair from fastmcp.server.dependencies import get_access_token from fastmcp.server.sessions import UserSession, current_principal from fastmcp.utilities.tests import asgi_server from tests.server.http.test_session_owner_enforcement import ( INITIALIZE_REQUEST, MCP_HEADERS, TOOLS_LIST_REQUEST, ) from tests.server.test_mrtr_guards import _ask, _elicit class StoredVerifier(TokenVerifier): def __init__(self, bearer: str) -> None: super().__init__() self.bearer = bearer self.result = AccessToken( token=bearer, client_id="shared", subject="reader", scopes=[], claims={"iss": "same", "sub": "reader"}, original_client_id="supplied", ) async def verify_token(self, token: str) -> AccessToken | None: return self.result if token == self.bearer else None async def test_named_sources_preserve_raw_fields_and_source_token() -> None: first, second = StoredVerifier("one"), StoredVerifier("two") auth = MultiAuth(verifiers={"company": first, "partner": second}) a, b = await auth.verify_token("one"), await auth.verify_token("two") assert a is not None and b is not None assert a is not first.result assert first.result.client_id == "shared" assert first.result.original_client_id == "supplied" assert json.loads(a.client_id) == ["company", "shared", "same", "reader"] assert json.loads(b.client_id) == ["partner", "shared", "same", "reader"] assert a.original_client_id == b.original_client_id == "shared" assert principal_components(a) != principal_components(b) assert a.claims == first.result.claims assert a.subject == first.result.subject assert a.token == first.result.token async def test_source_order_and_reconstruction_keep_identities() -> None: first, second = StoredVerifier("one"), StoredVerifier("two") left = MultiAuth(verifiers={"company": first, "partner": second}) right = MultiAuth(verifiers={"partner": second, "company": first}) a, b = await left.verify_token("one"), await right.verify_token("one") assert a is not None and b is not None assert principal_components(a) == principal_components(b) async def test_nested_sources_keep_the_inner_namespace() -> None: inner = MultiAuth(verifiers={"company": StoredVerifier("one")}) outer = MultiAuth(server=inner, server_source_id="interactive") result = await outer.verify_token("one") assert result is not None outer_name, inner_id, issuer, subject = json.loads(result.client_id) assert subject == "reader" assert issuer == "same" assert outer_name == "interactive" assert json.loads(inner_id) == ["company", "shared", "same", "reader"] assert result.original_client_id == "shared" @pytest.mark.parametrize("name", ["", " ", " company"]) def test_names_are_nonempty_and_unambiguous(name: str) -> None: with pytest.raises(ValueError, match="source IDs"): MultiAuth(verifiers={name: StoredVerifier("one")}) def test_opaque_sources_are_stable_and_duplicates_require_names() -> None: first = MultiAuth(verifiers=StoredVerifier("one")) restarted = MultiAuth(verifiers=StoredVerifier("new-credentials")) assert first._source_ids == restarted._source_ids with pytest.raises(ValueError, match="unique"): MultiAuth(verifiers=[StoredVerifier("one"), StoredVerifier("two")]) def test_configured_source_ids_are_stable_and_duplicates_require_names() -> None: first = JWTVerifier( jwks_uri="https://one.example/keys", issuer="https://one.example" ) second = JWTVerifier( jwks_uri="https://two.example/keys", issuer="https://two.example" ) a, b = MultiAuth(verifiers=[first, second]), MultiAuth(verifiers=[second, first]) assert a._source_ids == list(reversed(b._source_ids)) assert MultiAuth(verifiers=first)._source_ids == a._source_ids[:1] with pytest.raises(ValueError, match="unique"): MultiAuth(verifiers=[first, first]) async def test_ownership_consumers_agree_and_snapshot_preserves_raw_token() -> None: auth = MultiAuth( verifiers={"company": StoredVerifier("one"), "partner": StoredVerifier("two")} ) tokens = [await auth.verify_token(t) for t in ["one", "two"]] contexts = [] identities = [] task_scopes = [] for token in tokens: assert token is not None marker = auth_context_var.set(AuthenticatedUser(token)) try: assert get_access_token() is token identities.append(current_principal()) assert current_principal() == authenticated_principal( cast(ServerRequestContext[Any, Any], SimpleNamespace()) ) task_scopes.append(get_task_scope()) contexts.append(authorization_context(AuthenticatedUser(token))) restored = AccessToken.model_validate_json(token.model_dump_json()) assert principal_components(restored) == principal_components(token) snapshot = TaskContextSnapshot(access_token_json=token.model_dump_json()) _apply_snapshot_to_context(snapshot) assert current_principal() == identities[-1] restored_access_token = get_access_token() assert restored_access_token is not None assert restored_access_token.claims == token.claims assert restored_access_token.original_client_id == "shared" assert get_task_scope() == task_scopes[-1] finally: auth_context_var.reset(marker) assert identities[0] != identities[1] assert contexts[0] != contexts[1] assert task_scopes[0] != task_scopes[1] def test_plain_token_identity_retains_legacy_representation() -> None: token = AccessToken( token="one", client_id="client", subject="reader", scopes=[], claims={"iss": "issuer", "sub": "reader"}, ) assert principal_components(token) == ("client", "issuer", "reader") marker = auth_context_var.set(AuthenticatedUser(token)) try: assert get_task_scope() == "client|reader" assert authorization_context(AuthenticatedUser(token)) == { "client_id": "client", "issuer": "issuer", "subject": "reader", } finally: auth_context_var.reset(marker) async def test_introspection_endpoints_keep_session_and_transport_ownership_separate() -> ( None ): def respond(request: httpx2.Request) -> httpx2.Response: active = ( request.url.host == "one.example" and b"token=one" in request.content ) or (request.url.host == "two.example" and b"token=two" in request.content) return httpx2.Response( 200, json={"active": active, "client_id": "shared", "scope": ""} ) async with httpx2.AsyncClient(transport=httpx2.MockTransport(respond)) as http: sources = [ IntrospectionTokenVerifier( introspection_url="https://" + host + "/introspect", client_id="application", client_secret="synthetic", http_client=http, ) for host in ["one.example", "two.example"] ] auth = MultiAuth(verifiers=sources) server = FastMCP("Configured sources", auth=auth) @server.tool async def store(value: str, session: UserSession) -> str: await session.set("value", value) return value @server.tool async def read(session: UserSession) -> str | None: return await session.get("value") app = server.http_app(stateless_http=True, json_response=True) with TestClient(app, base_url="http://127.0.0.1") as client: def call(token: str, name: str, args: dict[str, Any]) -> Any: response = client.post( "/mcp", headers={**MCP_HEADERS, "authorization": "Bearer " + token}, json={ "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": name, "arguments": args}, }, ) assert response.status_code == 200 return response.json()["result"]["structuredContent"]["result"] assert call("one", "store", {"value": "saved"}) == "saved" assert call("two", "read", {}) is None assert call("two", "store", {"value": "other"}) == "other" assert call("one", "read", {}) == "saved" app = server.http_app(stateless_http=False, json_response=True) with TestClient(app, base_url="http://127.0.0.1") as client: initialized = client.post( "/mcp", headers={**MCP_HEADERS, "authorization": "Bearer one"}, json=INITIALIZE_REQUEST, ) assert initialized.status_code == 200 session_id = initialized.headers["mcp-session-id"] response = client.post( "/mcp", headers={ **MCP_HEADERS, "authorization": "Bearer two", "mcp-session-id": session_id, "mcp-protocol-version": "2024-11-05", }, json=TOOLS_LIST_REQUEST, ) assert response.status_code == 404 async def test_claims_only_subjects_keep_user_ownership_separate() -> None: verifier = StoredVerifier("one") verifier.result.subject = None verifier.result.claims["sub"] = "alice" auth = MultiAuth(verifiers={"company": verifier}) alice = await auth.verify_token("one") verifier.result.claims = {**verifier.result.claims, "sub": "bob"} bob = await auth.verify_token("one") assert alice is not None and bob is not None assert alice.subject is None and bob.subject is None scopes = [] for token in (alice, bob): marker = auth_context_var.set(AuthenticatedUser(token)) try: scopes.append(get_task_scope()) finally: auth_context_var.reset(marker) assert scopes[0] != scopes[1] @pytest.fixture async def configured_auth() -> AsyncGenerator[MultiAuth, None]: def respond(request: httpx2.Request) -> httpx2.Response: bearer = parse_qs(request.content.decode())["token"][0] source = (request.url.host or "").split(".")[0] return httpx2.Response( 200, json={ "active": bearer == source, "client_id": "shared", "scope": "", }, ) async with httpx2.AsyncClient(transport=httpx2.MockTransport(respond)) as http: yield MultiAuth( verifiers=[ IntrospectionTokenVerifier( introspection_url=f"https://{name}.example/introspect", client_id="application", client_secret="synthetic", http_client=http, ) for name in ("one", "two") ] ) async def test_configured_sources_keep_sse_message_ownership_separate( configured_auth: MultiAuth, ) -> None: server = FastMCP("SSE identities", auth=configured_auth) async with asgi_server(server, transport="sse") as running: async with running.http_client() as http: async with http.stream( "GET", running.url, headers={"accept": "text/event-stream", "authorization": "Bearer one"}, ) as response: assert response.status_code == 200 message_path = None async for line in response.aiter_lines(): if line.startswith("data: "): message_path = line.removeprefix("data: ") break assert message_path is not None message_url = str(httpx2.URL(running.url).join(message_path)) other = await http.post( message_url, headers={"authorization": "Bearer two"}, json=INITIALIZE_REQUEST, ) assert other.status_code == 404 same = await http.post( message_url, headers={"authorization": "Bearer one"}, json=INITIALIZE_REQUEST, ) assert same.status_code == 202 async def test_configured_sources_keep_saved_request_state_separate( configured_auth: MultiAuth, ) -> None: server = FastMCP("Continued requests", auth=configured_auth) @server.tool async def continue_work(ctx: Context) -> str | InputRequiredResult: if ctx.request_state is not None: return ctx.request_state return _ask( _elicit("confirmation", "Continue?", "answer"), "confirmation", "saved" ) app = server.http_app(stateless_http=True, json_response=True) with TestClient(app, base_url="http://127.0.0.1") as client: headers = { **MCP_HEADERS, "mcp-protocol-version": "2026-07-28", "mcp-method": "tools/call", "mcp-name": "continue_work", } request = { "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": { "name": "continue_work", "arguments": {}, "_meta": { "io.modelcontextprotocol/protocolVersion": "2026-07-28", "io.modelcontextprotocol/clientCapabilities": { "elicitation": {"form": {}} }, }, }, } first = client.post( "/mcp", headers={**headers, "authorization": "Bearer one"}, json=request ) assert first.status_code == 200, first.text saved = first.json()["result"]["requestState"] assert saved != "saved" request["params"]["requestState"] = saved other = client.post( "/mcp", headers={**headers, "authorization": "Bearer two"}, json=request ) assert "error" in other.json() same = client.post( "/mcp", headers={**headers, "authorization": "Bearer one"}, json=request ) assert same.json()["result"]["content"][0]["text"] == "saved" async def test_sdk_verifier_tokens_retain_original_client_id_after_serialization() -> ( None ): sdk_result = SDKAccessToken(token="one", client_id="client", scopes=[]) class SDKVerifier(TokenVerifier): async def verify_token(self, token: str) -> AccessToken | None: return cast(AccessToken, sdk_result) auth = MultiAuth(verifiers=SDKVerifier()) token = await auth.verify_token("one") assert token is not None assert sdk_result.client_id == "client" restored = AccessToken.model_validate_json(token.model_dump_json()) assert restored.original_client_id == "client" assert restored.client_id == token.client_id async def test_source_token_subclass_fields_survive_qualification() -> None: class DetailedToken(AccessToken): tenant: str verifier = StoredVerifier("one") verifier.result = DetailedToken( token="one", client_id="client", scopes=[], tenant="company" ) auth = MultiAuth(verifiers=verifier) token = await auth.verify_token("one") assert isinstance(token, DetailedToken) assert token.tenant == "company" assert verifier.result.client_id == "client" assert verifier.result.original_client_id is None assert token.original_client_id == "client" async def test_configured_issuer_list_keeps_task_ownership_distinct( rsa_key_pair: RSAKeyPair, ) -> None: issuers = ["https://company.example", "https://partner.example"] verifier = JWTVerifier(public_key=rsa_key_pair.public_key, issuer=issuers) auth = MultiAuth(verifiers=verifier) scopes = [] for issuer in issuers: bearer = rsa_key_pair.create_token( issuer=issuer, subject="reader", additional_claims={"client_id": "shared"} ) token = await auth.verify_token(bearer) assert token is not None assert token.original_client_id == "shared" assert token.claims["iss"] == issuer marker = auth_context_var.set(AuthenticatedUser(token)) try: scopes.append(get_task_scope()) finally: auth_context_var.reset(marker) assert scopes[0] != scopes[1] async def test_explicit_subjects_keep_task_ownership_distinct() -> None: verifier = StoredVerifier("one") verifier.result.claims = {"iss": "same"} auth = MultiAuth(verifiers={"company": verifier}) scopes = [] for subject in ("alice", "bob"): verifier.result.subject = subject token = await auth.verify_token("one") assert token is not None assert token.subject == subject assert "sub" not in token.claims marker = auth_context_var.set(AuthenticatedUser(token)) try: scopes.append(get_task_scope()) finally: auth_context_var.reset(marker) assert scopes[0] != scopes[1] @pytest.mark.parametrize("kind", ["introspection", "jwt"]) def test_remote_provider_identity_tracks_its_configured_verifier(kind: str) -> None: def configured(endpoint: str, source_id: str | None = None) -> MultiAuth: verifier = ( IntrospectionTokenVerifier( introspection_url=endpoint, client_id="application", client_secret="synthetic", ) if kind == "introspection" else JWTVerifier(jwks_uri=endpoint) ) server = RemoteAuthProvider( token_verifier=verifier, authorization_servers=[], base_url="https://service.example", ) return MultiAuth(server=server, server_source_id=source_id) first = configured("https://company.example/verify") rebuilt = configured("https://company.example/verify") other = configured("https://partner.example/verify") assert first._source_ids == rebuilt._source_ids assert first._source_ids != other._source_ids assert configured("https://company.example/verify", "company")._source_ids == [ "company" ] def test_proxy_source_identity_tracks_upstream_configuration() -> None: def configured( authority: str, secret: str = "synthetic", source_id: str | None = None ) -> MultiAuth: proxy = OAuthProxy( upstream_authorization_endpoint=f"https://{authority}/authorize", upstream_token_endpoint=f"https://{authority}/token", upstream_client_id="application", upstream_client_secret=secret, token_verifier=StoredVerifier("one"), base_url="https://service.example", ) return MultiAuth(server=proxy, server_source_id=source_id) first = configured("company.example") restarted = configured("company.example", "rotated") other = configured("partner.example") assert first._source_ids == restarted._source_ids assert first._source_ids != other._source_ids assert configured("company.example", source_id="company")._source_ids == ["company"]