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>

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)

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(

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,

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

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