From 4b2a89c3fdb6e9882fbf0a104f21b588104042de Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 31 Jul 2026 12:38:06 -0700 Subject: [PATCH] feat(sso): surface the ID-JAG capture gap on the SSO debug page /sso/debug/callback is where an operator lands when they are already trying to work out why ID-JAG is failing, so the reason belongs on it. The annotation appears only when the active SSO provider captures no identity assertion AND an oauth2_id_jag server is registered for that gap to break; a deployment without both renders the page it rendered before, byte for byte. Only the provider name and the remedy are rendered, never a configured value, and an unreachable MCP table costs the page its annotation rather than the page itself. The payload carries the one mutable-ok in this work. Conditionally including a member of a JSON document has to construct a mapping, and the rejected alternatives are recorded on the helper so the next reader does not rediscover them. Held out of the diagnosability PR deliberately: that PR is already reviewed and green, and this surface ships with the remaining config-load warning as one follow-up. --- litellm/proxy/management_endpoints/ui_sso.py | 25 ++++ .../proxy/management_endpoints/test_ui_sso.py | 126 ++++++++++++++++++ 2 files changed, 151 insertions(+) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 3155a0da73b..a4101140c1d 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1553,6 +1553,28 @@ async def warn_if_id_jag_assertion_uncaptured(assertion: SSOIdentityAssertion | ) +async def id_jag_capture_gap_to_surface() -> str | None: + """The capture gap worth putting in front of an operator: a real gap AND an ``oauth2_id_jag`` + server registered for it to break. A storage failure surfaces nothing rather than guessing. + + Returning ``None`` rather than an empty section is what lets the caller keep the debug page + byte-identical for every deployment without one, and it is why the caller's payload carries a + ``mutable-ok`` marker: an optional member of a JSON document has to be spelled as a mapping. + The immutable rewrites all fail against that. A dict comprehension is the same construction + rule, ``dict()`` is a flagged constructor, choosing between two whole payload literals doubles + the construction it was meant to avoid, and a frozen model dumped with ``exclude_none`` would + change how the three pre-existing keys serialize, which is the byte-identity this protects. + """ + gap = id_jag_assertion_capture_gap() + if gap is None: + return None + try: + return gap if await ema_assertion_retention_enabled() else None + except Exception as exc: # noqa: BLE001 # diagnostics must never break the page they annotate + verbose_proxy_logger.debug("Could not check for oauth2_id_jag MCP servers: %s", exc) + return None + + async def create_team_member_add_task(team_id, user_info): """Create a task for adding a member to a team.""" try: @@ -4520,10 +4542,13 @@ async def debug_sso_callback(request: Request): safe_raw_claims = {k: v for k, v in (received_response or {}).items() if k not in _OAUTH_TOKEN_FIELDS} safe_access_token_claims = {k: v for k, v in (access_token_payload or {}).items() if k not in _OAUTH_TOKEN_FIELDS} + gap = await id_jag_capture_gap_to_surface() + id_jag_section = {"id_jag_assertion_capture": gap} if gap is not None else {} # mutable-ok: optional JSON member sso_payload = { "parsed_by_proxy": filtered_result, "raw_claims": safe_raw_claims, "access_token_claims": safe_access_token_claims, + **id_jag_section, } # Replace the placeholder in the template with the actual data diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 4a1f5e1f33e..f4a6e0af83e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2,6 +2,7 @@ import asyncio import json import os import sys +from contextlib import ExitStack from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -7316,6 +7317,131 @@ async def test_debug_sso_callback_handles_missing_raw_response(): assert "user@example.com" in body +# ── The debug page is where an operator lands when ID-JAG is failing ────────── + +_GOOGLE_DEBUG_CLIENT_ID = "debug-google-client-id" +_GENERIC_DEBUG_CLIENT_ID = "debug-generic-client-id" + + +async def _render_debug_page(provider_env, id_jag_registered, gap_override=None): + """Drive /sso/debug/callback and return the raw response body.""" + from litellm.proxy.management_endpoints.ui_sso import GoogleSSOHandler, debug_sso_callback + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://proxy.example.com/" + mock_request.cookies = {} + mock_request.query_params = {} + + parsed = {"sub": "user_123", "email": "u@example.com"} + + async def fake_generic(**kwargs): + return parsed, {"sub": "user_123"}, {"scope": "openid"}, None + + async def fake_google(**kwargs): + return parsed + + stack = [ + patch.dict(os.environ, provider_env, clear=False), + patch("litellm.proxy.management_endpoints.ui_sso.get_generic_sso_response", side_effect=fake_generic), + patch.object(GoogleSSOHandler, "get_google_callback_response", side_effect=fake_google), + patch( + "litellm.proxy.management_endpoints.ui_sso.ema_assertion_retention_enabled", + AsyncMock(return_value=id_jag_registered), + ), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), + ] + if gap_override is not None: + stack.append( + patch( + "litellm.proxy.management_endpoints.ui_sso.id_jag_capture_gap_to_surface", + AsyncMock(return_value=gap_override["value"]), + ) + ) + + with ExitStack() as es: + for ctx in stack: + es.enter_context(ctx) + for var in ("MICROSOFT_CLIENT_ID", "GOOGLE_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"): + if var not in provider_env: + os.environ.pop(var, None) + response = await debug_sso_callback(mock_request) + + return response.body.decode() + + +@pytest.mark.asyncio +async def test_debug_page_surfaces_the_capture_gap_to_the_operator_reading_it(): + """This page is exactly where someone lands when ID-JAG is failing, so the reason has to be + on it; a diagnostic that cannot reach the surface the operator is staring at has a hole.""" + body = await _render_debug_page({"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=True) + + assert "id_jag_assertion_capture" in body + assert "google" in body + assert "GENERIC_CLIENT_ID" in body + + +@pytest.mark.asyncio +async def test_debug_page_gap_text_carries_no_configuration_values(): + """Only the provider name and the remedy belong in rendered HTML; the client id the operator + configured is not ours to echo back onto a page.""" + body = await _render_debug_page({"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=True) + + assert _GOOGLE_DEBUG_CLIENT_ID not in body + + +@pytest.mark.asyncio +async def test_debug_page_is_byte_identical_when_the_provider_captures(): + """A deployment with no gap must get the page it got before this change, to the byte. The + comparison is against the endpoint with the diagnostic forced inert, not against a guess.""" + with_feature = await _render_debug_page( + {"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, id_jag_registered=True + ) + pre_change = await _render_debug_page( + {"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, + id_jag_registered=True, + gap_override={"value": None}, + ) + + assert with_feature == pre_change + assert "id_jag" not in with_feature + + +@pytest.mark.asyncio +async def test_debug_page_is_byte_identical_when_no_id_jag_server_is_registered(): + """Most deployments run Google SSO and no id_jag server at all; their debug page must not + grow an ID-JAG section about a feature they do not use.""" + with_feature = await _render_debug_page( + {"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=False + ) + pre_change = await _render_debug_page( + {"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, + id_jag_registered=False, + gap_override={"value": None}, + ) + + assert with_feature == pre_change + assert "id_jag" not in with_feature + + +@pytest.mark.asyncio +async def test_debug_page_survives_a_store_outage(): + """The page's job is to render claims; an unreachable MCP table must cost it the annotation, + not the page.""" + from litellm.proxy.management_endpoints.ui_sso import id_jag_capture_gap_to_surface + + with ( + patch.dict(os.environ, {"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, clear=False), + patch( + "litellm.proxy.management_endpoints.ui_sso.ema_assertion_retention_enabled", + AsyncMock(side_effect=Exception("db down")), + ), + ): + assert await id_jag_capture_gap_to_surface() is None + + async def _render_legacy_login_page(env_overrides, general_settings): from litellm.proxy.management_endpoints.ui_sso import google_login