From fc484be77add10f37eeb9e86c7c577bdda05d18b Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 16 Oct 2026 18:28:20 -0700 Subject: [PATCH] fix(proxy): keep the OIDC login refresh token as the ID-JAG subject when the SSO client is the ID-JAG client (#47004) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../login_assertion_exchange.py | 56 +++++++++++- .../outbound_credentials/resolver.py | 9 +- .../sso_assertion_store.py | 13 ++- .../test_login_assertion_exchange.py | 85 +++++++++++++++++++ .../test_sso_assertion_store.py | 21 +++++ 5 files changed, 174 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py index 224f956bcfd..fcd8ae22ce3 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/login_assertion_exchange.py @@ -5,7 +5,9 @@ login assertion on ID-JAG leg 1. It presents a refresh token that the IdP issued one RFC 8693 exchange (``requested_token_type`` refresh_token), which is how a SAML requesting app works at an IdP whose ID-JAG exchange does not take a ``saml2`` subject directly. That exchange has to happen at login: a SAML assertion is valid for minutes, the first MCP call that needs it can come hours later, and -the refresh token outlives the assertion by as long as the IdP decides. +the refresh token outlives the assertion by as long as the IdP decides. An OIDC login signed in through the very +IdP client the ID-JAG server uses already holds such a refresh token: it is bound to that client, so it is +presented as is and no further exchange is posted, which Okta refuses for an id_token subject anyway. The servers consulted are the two configuration authorities, the pod-local config declaration and the DB rows, like the retention gate, not the in-memory registry, which is a per-process snapshot that can lag a @@ -19,12 +21,17 @@ the ID-JAG call reports the missing refresh token as the precondition it is. from __future__ import annotations import asyncio +import os from collections.abc import Iterable, Mapping, Sequence from datetime import datetime, timezone from typing import Final, Protocol from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import ( + SSOClientConfig, + sso_client_config, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( SSOIdentityAssertion, with_login_refresh_token, @@ -123,12 +130,31 @@ async def exchange_login_assertion_for_refresh_token( token_endpoint: TokenExchangeEndpoint, now: datetime, deadline_seconds: float = _LOGIN_EXCHANGE_DEADLINE_SECONDS, + sso_client: SSOClientConfig | None = None, ) -> SSOIdentityAssertion: """The assertion as captured when no server wants a refresh token, otherwise the assertion plus the - refresh token the first target's IdP client exchanged it for. An OIDC login's own SSO refresh token - does not count, it is bound to the SSO client and the ID-JAG client cannot redeem it.""" + refresh token one target's IdP client can redeem. An OIDC login's own refresh token is that token + when the SSO client that received it is the target's client at the same token endpoint; otherwise it + is bound to a client the ID-JAG client is not, and the assertion is exchanged for one at the first + target's client.""" if not targets or assertion.exchanged_refresh_token is not None: return assertion + reusable: Final = _target_redeeming_the_login_refresh_token(assertion, targets, sso_client) + if reusable is not None and assertion.refresh_token is not None: + verbose_proxy_logger.info( + "ID-JAG: the refresh token issued at OIDC login belongs to client_id=%s at %s, the IdP client of a server " + "with id_jag_subject_source: refresh_token, so it is kept as that subject and no further exchange is posted", + reusable.client_id, + reusable.org_token_endpoint, + ) + return with_login_refresh_token( + assertion, + assertion.refresh_token.get_secret_value(), + None, + now, + endpoint=reusable.org_token_endpoint, + client_id=reusable.client_id, + ) config: Final = targets[0] if len(targets) > 1: verbose_proxy_logger.warning( @@ -185,6 +211,24 @@ async def exchange_login_assertion_for_refresh_token( ) +def _target_redeeming_the_login_refresh_token( + assertion: SSOIdentityAssertion, targets: tuple[IdJagConfig, ...], sso_client: SSOClientConfig | None +) -> IdJagConfig | None: + """The target whose IdP client is the SSO client the login's refresh token was issued to. A refresh + token is redeemable only by the client that received it at the endpoint that issued it, so both have + to match exactly before the login token is presented anywhere.""" + if assertion.refresh_token is None or sso_client is None: + return None + return next( + ( + config + for config in targets + if config.org_token_endpoint == sso_client.token_endpoint and config.client_id == sso_client.client_id + ), + None, + ) + + async def exchange_login_assertion_if_configured( assertion: SSOIdentityAssertion | None, ) -> SSOIdentityAssertion | None: @@ -209,7 +253,11 @@ async def exchange_login_assertion_if_configured( servers: Final = await login_exchange_servers(global_mcp_server_manager, db_rows) targets: Final = refresh_token_exchange_targets(to_server_spec(server) for server in servers) return await exchange_login_assertion_for_refresh_token( - assertion, targets, TokenEndpointClient(), datetime.now(timezone.utc) + assertion, + targets, + TokenEndpointClient(), + datetime.now(timezone.utc), + sso_client=sso_client_config(os.environ), ) except Exception as exc: # noqa: BLE001 # the login itself must not fail on an egress-side exchange verbose_proxy_logger.warning("ID-JAG: the login-time refresh token exchange failed unexpectedly: %s", exc) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index bd0813dd966..ea0cdee5c1c 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -429,10 +429,11 @@ def _login_refresh_token_subject( assertion: SSOIdentityAssertion, config: IdJagConfig, now: datetime ) -> Result[_IdJagSubjectToken, CredError]: """The subject for a server configured with ``id_jag_subject_source: refresh_token``: the refresh - token the ID-JAG client exchanged the login assertion for when the user signed in. Never the SSO - refresh token an OIDC login holds, that one is bound to the SSO client. The assertion's own expiry - does not apply here, a SAML assertion is spent minutes after login while the refresh token it - bought lives as long as the IdP said.""" + token the ID-JAG client holds from the user's sign-in, exchanged from the login assertion or, when + the SSO client is this server's IdP client, the one the OIDC login itself returned. The login-time + hook records which client and endpoint it is bound to, and a token bound elsewhere is never sent + here. The assertion's own expiry does not apply, a SAML assertion is spent minutes after login + while the refresh token it bought lives as long as the IdP said.""" if assertion.exchanged_refresh_token is None: return Error( CredError.of_precondition_required( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index 58494654082..19b3bc294d5 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -229,10 +229,19 @@ def carry_exchanged_refresh_token( ) -> SSOIdentityAssertion: """The renewed OIDC assertion still carrying the refresh token the ID-JAG client exchanged at login. Renewal only replaces the id_token and the SSO refresh token, the exchanged one is unrelated to - them and stays valid for as long as the IdP said.""" + them and stays valid for as long as the IdP said. When the login's own refresh token was kept as + that subject (the SSO client is the ID-JAG client), a rotation the renewal just received replaces + it too, since the IdP has retired the token the row held.""" + login_token_reused: Final = ( + previous.refresh_token is not None + and previous.exchanged_refresh_token is not None + and previous.exchanged_refresh_token.get_secret_value() == previous.refresh_token.get_secret_value() + ) return renewed.model_copy( update={ - "exchanged_refresh_token": previous.exchanged_refresh_token, + "exchanged_refresh_token": renewed.refresh_token + if login_token_reused + else previous.exchanged_refresh_token, "exchanged_refresh_token_expires_at": previous.exchanged_refresh_token_expires_at, "exchanged_refresh_token_endpoint": previous.exchanged_refresh_token_endpoint, "exchanged_refresh_token_client_id": previous.exchanged_refresh_token_client_id, diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py index bb228544e14..ab2b5b3d966 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_login_assertion_exchange.py @@ -13,6 +13,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.login_assertion refresh_token_exchange_targets, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import SSOClientConfig from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( SAML2_SUBJECT_TOKEN_TYPE, SSOIdentityAssertion, @@ -163,6 +164,90 @@ async def test_an_oidc_login_is_exchanged_at_the_id_jag_client_and_keeps_its_sso assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" +def _sso_client(client_id: str = "agent", endpoint: str = "https://idp.example.com/token") -> SSOClientConfig: + return SSOClientConfig( + token_endpoint=endpoint, client_id=client_id, client_secret=SecretStr("s"), auth_method="client_secret_basic" + ) + + +@pytest.mark.asyncio +async def test_an_oidc_login_through_the_id_jag_client_keeps_its_own_refresh_token_without_an_exchange(): + """Okta refuses requested_token_type=refresh_token for an id_token subject. When the SSO app is the + ID-JAG client itself, the refresh token the login returned is already that client's, so it is bound + to the target and nothing is posted to the IdP.""" + endpoint = _Endpoint( + Error( + CredError.of_upstream_unavailable( + "token exchange failed with status 400: 'requested_token_type' is invalid or not supported." + ) + ) + ) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + + kept = await exchange_login_assertion_for_refresh_token( + oidc, (_config(),), endpoint, _NOW, sso_client=_sso_client() + ) + + assert endpoint.calls == [] + assert kept.exchanged_refresh_token is not None and kept.exchanged_refresh_token.get_secret_value() == "sso-rt" + assert kept.exchanged_refresh_token_expires_at is None + assert kept.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert kept.exchanged_refresh_token_client_id == "agent" + assert kept.id_token == oidc.id_token and kept.refresh_token == oidc.refresh_token + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sso_client", + [ + pytest.param(_sso_client(client_id="sso-app"), id="other-client"), + pytest.param(_sso_client(endpoint="https://other-idp.example.com/token"), id="other-endpoint"), + pytest.param(None, id="unknown-client"), + ], +) +async def test_a_login_refresh_token_from_another_client_or_issuer_is_never_presented( + sso_client: SSOClientConfig | None, +): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + + exchanged = await exchange_login_assertion_for_refresh_token( + oidc, (_config(),), endpoint, _NOW, sso_client=sso_client + ) + + assert [(call[0], call[1], call[2]["subject_token"]) for call in endpoint.calls] == [ + ("https://idp.example.com/token", "agent", "id-token") + ] + assert exchanged.exchanged_refresh_token is not None + assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" + + +@pytest.mark.asyncio +async def test_the_login_refresh_token_binds_to_the_target_whose_client_issued_it_not_the_first(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + oidc = SSOIdentityAssertion(id_token=SecretStr("id-token"), refresh_token=SecretStr("sso-rt")) + targets = (_config(client_id="other-agent"), _config(client_id="agent")) + + kept = await exchange_login_assertion_for_refresh_token(oidc, targets, endpoint, _NOW, sso_client=_sso_client()) + + assert endpoint.calls == [] + assert kept.exchanged_refresh_token_client_id == "agent" + assert kept.exchanged_refresh_token is not None and kept.exchanged_refresh_token.get_secret_value() == "sso-rt" + + +@pytest.mark.asyncio +async def test_a_saml_login_through_the_id_jag_client_still_exchanges_its_assertion(): + endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-1", expires_in=3600))) + + exchanged = await exchange_login_assertion_for_refresh_token( + _saml(), (_config(),), endpoint, _NOW, sso_client=_sso_client() + ) + + assert endpoint.calls[0][2]["subject_token_type"] == "urn:ietf:params:oauth:token-type:saml2" + assert exchanged.exchanged_refresh_token is not None + assert exchanged.exchanged_refresh_token.get_secret_value() == "rt-1" + + @pytest.mark.asyncio async def test_a_login_already_exchanged_is_not_exchanged_again(): endpoint = _Endpoint(Ok(ExchangedToken(access_token="rt-2", expires_in=3600))) diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py index af62594e3b8..3c8e7a9b448 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py @@ -612,6 +612,27 @@ async def test_login_refresh_token_round_trips_beside_the_saml_assertion(): assert login_refresh_token_expired(fetched, now + timedelta(seconds=3600)) +def test_renewal_follows_a_rotated_login_refresh_token_that_is_the_id_jag_subject(): + """When the OIDC login's own refresh token is the ID-JAG subject (the SSO client is the ID-JAG + client), a renewal that rotated it leaves the old token revoked at the IdP, so the subject moves + to the token the renewal returned.""" + now = datetime.now(timezone.utc).replace(microsecond=0) + login = assertion_from_sso_login(_make_id_token(), "sso-rt") + assert login is not None + previous = with_login_refresh_token( + login, "sso-rt", None, now, endpoint="https://idp.example.com/token", client_id="agent" + ) + renewed = assertion_from_sso_login(_make_id_token(7200), "sso-rt-2") + assert renewed is not None + + carried = carry_exchanged_refresh_token(previous, renewed) + + assert carried.exchanged_refresh_token is not None + assert carried.exchanged_refresh_token.get_secret_value() == "sso-rt-2" + assert carried.exchanged_refresh_token_endpoint == "https://idp.example.com/token" + assert carried.exchanged_refresh_token_client_id == "agent" + + def test_renewal_carries_the_exchanged_refresh_token_forward(): """An OIDC renewal rebuilds the row from the new id_token. The refresh token the ID-JAG client exchanged at login is unrelated to that renewal and must survive it."""