From eac5dc10f338e7faf751dd198d10e0bd7a2c024b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 22:11:36 -0700 Subject: [PATCH 01/23] fix(guardrails): apply PUT /guardrails/{id} to the serving worker immediately and reject invalid configs with 422 --- .../proxy/guardrails/guardrail_endpoints.py | 20 +++- .../proxy/guardrails/guardrail_registry.py | 37 ++---- .../guardrails/test_guardrail_endpoints.py | 53 +++++++-- .../guardrails/test_guardrail_registry.py | 108 ++++++++++++++---- 4 files changed, 158 insertions(+), 60 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 2b04828f0f2..3d2ed641a30 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -514,12 +514,26 @@ async def update_guardrail( guardrail_name: Final = result.get("guardrail_name", "Unknown") try: - IN_MEMORY_GUARDRAIL_HANDLER.update_in_memory_guardrail( - guardrail_id=guardrail_id, guardrail=cast(Guardrail, result) - ) + IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=cast(Guardrail, result)) verbose_proxy_logger.info( "Immediate sync: Successfully updated guardrail '%s' (ID: %s)", guardrail_name, guardrail_id ) + except (ValueError, TypeError) as update_error: + # The new config is invalid (a raising guardrail __init__): + # reinitialize_guardrail already restored the previous live instance, but + # update_guardrail_in_db above already persisted the rejected config to + # the DB. Roll that back too, so the DB and the live guardrail never + # disagree about what's actually enforcing, and surface the rejection to + # the caller instead of a misleading 200. + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=existing_guardrail, + prisma_client=prisma_client, + ) + raise HTTPException( + status_code=422, + detail=f"Invalid guardrail configuration, update rejected: {update_error}", + ) from update_error except Exception as update_error: verbose_proxy_logger.warning( "Immediate sync: Failed to update '%s' (ID: %s) in memory: %s", diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index dc13c09dd38..60873a1eeb1 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -6,7 +6,7 @@ import os from collections.abc import Callable, Iterator, Mapping from datetime import datetime, timezone from itertools import chain, count -from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, cast +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol from pydantic import ValidationError @@ -615,26 +615,6 @@ class InMemoryGuardrailHandler: return _guardrail_callback - def update_in_memory_guardrail( - self, - guardrail_id: str, - guardrail: Guardrail, - source: Literal["db", "config"] = "db", - ) -> None: - """ - Update a guardrail in memory - - - updates the guardrail in memory - - updates the guardrail params in litellm.callback_manager - """ - self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail - self._sources[guardrail_id] = source - - custom_guardrail_callback: Final = self.guardrail_id_to_custom_guardrail.get(guardrail_id) - if custom_guardrail_callback: - updated_litellm_params: Final = cast(LitellmParams, guardrail.get("litellm_params", {})) - custom_guardrail_callback.update_in_memory_litellm_params(litellm_params=updated_litellm_params) - def delete_in_memory_guardrail(self, guardrail_id: str) -> None: """ Delete a guardrail in memory and remove from litellm callbacks. @@ -789,11 +769,12 @@ class InMemoryGuardrailHandler: Removes old callback from litellm.callbacks and creates fresh instance. If the new config fails to initialize (e.g. an invalid on_flagged - combination), the previous instance is restored rather than left - deleted: initialize_guardrail's own ValueError/TypeError propagate - uncaught, so a caller reaching this point after already deleting the - old instance would otherwise leave the guardrail providing no - protection at all, not merely "still enforcing the old config." + combination or an invalid regex), the previous instance is restored + rather than left deleted, and the failure is re-raised as ValueError so + every init failure reaches callers as one exception type: a caller + reaching this point after already deleting the old instance would + otherwise leave the guardrail providing no protection at all, not + merely "still enforcing the old config." """ guardrail_id: Final = guardrail.get("guardrail_id") if not guardrail_id: @@ -812,7 +793,7 @@ class InMemoryGuardrailHandler: # that was enforcing must never fail open because an update was bad. try: return self.initialize_guardrail(guardrail=guardrail, config_file_path=config_file_path, source=source) - except Exception: + except Exception as init_error: if previous_guardrail is not None: verbose_proxy_logger.exception( "Reinitializing guardrail %s with updated params failed; restoring the previous configuration", @@ -824,7 +805,7 @@ class InMemoryGuardrailHandler: ) except Exception: # noqa: BLE001 # the original failure must propagate even if the restore breaks verbose_proxy_logger.exception("Restoring previous guardrail %s also failed", guardrail_id) - raise + raise ValueError(f"Guardrail initialization failed: {init_error}") from init_error def sync_guardrail_from_db(self, guardrail: Guardrail, config_file_path: str | None = None) -> Guardrail | None: """ diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 9b2117b7647..320e51203f6 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -104,7 +104,7 @@ def mock_in_memory_handler(mocker): mock_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL mock_handler.get_source.return_value = "config" mock_handler.initialize_guardrail = mocker.Mock() - mock_handler.update_in_memory_guardrail = mocker.Mock() + mock_handler.sync_guardrail_from_db = mocker.Mock() mock_handler.delete_in_memory_guardrail = mocker.Mock() mock_handler.reconcile_db_guardrails = mocker.Mock(return_value=[]) return mock_handler @@ -1045,13 +1045,15 @@ async def test_create_guardrail_endpoint( "scenario,expected_result,expected_exception", [ ("success_with_sync", "test-db-guardrail", None), - ("success_sync_fails", "test-db-guardrail", None), + ("success_sync_fails_unexpected_error", "test-db-guardrail", None), + ("sync_fails_invalid_config", None, HTTPException), ("database_failure", None, HTTPException), ("no_prisma_client", None, HTTPException), ], ids=[ "success_with_immediate_sync", - "success_but_sync_fails", + "success_but_sync_fails_with_unexpected_error", + "sync_rejects_invalid_config", "database_error", "missing_prisma_client", ], @@ -1071,6 +1073,7 @@ async def test_update_guardrail_endpoint( mock_logger = None if scenario == "success_with_sync": mock_prisma_client = mocker.Mock() + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", @@ -1081,10 +1084,13 @@ async def test_update_guardrail_endpoint( mock_in_memory_handler, ) - elif scenario == "success_sync_fails": + elif scenario == "success_sync_fails_unexpected_error": + # A non-ValueError/TypeError failure is not a config-rejection signal, + # so it keeps the pre-existing swallow-and-warn behavior rather than + # rolling back the DB write. mock_prisma_client = mocker.Mock() - mock_in_memory_handler.update_in_memory_guardrail.side_effect = Exception( - "Sync failed" + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( + side_effect=Exception("Sync failed") ) mock_logger = mocker.patch( "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" @@ -1100,6 +1106,25 @@ async def test_update_guardrail_endpoint( mock_in_memory_handler, ) + elif scenario == "sync_fails_invalid_config": + # Regression for the PUT half of the fix: a TypeError from the sync (the + # deleted update_in_memory_guardrail raised exactly this on every PUT) + # must roll back the DB write and surface a 422, not persist the + # rejected config with a 200. + mock_prisma_client = mocker.Mock() + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( + side_effect=TypeError("vars() argument must have __dict__ attribute") + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: reused pattern + mocker.patch( # test-quality-ok: reused pattern + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( # test-quality-ok: reused pattern + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + elif scenario == "database_failure": mock_prisma_client = mocker.Mock() mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception( @@ -1128,6 +1153,16 @@ async def test_update_guardrail_endpoint( assert "Database error" in str(exc_info.value.detail) elif scenario == "no_prisma_client": assert "Prisma client not initialized" in str(exc_info.value.detail) + elif scenario == "sync_fails_invalid_config": + assert exc_info.value.status_code == 422 + assert "update rejected" in str(exc_info.value.detail) + # Rolled back: update_guardrail_in_db is called once for the + # rejected write and once more to restore the previous config. + assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2 + assert ( + mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"] + == MOCK_DB_GUARDRAIL + ) else: result = await update_guardrail( @@ -1143,11 +1178,11 @@ async def test_update_guardrail_endpoint( prisma_client=mocker.ANY, ) - mock_in_memory_handler.update_in_memory_guardrail.assert_called_once_with( - guardrail_id="test-guardrail-id", guardrail=mocker.ANY + mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with( + guardrail=mocker.ANY ) - if scenario == "success_sync_fails": + if scenario == "success_sync_fails_unexpected_error": assert mock_logger is not None mock_logger.warning.assert_called_once() assert "Failed to update" in str(mock_logger.warning.call_args) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 2c0735970d3..beaffa73100 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -154,29 +154,51 @@ def test_duplicate_config_guardrail_names_get_distinct_stable_ids(): registry_module.guardrail_initializer_registry.pop("dup_name_test", None) -def test_update_in_memory_guardrail(): +def test_sync_guardrail_from_db_applies_db_dict_params_to_live_instance(): + """ + Regression for PUT /guardrails/{id}: the DB row arrives with litellm_params as + a plain jsonb dict, and the deleted update_in_memory_guardrail cast it to + LitellmParams without constructing one, so vars() raised and the running proxy + kept enforcing the stale config forever. The PUT endpoint now routes through + sync_guardrail_from_db, which must rebuild the live instance from the dict: + new blocked words compiled in, old ones gone, and the event hook re-derived + from mode (the base-class setattr path wrote self.mode while dispatch reads + self.event_hook, so only a full re-init applies a mode change). + """ + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + handler = InMemoryGuardrailHandler() - handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail( - guardrail_name="test-guardrail", - default_on=False, - event_hook=GuardrailEventHooks.pre_call, - ) + gid = "66666666-6666-6666-6666-666666666666" - handler.update_in_memory_guardrail( - "123", - Guardrail( - guardrail_name="test-guardrail", - litellm_params=LitellmParams(guardrail="test-guardrail", mode="pre_call", default_on=True), - ), - ) - - assert ( - handler.guardrail_id_to_custom_guardrail["123"].should_run_guardrail( - data={}, event_type=GuardrailEventHooks.pre_call + def db_guardrail(word: str, mode: str) -> Guardrail: + return Guardrail( + guardrail_id=gid, + guardrail_name="cf-put-sync", + litellm_params={ + "guardrail": "litellm_content_filter", + "mode": mode, + "default_on": True, + "blocked_words": [{"keyword": word, "action": "BLOCK"}], + }, ) - is True - ) - assert handler.guardrail_id_to_custom_guardrail["123"].event_hook is GuardrailEventHooks.pre_call + + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler.sync_guardrail_from_db(db_guardrail("foobarblock", "pre_call")) + handler.sync_guardrail_from_db(db_guardrail("quxnewblock", "during_call")) + + instance = handler.guardrail_id_to_custom_guardrail[gid] + assert isinstance(instance, ContentFilterGuardrail) + assert instance._check_blocked_words("hello QUXNEWBLOCK") is not None + assert instance._check_blocked_words("hello FOOBARBLOCK") is None + assert instance.event_hook == GuardrailEventHooks.during_call + assert instance.should_run_guardrail(data={}, event_type=GuardrailEventHooks.during_call) is True + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot def _make_guardrail(guardrail_id: str, name: str = "g") -> Guardrail: @@ -774,3 +796,49 @@ def test_reinitialize_guardrail_restores_previous_on_failure(): assert restored.guardrail_name == "restore-me" finally: registry_module.guardrail_initializer_registry.pop("restore_test", None) + + +def test_reinitialize_guardrail_raises_value_error_for_non_value_error_init_failures(): + """Regression for the LIT-6479 fix's 422 path: a constructor failure that is not + already a ValueError/TypeError (re.error from an invalid regex has neither in its + MRO) must still surface as ValueError, so the PUT/PATCH endpoints' rollback+422 + catch is exhaustive instead of warn-and-200 persisting a broken config.""" + import re + + from litellm.proxy.guardrails import guardrail_registry as registry_module + + def _initializer(litellm_params, guardrail): + if litellm_params.api_key == "bad-regex": + re.compile("([") + return CustomGuardrail( + guardrail_name=guardrail["guardrail_name"], + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + + registry_module.guardrail_initializer_registry["regex_test"] = _initializer + try: + handler = InMemoryGuardrailHandler() + created = handler.initialize_guardrail( + guardrail={ + "guardrail_name": "regex-me", + "litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "ok"}, + }, + ) + guardrail_id = created["guardrail_id"] + + with pytest.raises(ValueError, match="Guardrail initialization failed") as excinfo: + handler.reinitialize_guardrail( + guardrail={ + "guardrail_id": guardrail_id, + "guardrail_name": "regex-me", + "litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "bad-regex"}, + }, + ) + + assert isinstance(excinfo.value.__cause__, re.error) + assert guardrail_id in handler.IN_MEMORY_GUARDRAILS + restored = handler.guardrail_id_to_custom_guardrail[guardrail_id] + assert restored is not None and restored.guardrail_name == "regex-me" + finally: + registry_module.guardrail_initializer_registry.pop("regex_test", None) From 8c646a916c8860a04442645973974f71ee9c1c0b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:19:17 -0700 Subject: [PATCH 02/23] feat(cli): pre-fill the SSO verification code in the browser when the proxy allows it --- litellm/proxy/client/cli/commands/auth.py | 22 ++++-- .../proxy/client/cli/test_auth_commands.py | 70 +++++++++++++++++++ type-discipline-budget.json | 6 +- 3 files changed, 90 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 2fad9f933c1..a96d3fb9c85 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -100,9 +100,10 @@ class CliPollData(TypedDict, total=False): class CliSsoStartData(TypedDict): - login_id: str - poll_secret: str - user_code: str + login_id: ReadOnly[str] + poll_secret: ReadOnly[str] + user_code: ReadOnly[str] + verification_uri_complete: ReadOnly[NotRequired[str]] class CliAuthResult(TypedDict): @@ -860,11 +861,22 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool) -> None: poll_secret: Final = cli_sso_flow["poll_secret"] user_code: Final = cli_sso_flow["user_code"] - sso_url = f"{base_url}/sso/key/generate?" + urlencode({"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": key_id}) + browser_prefills_code: Final = isinstance(cli_sso_flow.get("verification_uri_complete"), str) + sso_url: Final = f"{base_url}/sso/key/generate?" + urlencode( + ( + ("source", LITELLM_CLI_SOURCE_IDENTIFIER), + ("key", key_id), + *((("user_code", user_code),) if browser_prefills_code else ()), + ) + ) click.echo(f"Opening browser to: {sso_url}") click.echo("Please complete the SSO authentication in your browser...") - click.echo(f"Verification code: {user_code}") + click.echo( + f"Verification code: {user_code} (pre-filled in the browser, check it matches)" + if browser_prefills_code + else f"Verification code: {user_code}" + ) click.echo(f"Session ID: {key_id}") # Open browser diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 1d0a99b8e0a..821323e722c 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -59,6 +59,7 @@ def _mock_cli_sso_start_response( login_id: str = "cli-session-uuid-456", poll_secret: str = "poll-secret", user_code: str = "ABCD-EFGH", + **extra_fields: object, ) -> Mock: mock_response = Mock() mock_response.status_code = 200 @@ -66,6 +67,7 @@ def _mock_cli_sso_start_response( "login_id": login_id, "poll_secret": poll_secret, "user_code": user_code, + **extra_fields, } mock_response.raise_for_status = Mock() return mock_response @@ -333,7 +335,9 @@ class TestLoginCommand: call_args = mock_browser.call_args[0][0] assert "https://test.example.com/sso/key/generate" in call_args assert "cli-test-uuid-123" in call_args + assert "user_code" not in call_args assert "Verification code: ABCD-EFGH" in result.output + assert "pre-filled in the browser" not in result.output mock_post.assert_called_once() mock_get.assert_called() assert mock_get.call_args.kwargs["headers"] == {"x-litellm-cli-poll-secret": "poll-secret"} @@ -347,6 +351,72 @@ class TestLoginCommand: # Verify commands were shown mock_show_commands.assert_called_once() + def test_login_prefills_the_code_in_the_browser_when_the_proxy_advertises_it( + self, isolated_home, secret_vault_factory + ) -> None: + vault = secret_vault_factory() + poll_response = Mock() + poll_response.status_code = 200 + poll_response.json.return_value = { + "status": "ready", + "key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt", + "user_id": "test-user-123", + "team_id": "team-1", + "teams": ["team-1"], + } + start_response = _mock_cli_sso_start_response( + login_id="cli-test-uuid-123", + verification_uri_complete=( + "https://internal-hostname.example.com/sso/key/generate" + "?source=litellm-cli&key=cli-test-uuid-123&user_code=ABCD-EFGH" + ), + ) + + with ( + patch("webbrowser.open") as mock_browser, + patch("requests.post", return_value=start_response), + patch("requests.get", return_value=poll_response), + ): + result = self.runner.invoke(login, obj={"base_url": "https://test.example.com", "secret_vault": vault}) + + assert result.exit_code == 0, result.output + assert json.loads(vault.blob)["key"] == "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt" + assert json.loads((isolated_home / ".litellm" / "token.json").read_text())["user_id"] == "test-user-123" + opened_url = mock_browser.call_args[0][0] + assert opened_url.startswith("https://test.example.com/sso/key/generate?") + assert "internal-hostname" not in opened_url + assert "key=cli-test-uuid-123" in opened_url + assert "user_code=ABCD-EFGH" in opened_url + assert "Verification code: ABCD-EFGH (pre-filled in the browser, check it matches)" in result.output + + def test_login_keeps_the_code_out_of_the_url_when_the_proxy_sends_a_non_url_verification_uri( + self, secret_vault_factory + ) -> None: + poll_response = Mock() + poll_response.status_code = 200 + poll_response.json.return_value = { + "status": "ready", + "key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt", + "user_id": "test-user-123", + "team_id": "team-1", + "teams": ["team-1"], + } + for advertised in (None, True): + start_response = _mock_cli_sso_start_response(verification_uri_complete=advertised) + + with ( + patch("webbrowser.open") as mock_browser, + patch("requests.post", return_value=start_response), + patch("requests.get", return_value=poll_response), + ): + result = self.runner.invoke( + login, obj={"base_url": "https://test.example.com", "secret_vault": secret_vault_factory()} + ) + + assert result.exit_code == 0, result.output + assert "user_code" not in mock_browser.call_args[0][0] + assert "pre-filled in the browser" not in result.output + def test_login_timeout(self): """Test login timeout scenario""" mock_context = Mock() diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 52cb9628252..2836b22c3d0 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22364 }, "LIT002": { - "limit": 26777 + "limit": 26776 }, "LIT003": { "limit": 269 @@ -27,12 +27,12 @@ "limit": 0 }, "LIT010": { - "limit": 16507 + "limit": 16506 }, "LIT011": { "limit": 5535 }, "LIT012": { - "limit": 4495 + "limit": 4492 } } From 8e3566d2f7564438c8c33ab66ce831aee416059e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:38:43 -0700 Subject: [PATCH 03/23] fix(guardrails): forward mode and streaming params to crowdstrike_aidr handler (#39317) * fix(guardrails): forward mode and streaming params to crowdstrike_aidr handler Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(guardrails): drop stream_chunk_builder patch from crowdstrike cadence test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(guardrails): type test params and cover unsupported crowdstrike mode rejection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crowdstrike_aidr/__init__.py | 13 +- .../crowdstrike_aidr/crowdstrike_aidr.py | 40 ++++- .../guardrail_hooks/crowdstrike_aidr.py | 13 +- .../guardrail_hooks/test_crowdstrike_aidr.py | 142 +++++++++++++++++- 4 files changed, 197 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py index 5e75b7d4d94..c88e6e97a96 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py @@ -1,8 +1,8 @@ from typing import TYPE_CHECKING, Final -from litellm.types.guardrails import GuardrailEventHooks, SupportedGuardrailIntegrations +from litellm.types.guardrails import SupportedGuardrailIntegrations -from .crowdstrike_aidr import CrowdStrikeAIDRHandler +from .crowdstrike_aidr import CrowdStrikeAIDRHandler, streaming_params_from_litellm_params if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams @@ -15,17 +15,16 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" if not guardrail_name: raise ValueError("CrowdStrike AIDR guardrail name is required") + streaming_params: Final = streaming_params_from_litellm_params(litellm_params) _crowdstrike_aidr_callback: Final = CrowdStrikeAIDRHandler( guardrail_name=guardrail_name, api_base=litellm_params.api_base, api_key=litellm_params.api_key, - # Exclude during_call to prevent duplicate input events - event_hook=[ - GuardrailEventHooks.pre_call.value, - GuardrailEventHooks.post_call.value, - ], + event_hook=litellm_params.mode, default_on=litellm_params.default_on, fail_on_error=litellm_params.fail_on_error, + streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only, + streaming_sampling_rate=streaming_params.streaming_sampling_rate, ) litellm.logging_callback_manager.add_litellm_callback(_crowdstrike_aidr_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index c8284fac440..f7f500b1adc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -24,8 +24,11 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionToolParam +from litellm.types.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import ( + CrowdStrikeAIDRGuardrailConfigModelOptionalParams, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -153,6 +156,21 @@ def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Mapping[str, Any] | return merged if present else None +def streaming_params_from_litellm_params( + litellm_params: LitellmParams, +) -> CrowdStrikeAIDRGuardrailConfigModelOptionalParams: + extras: Final[Mapping[str, object]] = litellm_params.model_extra or {} + nested: Final = litellm_params.optional_params + optional_params: Final[Mapping[str, object]] = {} if nested is None else nested.model_dump() + return CrowdStrikeAIDRGuardrailConfigModelOptionalParams.model_validate( + { + name: value + for name in CrowdStrikeAIDRGuardrailConfigModelOptionalParams.model_fields + if (value := optional_params.get(name, extras.get(name))) is not None + } + ) + + def _messages_since_last_assistant( messages: Sequence[AllMessageValues], ) -> _FilteredMessages: @@ -241,6 +259,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): api_key: str | None = None, api_base: str | None = None, fail_on_error: bool | None = True, + streaming_end_of_stream_only: bool | None = None, + streaming_sampling_rate: int | None = None, **kwargs, ) -> None: """ @@ -250,10 +270,19 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): guardrail_name (str): The name of the guardrail instance. api_key (str | None): The CrowdStrike AIDR API key. Reads from CS_AIDR_TOKEN env var if None. api_base (str | None): The CrowdStrike AIDR API base URL. Reads from CS_AIDR_BASE_URL env var if None. + streaming_end_of_stream_only (bool | None): Scan streamed output once at end of stream instead of + every streaming_sampling_rate chunks. Defaults to False. + streaming_sampling_rate (int | None): Scan the accumulated streamed output every Nth chunk. Defaults to 5. **kwargs: Additional arguments passed to the CustomGuardrail base class. """ self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.fail_on_error = True if fail_on_error is None else fail_on_error + self._set_streaming_params( + CrowdStrikeAIDRGuardrailConfigModelOptionalParams( + streaming_end_of_stream_only=streaming_end_of_stream_only, + streaming_sampling_rate=streaming_sampling_rate, + ) + ) self.api_key = api_key or os.environ.get("CS_AIDR_TOKEN") if not self.api_key: @@ -274,6 +303,15 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): "Initialized CrowdStrike AIDR Guardrail: name=%s, api_base=%s", guardrail_name, self.api_base ) + def _set_streaming_params(self, streaming_params: CrowdStrikeAIDRGuardrailConfigModelOptionalParams) -> None: + self.streaming_end_of_stream_only: bool = streaming_params.streaming_end_of_stream_only or False + self.streaming_sampling_rate: int = streaming_params.streaming_sampling_rate or 5 + + @override + def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: + super().update_in_memory_litellm_params(litellm_params) + self._set_streaming_params(streaming_params_from_litellm_params(litellm_params)) + async def _call_crowdstrike_aidr_guard( self, payload: dict[str, Any], hook_name: str ) -> _GuardChatCompletionsResult: diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py index f47c38af3e3..6beca030a3a 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py @@ -4,7 +4,18 @@ from .base import GuardrailConfigModel class CrowdStrikeAIDRGuardrailConfigModelOptionalParams(BaseModel): - pass + streaming_end_of_stream_only: bool | None = Field( + default=None, + description="If False (default when unset), post_call scans the accumulated streamed response every " + "streaming_sampling_rate chunks and an in-flight block stops the stream. If True, the guard runs once " + "over the assembled response at end of stream, so flagged content may already have reached the client.", + ) + streaming_sampling_rate: int | None = Field( + default=None, + ge=1, + description="When streaming_end_of_stream_only is False, scan the accumulated streamed response every Nth " + "chunk. Defaults to 5 when unset.", + ) class CrowdStrikeAIDRGuardrailConfigModel(GuardrailConfigModel[CrowdStrikeAIDRGuardrailConfigModelOptionalParams]): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index ec7854b9a35..1b50ea53db2 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -3,7 +3,9 @@ from unittest.mock import patch import httpx import pytest from fastapi import HTTPException +from pydantic import ValidationError +import litellm from litellm.exceptions import Timeout from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import initialize_guardrail @@ -12,8 +14,8 @@ from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr CrowdStrikeAIDRHandler, ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 -from litellm.types.guardrails import Guardrail, LitellmParams -from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse +from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams +from litellm.types.utils import Delta, GenericGuardrailAPIInputs, ModelResponse, ModelResponseStream @pytest.fixture @@ -1578,3 +1580,139 @@ async def test_unparseable_transformed_response_fails_closed_under_fail_open() - assert exc_info.value.status_code == 500 assert "failing closed" in exc_info.value.detail["error"] + + +def _initialize_from_config(**litellm_params_kwargs: object) -> CrowdStrikeAIDRHandler: + litellm_params = LitellmParams( + guardrail="crowdstrike_aidr", + api_key="pts_crowdstrike_tokenid", + api_base="https://api.crowdstrike.com/aidr/aiguard", + default_on=True, + **litellm_params_kwargs, + ) + guardrail = Guardrail(guardrail_name="crowdstrike-aidr-guard", litellm_params=litellm_params) + return initialize_guardrail(litellm_params=litellm_params, guardrail=guardrail) + + +@pytest.mark.parametrize( + ("mode", "runs_pre_call", "runs_post_call"), + [("post_call", False, True), ("pre_call", True, False), (["pre_call", "post_call"], True, True)], +) +def test_initialize_guardrail_honors_configured_mode( + mode: str | list[str], runs_pre_call: bool, runs_post_call: bool +) -> None: + handler = _initialize_from_config(mode=mode) + + assert handler.should_run_guardrail({}, GuardrailEventHooks.pre_call) is runs_pre_call + assert handler.should_run_guardrail({}, GuardrailEventHooks.post_call) is runs_post_call + + +def test_initialize_guardrail_rejects_unsupported_mode_instead_of_running_other_hooks() -> None: + with pytest.raises(ValueError, match="during_call is not in the supported event hooks"): + _initialize_from_config(mode="during_call") + + +def test_initialize_guardrail_defaults_streaming_params() -> None: + handler = _initialize_from_config(mode="post_call") + + assert handler.streaming_end_of_stream_only is False + assert handler.streaming_sampling_rate == 5 + + +@pytest.mark.parametrize( + "configured", + [ + {"streaming_end_of_stream_only": True, "streaming_sampling_rate": 50}, + {"optional_params": {"streaming_end_of_stream_only": True, "streaming_sampling_rate": 50}}, + ], +) +def test_initialize_guardrail_forwards_streaming_params(configured: dict[str, object]) -> None: + handler = _initialize_from_config(mode="post_call", **configured) + + assert handler.streaming_end_of_stream_only is True + assert handler.streaming_sampling_rate == 50 + + +def test_initialize_guardrail_rejects_non_positive_sampling_rate() -> None: + with pytest.raises(ValidationError): + _initialize_from_config(mode="post_call", streaming_sampling_rate=0) + + +def test_update_in_memory_litellm_params_reapplies_streaming_params() -> None: + handler = _initialize_from_config(mode="post_call") + + handler.update_in_memory_litellm_params( + LitellmParams( + guardrail="crowdstrike_aidr", + mode="post_call", + streaming_end_of_stream_only=True, + streaming_sampling_rate=7, + ) + ) + + assert handler.streaming_end_of_stream_only is True + assert handler.streaming_sampling_rate == 7 + + +def _stream_chunk(content: str, finish_reason: str | None) -> ModelResponseStream: + return ModelResponseStream( + model="gpt-4", + choices=[ + litellm.StreamingChoices( + index=0, delta=Delta(role="assistant", content=content), finish_reason=finish_reason + ) + ], + ) + + +async def _guard_calls_for_stream(handler: CrowdStrikeAIDRHandler, chunk_texts: list[str]) -> int: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails + + async def stream(): + for i, content in enumerate(chunk_texts): + yield _stream_chunk(content, "stop" if i == len(chunk_texts) - 1 else None) + + calls = 0 + + def _allow(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + return httpx.Response( + status_code=200, json={"result": {"blocked": False, "transformed": False}}, request=request + ) + + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": handler, + "metadata": {"guardrails": ["crowdstrike-aidr-guard"]}, + } + async with httpx.AsyncClient(transport=httpx.MockTransport(_allow)) as client: + await handler.async_handler.close() + handler.async_handler.client = client + async for _ in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/chat/completions"), + response=stream(), + request_data=request_data, + ): + pass + return calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("configured", "expected_calls"), + [ + ({}, 3), + ({"streaming_sampling_rate": 2}, 6), + ({"streaming_end_of_stream_only": True}, 1), + ({"streaming_end_of_stream_only": True, "streaming_sampling_rate": 2}, 1), + ], +) +async def test_streaming_params_from_config_control_output_scan_cadence( + configured: dict[str, object], expected_calls: int +) -> None: + """10 chunks: default samples at 5 and 10 plus the final pass, rate 2 samples 5 times plus final, end-of-stream scans once.""" + handler = _initialize_from_config(mode="post_call", **configured) + + assert await _guard_calls_for_stream(handler, list("ABCDEFGHIJ")) == expected_calls From 8cbaba88635871b44096488d6264490eac700b9f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:39:36 -0700 Subject: [PATCH 04/23] fix(mcp): gate the connect-time OBO pre-flight on the key's allowed servers (#39447) The pre-flight token exchange resolved its target from the requested path, so a key not entitled to that server still drove an outbound exchange at the IdP and populated the credential cache before the later access check denied it. Resolve the target through the same allowed-server set the rest of the route uses and skip the exchange when the requested server is not in it. Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/server.py | 7 ++ .../mcp_server/test_mcp_server.py | 73 +++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 989b08b929a..af3ff6714af 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3855,6 +3855,13 @@ if MCP_AVAILABLE: and server.auth_type == MCPAuth.oauth2_token_exchange and oauth2_headers and len(mcp_servers or []) == 1 + and server.server_id + in frozenset( + allowed.server_id + for allowed in await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth, mcp_servers=mcp_servers, client_ip=client_ip + ) + ) ): await global_mcp_server_manager.preflight_token_exchange( server=server, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 82f74cda835..3f6d8f8837c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -8198,6 +8198,79 @@ class TestPreemptive401ModeAware: await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False) +def _make_obo_server(alias: str) -> MCPServer: + return MCPServer( + server_id=f"id-{alias}", + name=alias, + alias=alias, + server_name=alias, + url=f"https://{alias}.test/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.test/token", + client_id="cid", + client_secret="csecret", + mcp_info={"server_name": alias}, + ) + + +class TestOboPreflightScopedToAllowedServers: + """The connect-time OBO exchange is an outbound IdP call whose result is cached, so it must + only run for a server the caller's key resolves to through the allowed set, not for any + server the requested path happens to name.""" + + SUBJECT_HEADERS = {"Authorization": "Bearer upstream-subject-token"} + + async def _run(self, requested: MCPServer, allowed: list[MCPServer], user_api_key_auth: UserAPIKeyAuth | None): + from litellm.proxy._experimental.mcp_server import server as server_module + + allowed_lookup = AsyncMock(return_value=allowed) + preflight = AsyncMock() + with ( + patch.object( # test-quality-ok: route handler reads the module-level manager, no injection seam + server_module.global_mcp_server_manager, "get_mcp_server_by_name", return_value=requested + ), + patch.object( # test-quality-ok: the exchanger is the observable; a real one would call an IdP + server_module.global_mcp_server_manager, "preflight_token_exchange", preflight + ), + patch.object( # test-quality-ok: allowed-set resolution needs the DB; the test controls its answer + server_module, "_get_allowed_mcp_servers", allowed_lookup + ), + ): + await server_module._raise_preemptive_401_for_unauthenticated_servers( + scope={"type": "http", "method": "POST", "path": f"/mcp/{requested.alias}", "headers": []}, + mcp_servers=[requested.alias], + oauth2_headers=self.SUBJECT_HEADERS, + mcp_server_auth_headers=None, + user_api_key_auth=user_api_key_auth, + client_ip="10.0.0.7", + ) + return allowed_lookup, preflight + + @pytest.mark.asyncio + async def test_unentitled_key_never_reaches_the_exchanger(self): + requested = _make_obo_server("obo_tools") + key = UserAPIKeyAuth(api_key="sk-plain-only") + + allowed_lookup, preflight = await self._run( + requested, allowed=[_make_obo_server("plain_tools")], user_api_key_auth=key + ) + + preflight.assert_not_awaited() + allowed_lookup.assert_awaited_once_with( + user_api_key_auth=key, mcp_servers=[requested.alias], client_ip="10.0.0.7" + ) + + @pytest.mark.asyncio + async def test_entitled_key_still_exchanges_at_connect(self): + requested = _make_obo_server("obo_tools") + key = UserAPIKeyAuth(api_key="sk-obo") + + _, preflight = await self._run(requested, allowed=[requested], user_api_key_auth=key) + + preflight.assert_awaited_once_with(server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key) + + @pytest.mark.asyncio async def test_post_mcp_call_guardrails_return_the_rewritten_result(): """The result a post_mcp_call guardrail rewrote must be what the caller sends back.""" From c19d49d9192f26a7e15c9571d1c3a4e7ee45092a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:40:13 -0700 Subject: [PATCH 05/23] fix(responses): keep provider response headers in streaming logging callbacks (#38131) * fix(responses): keep provider response headers in streaming logging callbacks The responses streaming iterator captures the provider's HTTP response headers into its own _hidden_params, but never puts them on the completed response, and the model_validate(model_dump()) copy made for logging drops pydantic private attributes. Success callbacks and StandardLoggingPayload.hidden_params.additional_headers therefore saw an empty dict for streaming /v1/responses, so Azure's apim-request-id was unreadable from the callback payload. Restore the headers on the nested response of the logging copy, preferring any the provider transform already set (the fake_stream path) and falling back to the ones the iterator captured from the stream. Skipped when the copy fell back to the original event, so a serialization failure never leaves logging-only state on the caller's object. * fix: satisfy LIT002 mutable-collection gate in header restore --------- Co-authored-by: Yucheng Zhu --- litellm/responses/streaming_iterator.py | 40 +++++ .../responses/test_streaming_iterator.py | 161 +++++++++++++++++- 2 files changed, 200 insertions(+), 1 deletion(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7871c85220c..f271655f5e3 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -17,6 +17,7 @@ from typing_extensions import TypeIs import litellm from litellm.constants import ( + EMPTY_MAPPING, LITELLM_MAX_STREAMING_DURATION_SECONDS, STREAM_SSE_DONE_STRING, ) @@ -273,6 +274,9 @@ class BaseResponsesAPIStreamingIterator: self._hidden_params["additional_headers"] = process_response_headers( self.response.headers or {} ) # GUARANTEE OPENAI HEADERS IN RESPONSE + self._raw_response_headers: Mapping[str, str] = MappingProxyType( + dict(self.response.headers or {}) # mutable-ok: immediately frozen by MappingProxyType + ) def _check_max_streaming_duration(self) -> None: """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" @@ -446,6 +450,7 @@ class BaseResponsesAPIStreamingIterator: except Exception: # Fallback to original if serialization fails pass + self._restore_provider_response_headers(logging_response) end_time: Final = datetime.now() if is_async: @@ -480,6 +485,41 @@ class BaseResponsesAPIStreamingIterator: ) self._run_post_success_hooks(end_time=end_time) + def _restore_provider_response_headers(self, logging_response: object) -> None: + """Re-apply the provider's response headers to the copy handed to logging callbacks. + + ``model_validate(model_dump())`` above drops pydantic private attributes, so the + ``_hidden_params`` the provider transform set on the nested response are lost. Returns early + when that copy fell back to the original event, so logging-only state never lands on the + object the caller is iterating. + """ + if logging_response is self.completed_response: + return + target: Final[object] = getattr(logging_response, "response", None) + existing_hidden: Final[object] = getattr(target, "_hidden_params", None) + if not isinstance(existing_hidden, Mapping): + return + existing: Final[Mapping[str, object]] = existing_hidden + source_hidden: Final[object] = getattr( + getattr(self.completed_response, "response", None), "_hidden_params", None + ) + source: Final[Mapping[str, object]] = source_hidden if isinstance(source_hidden, Mapping) else EMPTY_MAPPING + processed: Final[object] = source.get("additional_headers") or self._hidden_params.get("additional_headers") + raw: Final[object] = source.get("headers") or self._raw_response_headers + headers: Final[Mapping[str, object]] = processed if isinstance(processed, Mapping) else EMPTY_MAPPING + raw_headers: Final[Mapping[str, object]] = raw if isinstance(raw, Mapping) else EMPTY_MAPPING + # rebuild by value and let existing keys win: sharing the source dicts would alias what the proxy + # splats into the client's HTTP headers, and copying non-header keys would carry response_cost + setattr( # noqa: B010 # target is typed object here, so a plain attribute store does not type check + target, + "_hidden_params", + { # mutable-ok: the cost calculator writes optional_params into _hidden_params + "additional_headers": {**headers}, # mutable-ok: fresh copy, logging callbacks may mutate it + "headers": {**raw_headers}, # mutable-ok: fresh copy, logging callbacks may mutate it + **existing, + }, + ) + def _handle_logging_completed_response(self): """Base implementation - should be overridden by subclasses""" diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 9edcaaef034..c226c0b4d09 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -6,7 +6,7 @@ completion_start_time = end_time.""" import json from datetime import datetime from typing import Optional -from unittest.mock import Mock +from unittest.mock import Mock, patch import httpx import pytest @@ -378,3 +378,162 @@ def test_stamp_responses_usage_cost_survives_calculator_failure(): _stamp_responses_usage_cost(response, logging_obj) assert getattr(response.usage, "cost", None) is None + + +def _capture_dispatch(logged: list): + """Record the object handed to the success handlers. + + ``Mock(spec=LiteLLMLoggingObj).dispatch_success_handlers`` is an AsyncMock whose side effect + only runs when the coroutine is awaited, so capture with a plain function instead. + """ + + async def _noop() -> None: + return None + + def _dispatch(result, **kwargs): + logged.append(result) + return _noop() + + return _dispatch + + +def _headers_config(*, transform_hidden_params: Optional[dict] = None) -> Mock: + """Config whose completed event carries a real ResponsesAPIResponse, so the logging copy + performs a genuine model_dump/model_validate round trip.""" + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def _transform(model, parsed_chunk, logging_obj): + evt_type = parsed_chunk.get("type") + if evt_type != "response.completed": + stub = Mock() + stub.type = evt_type + return stub + response = ResponsesAPIResponse( + id="resp_headers", + created_at=1, + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + if transform_hidden_params is not None: + response._hidden_params.update(transform_hidden_params) + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=response, + ) + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + + +def _make_header_iterator( + *, + headers: dict, + config: Mock, + logging_obj: LiteLLMLoggingObj, +) -> ResponsesAPIStreamingIterator: + async def aiter_bytes(): + yield _sse_event({"type": "response.completed"}) + + mock_response = Mock() + mock_response.headers = headers + mock_response.aiter_bytes = aiter_bytes + + return ResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-4o-mini", + responses_api_provider_config=config, + logging_obj=logging_obj, + litellm_metadata={}, + custom_llm_provider="azure", + ) + + +@pytest.mark.asyncio +async def test_streaming_logging_response_carries_provider_response_headers(): + """LIT-6055: the provider headers the iterator captured must reach the logged response, so + custom loggers can read Azure's apim-request-id from the callback payload.""" + logging_obj = _logging_obj_stub() + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator( + headers={"apim-request-id": "azure-correlation-1", "x-ms-region": "East US 2"}, + config=_headers_config(), + logging_obj=logging_obj, + ) + async for _ in iterator: + pass + + assert len(logged) == 1 + hidden_params = logged[0].response._hidden_params + assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "azure-correlation-1" + assert hidden_params["additional_headers"]["llm_provider-x-ms-region"] == "East US 2" + assert hidden_params["headers"]["apim-request-id"] == "azure-correlation-1" + # the proxy builds the client's response headers from the iterator's own dict, so the logged + # response must hold copies rather than alias it + assert hidden_params["additional_headers"] is not iterator._hidden_params["additional_headers"] + assert hidden_params["headers"] is not iterator._raw_response_headers + + +@pytest.mark.asyncio +async def test_streaming_logging_copy_preserves_transform_hidden_params(): + """LIT-6055: model_validate(model_dump()) drops pydantic private attributes, so headers a + provider transform already set on the response (fake_stream) must be re-applied.""" + logging_obj = _logging_obj_stub() + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator( + headers={}, + config=_headers_config( + transform_hidden_params={ + "additional_headers": {"llm_provider-apim-request-id": "from-transform"}, + "headers": {"apim-request-id": "from-transform"}, + "response_cost": 0.5, + } + ), + logging_obj=logging_obj, + ) + async for _ in iterator: + pass + + assert len(logged) == 1 + hidden_params = logged[0].response._hidden_params + assert hidden_params["additional_headers"]["llm_provider-apim-request-id"] == "from-transform" + assert hidden_params["headers"]["apim-request-id"] == "from-transform" + assert iterator.completed_response is not logged[0] + # only the header keys travel: response_cost would short-circuit the cost calculator + assert "response_cost" not in hidden_params + + +@pytest.mark.asyncio +async def test_streaming_logging_copy_fallback_leaves_caller_event_untouched(): + """LIT-6055: when the logging copy falls back to the original event, the header restore must + not stamp logging-only state onto the object the caller is iterating.""" + logging_obj = _logging_obj_stub() + logged: list[object] = [] + logging_obj.dispatch_success_handlers = _capture_dispatch(logged) + logging_obj._on_deferred_stream_complete = None + + iterator = _make_header_iterator( + headers={"apim-request-id": "azure-correlation-1"}, + config=_headers_config(), + logging_obj=logging_obj, + ) + async for _ in iterator: + pass + + assert len(logged) == 1 + iterator._completed_response_logged = False + logged.clear() + with patch.object(type(iterator.completed_response), "model_dump", side_effect=ValueError("cannot serialize")): + iterator._log_completed_response(is_async=True) + + assert logged == [iterator.completed_response] + assert iterator.completed_response.response._hidden_params == {} From 44a6c659fb831bc3da929f020f2cd4f3e508e384 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 2 Sep 2026 17:45:58 -0700 Subject: [PATCH 06/23] fix(mcp): fence an outbound-token write against an overlapping invalidation (#35398) get_or_compute single-flights concurrent misses under a per-key lock, but invalidate() deletes outside it. A mint already awaiting the IdP when an invalidation ran wrote its result into the slot afterwards, so a bearer minted before the invalidation was served after it for its full TTL; the upstream-401 retry then re-presented the bearer the server had just rejected. The cache now keeps a per-key generation beside the single-flight lock. invalidate() bumps it and get_or_compute snapshots it inside the lock right before minting, skipping the cache write when it no longer matches. The token is still returned to the caller it was minted for, so only the caching is fenced. Generation and lock share one object that a running mint binds to a local, which pins the weak map's entry for exactly as long as a write can still land; a generation held weakly on its own would be collected across the await and the re-check would silently pass. --- .../outbound_credentials/token_endpoint.py | 67 ++++++++--- .../test_token_endpoint.py | 105 ++++++++++++++++++ 2 files changed, 158 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py index 84f714db449..6f1e5baa109 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py @@ -121,6 +121,23 @@ class TokenEndpointClient: return Ok(ExchangedToken(access_token=parsed.access_token, expires_in=parsed.expires_in)) +class _KeyGuard: + """The per-key single-flight lock plus the invalidation generation that lock protects. + + Both live on one object so their lifetimes cannot diverge. `get_or_compute` binds the guard to + a local for its whole critical section, which keeps the weak map's entry alive for as long as + that compute could still write; an `invalidate` overlapping the compute therefore reaches the + very same object and its bump is guaranteed to be observed. Conversely a guard nobody holds is + collectible precisely because no write is outstanding for it to fence. + """ + + __slots__ = ("__weakref__", "generation", "lock") + + def __init__(self) -> None: + self.lock = asyncio.Lock() + self.generation = 0 + + class ExchangedTokenCache: """Memoizes the final token string per key, single-flighting concurrent misses on one lock.""" @@ -129,7 +146,7 @@ class ExchangedTokenCache: max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, ) - self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() + self._guards: weakref.WeakValueDictionary[str, _KeyGuard] = weakref.WeakValueDictionary() async def get_or_compute( self, @@ -144,28 +161,50 @@ class ExchangedTokenCache: guaranteeing the token it gets back was minted for the *current* inputs: a stored entry whose fingerprint differs reads as a miss and is re-minted over. That keeps eviction addressable without the key having to encode the credential material it protects. + + An `invalidate` landing while `compute` is in flight wins over that compute's write. The + token is still returned to the caller it was minted for, but it is not stored, so the next + resolution re-mints rather than serving a bearer that predates the invalidation for the + rest of its TTL. """ cached = self._get(cache_key, fingerprint) if cached is not None: return Ok(cached) - async with self._lock(cache_key): + guard = self._guard(cache_key) + async with guard.lock: cached = self._get(cache_key, fingerprint) if cached is not None: return Ok(cached) + generation = guard.generation match await compute(): case Ok(token): - self._cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped - cache_key, - (fingerprint, token.access_token), - ttl=_cache_ttl_seconds(token.expires_in), - ) + if guard.generation == generation: + self._store(cache_key, fingerprint, token) return Ok(token.access_token) case Error(err): return Error(err) def invalidate(self, cache_key: str) -> None: - """Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401).""" + """Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401). + + Bumping the guard's generation is what makes the eviction stick against a compute already + awaiting the token endpoint: that compute snapshotted the old generation and so skips its + write. No guard means no compute is in flight, since an in-flight one pins its own. + + Stays synchronous: callers invalidate from plain `def`s. + """ self._cache.delete_cache(cache_key) # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + guard = self._guards.get(cache_key) + if guard is None: + return + guard.generation += 1 + + def _store(self, cache_key: str, fingerprint: str, token: ExchangedToken) -> None: + self._cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + cache_key, + (fingerprint, token.access_token), + ttl=_cache_ttl_seconds(token.expires_in), + ) def _get(self, cache_key: str, fingerprint: str) -> str | None: """The stored token, or None when absent or minted for different inputs. @@ -180,12 +219,12 @@ class ExchangedTokenCache: return None return token if stored_fingerprint == fingerprint else None - def _lock(self, cache_key: str) -> asyncio.Lock: - lock = self._locks.get(cache_key) - if lock is None: - lock = asyncio.Lock() - self._locks[cache_key] = lock - return lock + def _guard(self, cache_key: str) -> _KeyGuard: + guard = self._guards.get(cache_key) + if guard is None: + guard = _KeyGuard() + self._guards[cache_key] = guard + return guard def _cache_ttl_seconds(expires_in: int | None) -> int: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py index 5f277db2f72..db3a1a386a3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py @@ -7,6 +7,7 @@ cache's hit/single-flight behavior. Each assertion fails under a real mutation o """ import asyncio +import gc import json from unittest.mock import AsyncMock, MagicMock, patch @@ -359,6 +360,110 @@ async def test_cache_invalidate_only_evicts_the_named_key(): assert calls == 2 +@pytest.mark.asyncio +async def test_cache_invalidate_mid_compute_is_not_overwritten_by_that_compute(): + """A bearer minted before an invalidation must never be served after it. + + The compute is suspended at the token endpoint when the invalidation lands, so its write is + the one that would resurrect the evicted bearer for the rest of its TTL. The caller it was + minted for still gets it; the *cache* is what the invalidation is about. + """ + cache = ExchangedTokenCache() + mint_started, release_mint = asyncio.Event(), asyncio.Event() + + async def slow_mint(): + mint_started.set() + await release_mint.wait() + return _ok_token("bearer-minted-before-invalidation") + + async def re_mint(): + return _ok_token("bearer-minted-after-invalidation") + + in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp")) + await mint_started.wait() + + assert not in_flight.done() + cache.invalidate("slot") + release_mint.set() + + raced = await in_flight + assert isinstance(raced, Ok) and raced.ok == "bearer-minted-before-invalidation" + + after = await cache.get_or_compute("slot", re_mint, fingerprint="fp") + assert isinstance(after, Ok) and after.ok == "bearer-minted-after-invalidation" + + +@pytest.mark.asyncio +async def test_cache_invalidate_mid_compute_survives_garbage_collection(): + """The record of an invalidation must outlive a collection cycle taken mid-compute. + + Per-key state is held weakly so idle keys do not accumulate. If the state a compute checks + before writing were collectible while that compute is suspended, the check would read as + "nothing was invalidated" and the stale write would land; the running compute has to pin it. + """ + cache = ExchangedTokenCache() + mint_started, release_mint = asyncio.Event(), asyncio.Event() + + async def slow_mint(): + mint_started.set() + await release_mint.wait() + return _ok_token("bearer-minted-before-invalidation") + + async def re_mint(): + return _ok_token("bearer-minted-after-invalidation") + + in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp")) + await mint_started.wait() + + assert not in_flight.done() + cache.invalidate("slot") + gc.collect() + release_mint.set() + await in_flight + + after = await cache.get_or_compute("slot", re_mint, fingerprint="fp") + assert isinstance(after, Ok) and after.ok == "bearer-minted-after-invalidation" + + +@pytest.mark.asyncio +async def test_cache_stores_a_compute_that_started_after_the_invalidation(): + """Only the mint that predates the invalidation loses its write. + + A caller queued behind the single-flight lock computes after the eviction, so its token is + fresh and must be cached; otherwise the fix would trade one stale bearer for re-minting on + every subsequent resolution. + """ + cache = ExchangedTokenCache() + mint_started, release_mint = asyncio.Event(), asyncio.Event() + + async def slow_mint(): + mint_started.set() + await release_mint.wait() + return _ok_token("bearer-minted-before-invalidation") + + async def re_mint(): + return _ok_token("bearer-minted-after-invalidation") + + async def must_not_run(): + pytest.fail("the mint that followed the invalidation should have been cached") + + in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp")) + await mint_started.wait() + queued = asyncio.create_task(cache.get_or_compute("slot", re_mint, fingerprint="fp")) + await asyncio.sleep(0) + + assert not queued.done() + cache.invalidate("slot") + release_mint.set() + + raced, fresh = await asyncio.gather(in_flight, queued) + assert isinstance(raced, Ok) and raced.ok == "bearer-minted-before-invalidation" + assert isinstance(fresh, Ok) and fresh.ok == "bearer-minted-after-invalidation" + + served = await cache.get_or_compute("slot", must_not_run, fingerprint="fp") + assert isinstance(served, Ok) and served.ok == "bearer-minted-after-invalidation" + + @pytest.mark.asyncio async def test_cache_does_not_store_a_failed_compute(): cache = ExchangedTokenCache() From ff1f21aea92574b3df1b9a00b2cdd70fa5e17dda Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 2 Sep 2026 17:54:17 -0700 Subject: [PATCH 07/23] fix(ui): paginate request logs by session groups server-side (#39257) * fix(ui): paginate request logs by session groups server-side The logs table server-paginated raw spend logs and then collapsed multi-call sessions client-side, so a page could render 3 rows while the footer claimed 25 and sessions straddled pages. Adds an opt-in group_by_session param to /spend/logs/ui that pages and counts one representative row per session (DISTINCT ON, newest non-MCP call), keeps the bounded count contract, enriches whole-session llm/agent composition counts, and deletes the client-side collapse pipeline. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QxT89fiygmzz2ALcjpu7Ve * feat(ui): add a 10 rows-per-page option and default request logs to it Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QxT89fiygmzz2ALcjpu7Ve * fix(ui): key session aggregates per api key in the logs enrichment Grouped pagination splits a reused session id into one row per api key, but the enrichment still aggregated by session_id alone, so both rows showed combined spend and counts. The aggregate query now groups by (session_id, api_key), the count folds into it (the separate group_by query is deleted), and each row reads its own key's totals. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QxT89fiygmzz2ALcjpu7Ve * fix(ui): treat an empty api_key as a real session group value The spend-log schema defaults api_key to an empty string; truthiness guards in the enrichment treated it as missing, so keyless multi-call sessions lost their count and spend. Only None means missing now. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QxT89fiygmzz2ALcjpu7Ve --------- Co-authored-by: Claude Fable 5 --- .../spend_management_endpoints.py | 118 ++++++++------ tests/e2e/ui/helpers/traffic.ts | 3 + .../e2e/ui/tests/logs/logsPagination.spec.ts | 150 ++++++++++++++++++ .../test_spend_management_endpoints.py | 145 +++++++++++++++-- .../test_spend_query_optimization.py | 106 +++++++++++++ .../src/components/networking.tsx | 1 + .../view_logs/RequestLogsPanel.test.tsx | 45 +++--- .../components/view_logs/RequestLogsPanel.tsx | 63 +------- .../components/view_logs/RequestLogsTable.tsx | 4 +- .../RequestLogsTableColumns.test.tsx | 14 +- .../view_logs/RequestLogsTableColumns.tsx | 6 +- .../src/components/view_logs/columns.tsx | 1 - .../src/components/view_logs/constants.ts | 3 + .../components/view_logs/log_filter_logic.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 15 files changed, 512 insertions(+), 152 deletions(-) create mode 100644 tests/e2e/ui/tests/logs/logsPagination.spec.ts diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index c48750cea72..bb7dfafb297 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -55,6 +55,10 @@ router: Final = APIRouter() SPEND_LOGS_PAGINATION_COUNT_CAP: Final = 10000 +_SESSION_GROUP_KEY_SQL: Final = "COALESCE(NULLIF(session_id, ''), request_id), api_key" +_MCP_CALL_TYPES_SQL: Final = "('call_mcp_tool', 'list_mcp_tools')" +_AGENT_CALL_TYPE_SQL: Final = "'asend_message'" + _INTERNAL_HEALTH_CHECK_API_KEYS: Final = ( LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, hash_token(token=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME), @@ -144,21 +148,16 @@ class _DailyTagSpendRow(TypedDict): total_spend: float -class _SessionCountAggregate(TypedDict): - session_id: int - - -class _SessionCountRow(TypedDict): - session_id: str - _count: _SessionCountAggregate - - class _SessionSpendRow(TypedDict): session_id: str + api_key: ReadOnly[str] + session_total_count: ReadOnly[int] session_total_spend: float mcp_tool_call_count: int mcp_tool_call_spend: float session_cache_hit_count: ReadOnly[int] + session_llm_count: ReadOnly[int] + session_agent_count: ReadOnly[int] class _SpendSumAggregate(TypedDict, total=False): @@ -242,18 +241,6 @@ async def _count_spend_logs(prisma_client: PrismaClient, where: Mapping[str, obj return await _spend_logs_table(prisma_client).count(where=where) -async def _count_logs_per_session( - prisma_client: PrismaClient, session_ids: Sequence[str | None] -) -> Sequence[_SessionCountRow]: - """Count spend log rows per session for the given session ids.""" - rows: Final = await _spend_logs_table(prisma_client).group_by( - by=["session_id"], - where={"session_id": {"in": session_ids}}, - count={"session_id": True}, - ) - return cast(Sequence[_SessionCountRow], rows) # cast-ok: group_by(count=) shape is fixed by the by/count args - - async def _find_team_row(prisma_client: PrismaClient, team_id: str) -> _SupportsModelDump | None: """Read a single team row as a Prisma model instance.""" return await _team_table(prisma_client).find_unique(where={"team_id": team_id}) @@ -2290,6 +2277,10 @@ async def ui_view_spend_logs( default=False, description="Exclude LiteLLM internal health check requests from results", ), + group_by_session: bool = fastapi.Query( + default=False, + description="Paginate over sessions instead of raw logs: one representative row per session, total counts sessions", + ), ): """ View spend logs with pagination support. @@ -2644,12 +2635,16 @@ async def ui_view_spend_logs( else: _order_expr = order_column + joined_conditions: Final = " AND ".join(sql_conditions) + session_grouping: Final = group_by_session is True + count_group_clause: Final = f"GROUP BY {_SESSION_GROUP_KEY_SQL}" if session_grouping else "" count_query: Final = f""" SELECT COUNT(*) AS total_count FROM ( SELECT 1 FROM "LiteLLM_SpendLogs" - WHERE {" AND ".join(sql_conditions)} + WHERE {joined_conditions} + {count_group_clause} LIMIT ${p} ) AS bounded_matches """ @@ -2660,21 +2655,36 @@ async def ui_view_spend_logs( total_is_capped: Final = raw_total > SPEND_LOGS_PAGINATION_COUNT_CAP total_records: Final = SPEND_LOGS_PAGINATION_COUNT_CAP if total_is_capped else raw_total - sql_query: Final = f""" - SELECT - request_id, call_type, api_key, spend, total_tokens, + select_columns: Final = """request_id, call_type, api_key, spend, total_tokens, prompt_tokens, completion_tokens, "startTime", "endTime", "completionStartTime", model, model_id, model_group, custom_llm_provider, api_base, "user", metadata, cache_hit, cache_key, request_tags, team_id, organization_id, end_user, requester_ip_address, session_id, status, mcp_namespaced_tool_name, agent_id, - COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms + COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms""" + sql_query: Final = ( + f""" + SELECT * FROM ( + SELECT DISTINCT ON ({_SESSION_GROUP_KEY_SQL}) + {select_columns} + FROM "LiteLLM_SpendLogs" + WHERE {joined_conditions} + ORDER BY {_SESSION_GROUP_KEY_SQL}, call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC + ) AS session_representatives + ORDER BY {_order_expr} {_sql_dir}{_nulls_clause}, request_id + LIMIT ${p} OFFSET ${p + 1} + """ + if session_grouping + else f""" + SELECT + {select_columns} FROM "LiteLLM_SpendLogs" - WHERE {" AND ".join(sql_conditions)} + WHERE {joined_conditions} ORDER BY {_order_expr} {_sql_dir}{_nulls_clause} LIMIT ${p} OFFSET ${p + 1} """ + ) sql_params.extend([page_size, skip]) data: Final = await prisma_client.db.query_raw(sql_query, *sql_params) @@ -4075,11 +4085,12 @@ async def _build_ui_spend_logs_response( Build the paginated response for the UI spend-logs endpoint. When ``enrich_session_counts`` is ``True`` (the default for the v1/UI - endpoint), each row is enriched with ``session_total_count`` so the - frontend knows which sessions are expandable (multi-call sessions). - For every row that carries a ``session_id``, a single ``GROUP BY`` query - fetches the total number of logs in each referenced session. Rows without - a ``session_id`` default to ``1``. + endpoint), each row is enriched with ``session_total_count`` plus spend + and call-type aggregates so the frontend knows which sessions are + expandable (multi-call sessions). One ``GROUP BY (session_id, api_key)`` + query serves every referenced session, keyed per api key so two callers + reusing a session id never see each other's totals. Rows without a + ``session_id`` default to ``1``. When ``enrich_session_counts`` is ``False`` (v2 endpoint), rows are serialised without the extra query. @@ -4101,7 +4112,6 @@ async def _build_ui_spend_logs_response( A dict with ``data`` (enriched rows), ``total``, ``page``, ``page_size``, ``total_pages``, and ``total_is_capped``. """ - count_map: dict[str, int] = {} if enrich_session_counts: session_ids: Final[Sequence[str | None]] = list( { @@ -4110,15 +4120,8 @@ async def _build_ui_spend_logs_response( if (row.get("session_id") if isinstance(row, dict) else getattr(row, "session_id", None)) } ) - if session_ids: - # NOTE: This GROUP BY runs on every v1/UI page load. The IN clause - # is bounded by page_size (typically 25-50 distinct session IDs). - # If performance degrades at scale, consider short-lived caching or - # folding the count into the main query via a window function. - counts: Final = await _count_logs_per_session(prisma_client, session_ids) - count_map = {r["session_id"]: r["_count"]["session_id"] for r in counts if r.get("session_id")} - session_spend_map: dict[str, dict[str, int | float]] = {} + session_spend_map: dict[tuple[str, str], dict[str, int | float]] = {} if enrich_session_counts and session_ids: from prisma.errors import PrismaError @@ -4130,38 +4133,46 @@ async def _build_ui_spend_logs_response( { (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) for row in data - if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) + if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) is not None } ) rows: Final[Sequence[_SessionSpendRow]] = await _query_raw( prisma_client, - """ - SELECT session_id, + f""" + SELECT session_id, api_key, + COUNT(*)::int AS session_total_count, COALESCE(SUM(spend), 0)::double precision AS session_total_spend, COUNT(*) FILTER ( - WHERE call_type IN ('call_mcp_tool', 'list_mcp_tools') + WHERE call_type IN {_MCP_CALL_TYPES_SQL} )::int AS mcp_tool_call_count, COALESCE(SUM(spend) FILTER ( - WHERE call_type IN ('call_mcp_tool', 'list_mcp_tools') + WHERE call_type IN {_MCP_CALL_TYPES_SQL} ), 0)::double precision AS mcp_tool_call_spend, - COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count + COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count, + COUNT(*) FILTER ( + WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL} + )::int AS session_llm_count, + COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count FROM "LiteLLM_SpendLogs" WHERE session_id = ANY($1::text[]) AND api_key = ANY($2::text[]) - GROUP BY session_id + GROUP BY session_id, api_key """, session_ids, authorized_api_keys, ) session_spend_map = { - row["session_id"]: { + (row["session_id"], row["api_key"]): { + "session_total_count": int(row.get("session_total_count") or 0), "session_total_spend": float(row.get("session_total_spend") or 0.0), "mcp_tool_call_count": int(row.get("mcp_tool_call_count") or 0), "mcp_tool_call_spend": float(row.get("mcp_tool_call_spend") or 0.0), "session_cache_hit_count": int(row.get("session_cache_hit_count") or 0), + "session_llm_count": int(row.get("session_llm_count") or 0), + "session_agent_count": int(row.get("session_agent_count") or 0), } for row in rows - if row.get("session_id") + if row.get("session_id") and row.get("api_key") is not None } except PrismaError: verbose_proxy_logger.debug( @@ -4174,14 +4185,17 @@ async def _build_ui_spend_logs_response( for row in data: row_dict = dict(row) if isinstance(row, dict) else row.model_dump() sid = row_dict.get("session_id") - row_dict["session_total_count"] = count_map.get(sid, 1) if sid else 1 - session_stats = session_spend_map.get(sid) if sid else None + row_api_key = row_dict.get("api_key") + session_stats = session_spend_map.get((sid, row_api_key)) if sid and row_api_key is not None else None + row_dict["session_total_count"] = int(session_stats["session_total_count"]) if session_stats else 1 if session_stats: row_dict["session_total_spend"] = session_stats["session_total_spend"] if session_stats["mcp_tool_call_count"]: row_dict["mcp_tool_call_count"] = session_stats["mcp_tool_call_count"] row_dict["mcp_tool_call_spend"] = session_stats["mcp_tool_call_spend"] row_dict["session_cache_hit_count"] = session_stats["session_cache_hit_count"] + row_dict["session_llm_count"] = session_stats["session_llm_count"] + row_dict["session_agent_count"] = session_stats["session_agent_count"] enriched.append(row_dict) response_data: list = enriched else: diff --git a/tests/e2e/ui/helpers/traffic.ts b/tests/e2e/ui/helpers/traffic.ts index 25eb671fd0e..7f8417cdffb 100644 --- a/tests/e2e/ui/helpers/traffic.ts +++ b/tests/e2e/ui/helpers/traffic.ts @@ -21,6 +21,8 @@ interface ChatOptions { apiKey?: string; /** Sent as `user`, which lands in the spend log's end_user column. */ endUser?: string; + /** Sent as `litellm_trace_id`, which lands in the spend log's session_id column. */ + traceId?: string; } /** POST /v1/chat/completions and return the completion id (the Logs Request ID). */ @@ -34,6 +36,7 @@ export async function sendChatCompletion(request: APIRequestContext, opts: ChatO model: opts.model, messages: [{ role: "user", content: opts.prompt }], ...(opts.endUser ? { user: opts.endUser } : {}), + ...(opts.traceId ? { litellm_trace_id: opts.traceId } : {}), }, }); expect(res.ok(), `chat completion for ${opts.model} failed (${res.status()}): ${await res.text()}`).toBe(true); diff --git a/tests/e2e/ui/tests/logs/logsPagination.spec.ts b/tests/e2e/ui/tests/logs/logsPagination.spec.ts new file mode 100644 index 00000000000..416e1c3171f --- /dev/null +++ b/tests/e2e/ui/tests/logs/logsPagination.spec.ts @@ -0,0 +1,150 @@ +import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { Page } from "../../fixtures/pages"; +import { CHAT_MODEL_A, createVirtualKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic"; + +/** + * Session-grouped pagination (#38060): a page of N rows must render exactly N session rows, a + * session must never straddle pages, and two callers reusing one session id stay separate rows. + * All traffic is generated per run behind a unique key alias or session id, so concurrent specs + * cannot decide the outcome. + */ + +const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; + +/** Every tab stays mounted, so the DOM holds four tables at once; scope to the visible one. */ +const requestLogsRows = (page: PlaywrightPage): Locator => + page.locator("table").filter({ visible: true }).first().locator("tbody tr"); + +const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true }); + +async function openLogs(page: PlaywrightPage): Promise { + await navigateToPage(page, Page.Logs); + await dismissFeedbackPopup(page); + await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 }); +} + +async function openFilterDrawer(page: PlaywrightPage): Promise { + await visibleTestId(page, "datatable-filters-trigger").click(); + const drawer = page.getByRole("dialog", { name: "Filters" }); + await expect(drawer).toBeVisible({ timeout: 10_000 }); + return drawer; +} + +async function applyKeyAliasFilter(page: PlaywrightPage, drawer: Locator, alias: string): Promise { + await drawer.getByRole("combobox", { name: "Search a key alias" }).click(); + await page.keyboard.type(alias); + await page.getByRole("option", { name: alias, exact: true }).first().click(); + await drawer.getByRole("button", { name: "Apply Filters" }).click(); + await expect(drawer).not.toBeVisible({ timeout: 10_000 }); +} + +async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): Promise { + await visibleTestId(page, "pagination-page-size").click(); + await page.getByRole("option", { name: size, exact: true }).click(); +} + +test.describe("Logs page session-grouped pagination", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("a 25-row page renders exactly 25 session rows and no session straddles pages", async ({ page, request }) => { + const suffix = uniqueSuffix(); + const alias = `e2e-logs-pgn-${suffix}`; + const mine = await createVirtualKey(request, { key_alias: alias }); + + const soloIds: string[] = []; + for (let i = 0; i < 26; i++) { + soloIds.push( + await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-pgn-solo-${i}-${suffix}`, + apiKey: mine.key, + }), + ); + } + const sessionA = `sess-pgn-a-${suffix}`; + const sessionB = `sess-pgn-b-${suffix}`; + let lastSessionCallId = ""; + for (let i = 0; i < 7; i++) { + lastSessionCallId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-pgn-a-${i}-${suffix}`, + apiKey: mine.key, + traceId: sessionA, + }); + } + for (let i = 0; i < 3; i++) { + lastSessionCallId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-pgn-b-${i}-${suffix}`, + apiKey: mine.key, + traceId: sessionB, + }); + } + await waitForSpendLog(request, lastSessionCallId); + await waitForSpendLog(request, soloIds[soloIds.length - 1]); + + // 36 calls in 28 session groups: 26 solos plus sessions of 7 and 3. + await openLogs(page); + const drawer = await openFilterDrawer(page); + await applyKeyAliasFilter(page, drawer, alias); + await setRowsPerPage(page, "25"); + + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 28", { timeout: 30_000 }); + await expect(requestLogsRows(page)).toHaveCount(25); + // The sessions are the newest groups, so their single representative rows sit on page 1. + await expect(requestLogsRows(page).filter({ hasText: sessionA })).toHaveCount(1); + await expect(requestLogsRows(page).filter({ hasText: sessionA })).toContainText("7"); + await expect(requestLogsRows(page).filter({ hasText: sessionB })).toHaveCount(1); + + await visibleTestId(page, "pagination-next").click(); + + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-28 of 28", { timeout: 30_000 }); + await expect(requestLogsRows(page)).toHaveCount(3); + await expect(requestLogsRows(page).filter({ hasText: sessionA })).toHaveCount(0); + await expect(requestLogsRows(page).filter({ hasText: sessionB })).toHaveCount(0); + }); + + test("two keys reusing one session id stay separate rows", async ({ page, request }) => { + const suffix = uniqueSuffix(); + const mine = await createVirtualKey(request, { key_alias: `e2e-logs-pgn-mine-${suffix}` }); + const theirs = await createVirtualKey(request, { key_alias: `e2e-logs-pgn-theirs-${suffix}` }); + const sharedSession = `sess-pgn-shared-${suffix}`; + + let lastId = ""; + for (let i = 0; i < 2; i++) { + lastId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-pgn-shared-mine-${i}-${suffix}`, + apiKey: mine.key, + traceId: sharedSession, + }); + } + lastId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `logs-pgn-shared-theirs-${suffix}`, + apiKey: theirs.key, + traceId: sharedSession, + }); + await waitForSpendLog(request, lastId); + + await openLogs(page); + const drawer = await openFilterDrawer(page); + await drawer.getByPlaceholder("Enter session ID…").fill(sharedSession); + await drawer.getByRole("button", { name: "Apply Filters" }).click(); + await expect(drawer).not.toBeVisible({ timeout: 10_000 }); + + // One row per caller: reusing a session id must not merge two keys' activity into one row. + await expect(requestLogsRows(page).filter({ hasText: sharedSession })).toHaveCount(2, { timeout: 30_000 }); + + // And each row carries ITS key's totals: two calls badge the first key's row, + // while the other key's single call renders as a plain LLM row. + const mineRow = requestLogsRows(page).filter({ hasText: sharedSession }).filter({ hasText: mine.token }); + const theirsRow = requestLogsRows(page).filter({ hasText: sharedSession }).filter({ hasText: theirs.token }); + await expect(mineRow).toHaveCount(1); + await expect(theirsRow).toHaveCount(1); + await expect(mineRow.getByText("2", { exact: true })).toBeVisible(); + await expect(theirsRow.getByText("LLM", { exact: true })).toBeVisible(); + }); +}); diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index a0dcbf802ef..30b086bab61 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3959,18 +3959,18 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts(): ] mock_prisma = MagicMock() - mock_prisma.db.litellm_spendlogs.group_by = AsyncMock( - return_value=[ - {"session_id": session_id, "_count": {"session_id": 2}}, - ] - ) + mock_prisma.db.litellm_spendlogs.group_by = AsyncMock() mock_prisma.db.query_raw = AsyncMock( return_value=[ { "session_id": session_id, + "api_key": api_key, + "session_total_count": 2, "session_total_spend": 15.0, "mcp_tool_call_count": 1, "mcp_tool_call_spend": 10.0, + "session_llm_count": 1, + "session_agent_count": 0, } ] ) @@ -3995,6 +3995,8 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts(): assert rows[0]["mcp_tool_call_spend"] == 10.0 assert rows[1]["mcp_tool_call_count"] == 1 assert rows[1]["mcp_tool_call_spend"] == 10.0 + assert rows[0]["session_llm_count"] == 1 + assert rows[0]["session_agent_count"] == 0 # Every row in the session carries the full session spend, not just its own assert rows[0]["session_total_spend"] == 15.0 @@ -4003,13 +4005,126 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts(): # Row without a session_id defaults to 1 assert rows[2]["session_total_count"] == 1 - # group_by should have been called with the session_id - mock_prisma.db.litellm_spendlogs.group_by.assert_called_once_with( - by=["session_id"], - where={"session_id": {"in": [session_id]}}, - count={"session_id": True}, + # The count is folded into the single aggregate query; no separate group_by call. + mock_prisma.db.litellm_spendlogs.group_by.assert_not_called() + + +@pytest.mark.asyncio +async def test_build_ui_spend_logs_response_key_split_session_gets_per_key_aggregates(): + """ + Two keys reusing one session id are separate rows under grouped pagination, + and each row must carry ITS key's totals, never the combined session's: + the aggregate query and its lookup are keyed by (session_id, api_key). + """ + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + _build_ui_spend_logs_response, ) + session_id = "sess-shared" + dict_rows = [ + {"request_id": "req-a", "session_id": session_id, "call_type": "completion", "api_key": "key-a"}, + {"request_id": "req-b", "session_id": session_id, "call_type": "completion", "api_key": "key-b"}, + ] + + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "session_id": session_id, + "api_key": "key-a", + "session_total_count": 2, + "session_total_spend": 0.2, + "mcp_tool_call_count": 0, + "mcp_tool_call_spend": 0.0, + "session_cache_hit_count": 1, + "session_llm_count": 2, + "session_agent_count": 0, + }, + { + "session_id": session_id, + "api_key": "key-b", + "session_total_count": 1, + "session_total_spend": 0.7, + "mcp_tool_call_count": 0, + "mcp_tool_call_spend": 0.0, + "session_cache_hit_count": 0, + "session_llm_count": 1, + "session_agent_count": 0, + }, + ] + ) + + result = await _build_ui_spend_logs_response( + prisma_client=mock_prisma, + data=dict_rows, + total_records=2, + page=1, + page_size=50, + total_pages=1, + enrich_session_counts=True, + ) + + rows = result["data"] + assert [(r["session_total_count"], r["session_total_spend"]) for r in rows] == [(2, 0.2), (1, 0.7)] + assert [r["session_cache_hit_count"] for r in rows] == [1, 0] + assert [r["session_llm_count"] for r in rows] == [2, 1] + + aggregate_sql = mock_prisma.db.query_raw.mock_calls[0][1][0] + assert "GROUP BY session_id, api_key" in aggregate_sql + + +@pytest.mark.asyncio +async def test_build_ui_spend_logs_response_empty_api_key_keeps_session_aggregates(): + """ + The spend-log schema defaults api_key to an empty string, which is a real + group value and not a missing one: a multi-call session logged under an + empty key must keep its count and spend instead of degrading to a plain + single-call row. + """ + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + _build_ui_spend_logs_response, + ) + + session_id = "sess-keyless" + dict_rows = [ + {"request_id": "req-1", "session_id": session_id, "call_type": "completion", "api_key": ""}, + ] + + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "session_id": session_id, + "api_key": "", + "session_total_count": 3, + "session_total_spend": 0.09, + "mcp_tool_call_count": 0, + "mcp_tool_call_spend": 0.0, + "session_cache_hit_count": 0, + "session_llm_count": 3, + "session_agent_count": 0, + } + ] + ) + + result = await _build_ui_spend_logs_response( + prisma_client=mock_prisma, + data=dict_rows, + total_records=1, + page=1, + page_size=50, + total_pages=1, + enrich_session_counts=True, + ) + + row = result["data"][0] + assert row["session_total_count"] == 3 + assert row["session_total_spend"] == 0.09 + + # The empty key must reach the aggregate's authorized-keys filter too. + _, call_args, _ = mock_prisma.db.query_raw.mock_calls[0] + assert call_args[2] == [""] + @pytest.mark.asyncio async def test_build_ui_spend_logs_response_sums_multi_round_session_spend(): @@ -4033,14 +4148,13 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend(): ] mock_prisma = MagicMock() - mock_prisma.db.litellm_spendlogs.group_by = AsyncMock( - return_value=[{"session_id": session_id, "_count": {"session_id": 3}}] - ) # The raw aggregate query returns the full session spend (0.01 + 0.02 + 0.03). mock_prisma.db.query_raw = AsyncMock( return_value=[ { "session_id": session_id, + "api_key": api_key, + "session_total_count": 3, "session_total_spend": 0.06, "mcp_tool_call_count": 0, "mcp_tool_call_spend": 0.0, @@ -4089,13 +4203,12 @@ async def test_build_ui_spend_logs_response_session_cache_hit_count(): ] mock_prisma = MagicMock() - mock_prisma.db.litellm_spendlogs.group_by = AsyncMock( - return_value=[{"session_id": session_id, "_count": {"session_id": 2}}] - ) mock_prisma.db.query_raw = AsyncMock( return_value=[ { "session_id": session_id, + "api_key": api_key, + "session_total_count": 2, "session_total_spend": 0.05, "mcp_tool_call_count": 0, "mcp_tool_call_spend": 0.0, diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index ef68d9ce178..27e633099f0 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -274,6 +274,9 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch): "the page query must not carry a window count that forces a full-window " f"scan. SQL was:\n{page_sql}" ) + assert "GROUP BY" not in count_sql and "DISTINCT ON" not in page_sql, ( + "without group_by_session the endpoint must keep raw per-call pagination" + ) assert response["total"] == 137 assert response["total_is_capped"] is False @@ -499,3 +502,106 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch): params = mock_prisma.db.query_raw.call_args[0][1:] assert "team_x" in params, "team_id must be forwarded into the DB query params" assert "sl.team_id = $3" in sql, f"team query must filter on team_id. SQL was:\n{sql}" + + +@pytest.mark.asyncio +async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): + """ + With group_by_session=true, /spend/logs/ui must page and count SESSIONS, + not raw calls: the page query returns one representative row per session + (DISTINCT ON the session group key, preferring non-MCP calls, newest + first) and the bounded count counts groups. Otherwise the UI collapses a + server page of N calls into fewer visible rows while the footer still + claims N (issue #38060). + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + SPEND_LOGS_PAGINATION_COUNT_CAP, + ui_view_spend_logs, + ) + + page_rows = [ + {"request_id": "req-1", "metadata": "{}", "session_id": None}, + {"request_id": "req-2", "metadata": "{}", "session_id": None}, + ] + mock_prisma = _make_ui_spend_logs_mock(count_total=12, page_rows=page_rows) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=1, + page_size=50, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + group_by_session=True, + ) + + group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key" + + count_call = mock_prisma.db.query_raw.call_args_list[0] + count_sql = count_call[0][0] + assert f"GROUP BY {group_key}" in count_sql, f"grouped total must count sessions. SQL was:\n{count_sql}" + assert "COUNT(*) OVER ()" not in count_sql + assert "LIMIT" in count_sql and "FROM (" in count_sql, "the grouped count must stay bounded" + assert count_call[0][-1] == SPEND_LOGS_PAGINATION_COUNT_CAP + 1 + + page_sql = mock_prisma.db.query_raw.call_args_list[1][0][0] + assert f"DISTINCT ON ({group_key})" in page_sql, f"page must return one row per session. SQL was:\n{page_sql}" + assert f"ORDER BY {group_key}, call_type IN ('call_mcp_tool', 'list_mcp_tools'), \"startTime\" DESC" in page_sql, ( + "the session representative must prefer the newest non-MCP call" + ) + assert "COUNT(*) OVER ()" not in page_sql + + assert response["total"] == 12 + assert response["total_is_capped"] is False + assert response["total_pages"] == 1 + + +@pytest.mark.asyncio +async def test_spend_logs_ui_request_id_lookup_with_grouping_returns_exact_row(monkeypatch): + """ + A request_id lookup with group_by_session=true must still resolve the + exact requested row: the filter runs before grouping, so the row is its + own group's representative and deep links keep working. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ui_view_spend_logs + + target_row = {"request_id": "req-deep-link", "metadata": "{}", "session_id": None} + mock_prisma = _make_ui_spend_logs_mock(count_total=1, page_rows=[target_row]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id="req-deep-link", + start_date=None, + end_date=None, + page=1, + page_size=1, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + group_by_session=True, + ) + + page_call = mock_prisma.db.query_raw.call_args_list[1] + assert "request_id = $" in page_call[0][0], "the request_id equality filter must survive grouping" + assert "req-deep-link" in page_call[0] + assert [row["request_id"] for row in response["data"]] == ["req-deep-link"] + assert response["total"] == 1 diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index db7baf5e171..63232afc46b 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2042,6 +2042,7 @@ interface UiSpendLogsParams { min_spend?: number; max_spend?: number; exclude_internal_health_checks?: boolean; + group_by_session?: boolean; } interface UiSpendLogsCallOptions { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx index 22e3f635b50..c368a43dd1e 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx @@ -138,33 +138,39 @@ describe("RequestLogsPanel", () => { respondWith([]); }); - describe("multi-call session collapsing", () => { - const sessionRows = [ - logEntry({ request_id: "req-mcp", call_type: "call_mcp_tool", session_id: "sess-1", session_total_count: 3 }), - logEntry({ request_id: "req-llm", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), - logEntry({ request_id: "req-llm-2", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), - ]; - - it("collapses a multi-call session to a single representative row", async () => { - respondWith(sessionRows); + describe("server-grouped session pagination (#38060)", () => { + it("requests session-grouped pages of 10 rows by default", async () => { renderPanel(); - await waitFor(() => expect(row("req-mcp") ?? row("req-llm") ?? row("req-llm-2")).not.toBeNull()); - - const rendered = ["req-mcp", "req-llm", "req-llm-2"].filter((id) => row(id) !== null); - expect(rendered).toHaveLength(1); + await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled()); + expect(lastCall()?.params?.group_by_session).toBe(true); + expect(lastCall()?.page_size).toBe(10); }); - it("prefers an LLM call over an MCP call as the session's representative", async () => { - respondWith(sessionRows); + it("renders every row the server returns without client-side collapsing", async () => { + respondWith([ + logEntry({ request_id: "req-a", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), + logEntry({ request_id: "req-b", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), + logEntry({ request_id: "req-c", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), + ]); renderPanel(); - await waitFor(() => expect(row("req-llm")).not.toBeNull()); - expect(row("req-mcp")).toBeNull(); + await waitFor(() => expect(row("req-a")).not.toBeNull()); + expect(row("req-b")).not.toBeNull(); + expect(row("req-c")).not.toBeNull(); }); - it("shows the session's call count and composition on the representative row", async () => { - respondWith(sessionRows); + it("shows the session's call count on the server-picked representative row", async () => { + respondWith([ + logEntry({ + request_id: "req-llm", + call_type: "acompletion", + session_id: "sess-1", + session_total_count: 3, + session_llm_count: 2, + mcp_tool_call_count: 1, + }), + ]); renderPanel(); await waitFor(() => expect(row("req-llm")).not.toBeNull()); @@ -296,6 +302,7 @@ describe("RequestLogsPanel", () => { if (!byIdCall) throw new Error("expected a by-id uiSpendLogsCall"); expect(byIdCall.page).toBe(1); expect(byIdCall.page_size).toBe(1); + expect(byIdCall.params?.group_by_session).toBeUndefined(); }); it("closing the drawer removes ?log_id= from the URL and closes the drawer", async () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx index 52ea78abf5e..2ebabf64ab0 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx @@ -10,7 +10,7 @@ import type { KeyResponse } from "../key_team_helpers/key_list"; import { keyInfoV1Call, uiSpendLogsCall } from "../networking"; import KeyInfoView from "../templates/key_info_view"; import type { LogEntry } from "./columns"; -import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "./constants"; +import { LOGS_PAGE_SIZE_OPTIONS } from "./constants"; import { DEFAULT_LOGS_SORTING, formatLogsWindow, @@ -24,7 +24,7 @@ import { LogDetailsDrawer } from "./LogDetailsDrawer"; import { LiveTailBanner, LogsTableToolbar } from "./LogsTableToolbar"; import { RequestLogsTable } from "./RequestLogsTable"; -const PAGE_SIZE = 50; +const PAGE_SIZE = LOGS_PAGE_SIZE_OPTIONS[0]; const DEFAULT_INTERVAL = { value: 24, unit: "hours" }; interface RequestLogsPanelProps { @@ -35,12 +35,6 @@ interface RequestLogsPanelProps { isActive: boolean; } -interface SessionComposition { - llm: number; - agent: number; - mcp: number; -} - export default function RequestLogsPanel({ accessToken, token, userRole, userID, isActive }: RequestLogsPanelProps) { const [pagination, setPagination] = useState({ pageIndex: 0, pageSize: PAGE_SIZE }); const [sorting, setSorting] = useState(DEFAULT_LOGS_SORTING); @@ -157,49 +151,7 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, const isDrawerOpen = displayLog !== null || displaySessionId !== null; - const rows = useMemo(() => { - const searchedLogs = filteredLogs.data; - - const sessionCompositionById = searchedLogs.reduce>((acc, log) => { - if (!log.session_id) return acc; - if (!acc[log.session_id]) { - acc[log.session_id] = { llm: 0, agent: 0, mcp: 0 }; - } - if (MCP_CALL_TYPES.includes(log.call_type)) { - acc[log.session_id].mcp += 1; - } else if (AGENT_CALL_TYPES.includes(log.call_type)) { - acc[log.session_id].agent += 1; - } else { - acc[log.session_id].llm += 1; - } - return acc; - }, {}); - - const sessionRepresentativeMap = new Map(); - for (const log of searchedLogs) { - if (!log.session_id || (log.session_total_count || 1) <= 1) continue; - const isMcp = MCP_CALL_TYPES.includes(log.call_type); - const existing = sessionRepresentativeMap.get(log.session_id); - if (!existing || (existing.isMcp && !isMcp)) { - sessionRepresentativeMap.set(log.session_id, { requestId: log.request_id, isMcp }); - } - } - - return searchedLogs - .map((log) => { - const sessionComposition = log.session_id ? sessionCompositionById[log.session_id] : undefined; - return { - ...log, - session_llm_count: sessionComposition?.llm ?? undefined, - session_mcp_count: sessionComposition?.mcp ?? undefined, - session_agent_count: sessionComposition?.agent ?? undefined, - }; - }) - .filter((log) => { - if (!log.session_id || (log.session_total_count || 1) <= 1) return true; - return sessionRepresentativeMap.get(log.session_id)?.requestId === log.request_id; - }); - }, [filteredLogs.data]); + const rows: LogEntry[] = filteredLogs.data; const searchTerm = useMemo(() => { const entry = columnFilters.find((filter) => filter.id === LOG_FILTER_IDS.REQUEST_ID); @@ -258,13 +210,12 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, ); const handleSessionClick = useCallback( - (sessionId: string) => { - if (!sessionId) return; - const log = rows.find((candidate) => candidate.session_id === sessionId) ?? null; + (log: LogEntry) => { + if (!log.session_id) return; setSelectedLog(log); - openSession(sessionId, log?.request_id ?? null); + openSession(log.session_id, log.request_id); }, - [rows, openSession], + [openSession], ); const handleSelectLog = useCallback( diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx index 7146cf33847..4159b3b699b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx @@ -8,6 +8,7 @@ import { DataTable, DataTableFilterDrawer, DataTableToolbar } from "@/components import type { Team } from "../key_team_helpers/key_list"; import type { LogEntry } from "./columns"; +import { LOGS_PAGE_SIZE_OPTIONS } from "./constants"; import { LOG_FILTER_LABELS, type LogsWindow } from "./log_filter_logic"; import { RequestLogsFilters } from "./RequestLogsFilters"; import { getRequestLogsTableColumns } from "./RequestLogsTableColumns"; @@ -28,7 +29,7 @@ interface RequestLogsTableProps { onRefresh: () => void; onRowClick: (log: LogEntry) => void; onKeyHashClick: (keyHash: string) => void; - onSessionClick: (sessionId: string) => void; + onSessionClick: (log: LogEntry) => void; teams: Team[]; logsWindow: LogsWindow; toolbarChildren?: ReactNode; @@ -91,6 +92,7 @@ export function RequestLogsTable({ paginationMode="server" pagination={pagination} onPaginationChange={onPaginationChange} + pageSizeOptions={LOGS_PAGE_SIZE_OPTIONS} rowCount={rowCount} filterMode="server" columnFilters={columnFilters} diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index 3d0ea03c5d5..ce59c62f1c4 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -85,13 +85,19 @@ describe("row action cells", () => { expect(deps.onKeyHashClick).toHaveBeenCalledWith("sk-hash-9"); }); - it("reports the session id from the session cell", async () => { + it("reports the clicked row from the session cell, so two rows sharing a session id stay distinguishable", async () => { const user = userEvent.setup(); const deps = { onKeyHashClick: vi.fn(), onSessionClick: vi.fn() }; - renderRows([logEntry({ request_id: "req-sess", session_id: "sess-42" })], deps); + renderRows( + [ + logEntry({ request_id: "req-key-a", session_id: "sess-42", api_key: "key-a" }), + logEntry({ request_id: "req-key-b", session_id: "sess-42", api_key: "key-b" }), + ], + deps, + ); - await user.click(screen.getByText("sess-42")); - expect(deps.onSessionClick).toHaveBeenCalledWith("sess-42"); + await user.click(screen.getAllByText("sess-42")[1]); + expect(deps.onSessionClick).toHaveBeenCalledWith(expect.objectContaining({ request_id: "req-key-b" })); }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index cf776515bd4..b9058d02a6b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -13,7 +13,7 @@ import { AgentBadge, AgentIcon, LlmBadge, McpBadge, SparkleIcon, WrenchIcon } fr export interface RequestLogsTableColumnsDeps { onKeyHashClick: (keyHash: string) => void; - onSessionClick: (sessionId: string) => void; + onSessionClick: (log: LogEntry) => void; } const readMetaString = (metadata: Record | undefined, key: string): string | undefined => { @@ -61,7 +61,7 @@ export const getRequestLogsTableColumns = ({ const isAgent = AGENT_CALL_TYPES.includes(log.call_type); const sessionLlmCount = log.session_llm_count ?? (isMcp || isAgent ? 0 : sessionCount); const sessionAgentCount = log.session_agent_count ?? (isAgent ? sessionCount : 0); - const sessionMcpCount = log.session_mcp_count ?? (isMcp ? sessionCount : 0); + const sessionMcpCount = log.mcp_tool_call_count ?? (isMcp ? sessionCount : 0); if (isMcp) return ; if (isAgent && sessionCount <= 1) return ; @@ -113,7 +113,7 @@ export const getRequestLogsTableColumns = ({ header: "Session ID", size: 120, enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => onSessionClick(row.original)} />, }, { id: "request_id", diff --git a/ui/litellm-dashboard/src/components/view_logs/columns.tsx b/ui/litellm-dashboard/src/components/view_logs/columns.tsx index eef957922d7..2f3a3681352 100644 --- a/ui/litellm-dashboard/src/components/view_logs/columns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/columns.tsx @@ -46,6 +46,5 @@ export type LogEntry = { mcp_tool_call_count?: number; mcp_tool_call_spend?: number; session_llm_count?: number; - session_mcp_count?: number; session_agent_count?: number; }; diff --git a/ui/litellm-dashboard/src/components/view_logs/constants.ts b/ui/litellm-dashboard/src/components/view_logs/constants.ts index 5b0b1d0fee3..1c17f398e35 100644 --- a/ui/litellm-dashboard/src/components/view_logs/constants.ts +++ b/ui/litellm-dashboard/src/components/view_logs/constants.ts @@ -12,6 +12,9 @@ export const ERROR_CODE_OPTIONS: { label: string; value: string }[] = [ { label: "529 - Overloaded", value: "529" }, ]; +/** Page sizes the logs tables offer; the first entry is the default. */ +export const LOGS_PAGE_SIZE_OPTIONS = [10, 25, 50, 100]; + /** Call types that represent MCP tool invocations (shared across columns, index, drawer). */ export const MCP_CALL_TYPES = ["call_mcp_tool", "list_mcp_tools"]; diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx index 3b8d96596de..acd5be06593 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx @@ -181,6 +181,7 @@ export function useLogFilterLogic({ sort_by: sortBy, sort_order: sortOrder, exclude_internal_health_checks: excludeInternalHealthChecks, + group_by_session: true, }, }); }, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index be81450b941..7583b6f6df1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -56755,6 +56755,8 @@ export interface operations { sort_order?: string | null; /** @description Exclude LiteLLM internal health check requests from results */ exclude_internal_health_checks?: boolean; + /** @description Paginate over sessions instead of raw logs: one representative row per session, total counts sessions */ + group_by_session?: boolean; }; header?: never; path?: never; @@ -56867,6 +56869,8 @@ export interface operations { sort_order?: string | null; /** @description Exclude LiteLLM internal health check requests from results */ exclude_internal_health_checks?: boolean; + /** @description Paginate over sessions instead of raw logs: one representative row per session, total counts sessions */ + group_by_session?: boolean; }; header?: never; path?: never; From 993766be0e1a1ea0f173dc9445f85e4eab3665bc Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 2 Sep 2026 18:03:22 -0700 Subject: [PATCH 08/23] feat(proxy): serve the auto-router preset catalog at runtime (#39412) The dashboard's template picker imported autorouter_presets.json at build time, so every catalog change needed a dashboard rebuild and artifacts refresh. The catalog now lives in litellm/proxy/public_endpoints/ and GET /public/autorouter_presets serves it, fetching litellm.autorouter_presets_url (GitHub raw on main, 1h in-process cache, bundled fallback) so a merged catalog change propagates to running proxies like the model cost map does. The dashboard fetches it at runtime via useAutoRouterPresets and keeps no local copy. Resolves LIT-6764 --- litellm/__init__.py | 4 + .../public_endpoints}/autorouter_presets.json | 0 .../public_endpoints/public_endpoints.py | 85 ++++++- .../public_endpoints/public_endpoints.py | 42 ++- .../public_endpoints/test_public_endpoints.py | 240 ++++++++++++++++++ .../hooks/autoRouter/useAutoRouterPresets.ts | 16 ++ .../add_model/add_auto_router_tab.test.tsx | 59 ++++- .../add_model/add_auto_router_tab.tsx | 28 +- .../src/components/networking.tsx | 10 + .../src/lib/autorouter_presets.test.ts | 13 +- .../src/lib/autorouter_presets.ts | 13 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 90 +++++++ .../tests/mocks/autoRouterPresets.ts | 16 ++ 13 files changed, 592 insertions(+), 24 deletions(-) rename {ui/litellm-dashboard/src => litellm/proxy/public_endpoints}/autorouter_presets.json (100%) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets.ts create mode 100644 ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts diff --git a/litellm/__init__.py b/litellm/__init__.py index 44f2e7c1f02..41a3789ab0d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -424,6 +424,10 @@ anthropic_beta_headers_url: str = os.getenv( "LITELLM_ANTHROPIC_BETA_HEADERS_URL", "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/anthropic_beta_headers_config.json", ) +autorouter_presets_url: str = os.getenv( + "LITELLM_AUTOROUTER_PRESETS_URL", + "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/autorouter_presets.json", +) suppress_debug_info: bool = False dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None diff --git a/ui/litellm-dashboard/src/autorouter_presets.json b/litellm/proxy/public_endpoints/autorouter_presets.json similarity index 100% rename from ui/litellm-dashboard/src/autorouter_presets.json rename to litellm/proxy/public_endpoints/autorouter_presets.json diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 4d58a974bb8..94a59828451 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -1,11 +1,13 @@ +import asyncio import json import os import re -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from importlib.resources import files from typing import TYPE_CHECKING, Final, Protocol from fastapi import APIRouter, HTTPException, Request +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm @@ -28,6 +30,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import ) from litellm.types.proxy.public_endpoints.public_endpoints import ( AgentCreateInfo, + AutoRouterPresetRecord, ComplexityScorerDefaults, ProviderCreateInfo, PublicModelHubInfo, @@ -464,6 +467,86 @@ async def get_litellm_blog_posts(): return BlogPostsResponse(posts=posts) +_AUTOROUTER_PRESETS_ADAPTER: Final = TypeAdapter(dict[str, AutoRouterPresetRecord]) + + +def _load_bundled_autorouter_presets() -> Mapping[str, AutoRouterPresetRecord]: + raw: Final = json.loads( + files("litellm.proxy.public_endpoints").joinpath("autorouter_presets.json").read_text(encoding="utf-8") + ) + return _AUTOROUTER_PRESETS_ADAPTER.validate_python(raw) + + +async def _fetch_remote_autorouter_presets(url: str) -> Mapping[str, AutoRouterPresetRecord]: + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.llms.custom_http import httpxSpecialProvider + + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.UI) + response: Final = await client.get(url, timeout=5.0) + response.raise_for_status() + presets: Final = _AUTOROUTER_PRESETS_ADAPTER.validate_python(response.json()) + if not presets: + raise ValueError("remote auto-router preset catalog is empty") + return presets + + +async def _resolve_autorouter_presets( + url: str, + fetch: Callable[[str], Awaitable[Mapping[str, AutoRouterPresetRecord]]], +) -> Mapping[str, AutoRouterPresetRecord]: + if os.getenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", "").lower() == "true": + return _load_bundled_autorouter_presets() + try: + return await fetch(url) + except Exception as e: + verbose_logger.warning( + "LiteLLM: failed to fetch auto-router presets from %s: %s. Serving the bundled catalog for the life of this process.", + url, + str(e), + ) + return _load_bundled_autorouter_presets() + + +class _AutoRouterPresetsCache: + presets: Mapping[str, AutoRouterPresetRecord] | None = None + lock: asyncio.Lock | None = None + + +async def get_autorouter_presets( + url: str, + fetch: Callable[[str], Awaitable[Mapping[str, AutoRouterPresetRecord]]] = _fetch_remote_autorouter_presets, +) -> Mapping[str, AutoRouterPresetRecord]: + cached: Final = _AutoRouterPresetsCache.presets + if cached is not None: + return cached + if _AutoRouterPresetsCache.lock is None: + _AutoRouterPresetsCache.lock = asyncio.Lock() + async with _AutoRouterPresetsCache.lock: + held: Final = _AutoRouterPresetsCache.presets + if held is not None: + return held + resolved: Final = await _resolve_autorouter_presets(url=url, fetch=fetch) + _AutoRouterPresetsCache.presets = resolved + return resolved + + +@router.get( + "/public/autorouter_presets", + tags=["public", "auto router"], # mutable-ok: FastAPI route tags take a list + response_model=dict[str, AutoRouterPresetRecord], +) +async def get_public_autorouter_presets() -> Mapping[str, AutoRouterPresetRecord]: + """ + Return the auto-router preset catalog the dashboard's template picker renders. + + Resolved once per process, like the model cost map: fetched from ``litellm.autorouter_presets_url`` + (override with ``LITELLM_AUTOROUTER_PRESETS_URL``) on the first request, falling back to the + catalog bundled with the package on any failure. Set ``LITELLM_LOCAL_AUTOROUTER_PRESETS=True`` + to serve the bundled catalog only. A restart picks up a newly published catalog. + """ + return await get_autorouter_presets(url=litellm.autorouter_presets_url) + + @router.get( "/public/endpoints", tags=["public"], diff --git a/litellm/types/proxy/public_endpoints/public_endpoints.py b/litellm/types/proxy/public_endpoints/public_endpoints.py index c7f80a61e0f..fa73926305b 100644 --- a/litellm/types/proxy/public_endpoints/public_endpoints.py +++ b/litellm/types/proxy/public_endpoints/public_endpoints.py @@ -1,7 +1,7 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Any, Literal -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict class PublicModelHubInfo(BaseModel): @@ -73,6 +73,44 @@ class SupportedEndpointsResponse(BaseModel): endpoints: list[SupportedEndpoint] +class AutoRouterPresetTiers(BaseModel): + """Exactly the four built-in tiers the dashboard's preset prefill can apply. + + extra="forbid" on purpose: a tier name this dashboard cannot apply would grey out or crash the + picker, so such a catalog is rejected wholesale and the bundled one serves instead. + """ + + model_config = ConfigDict(extra="forbid") + + SIMPLE: Sequence[str] + MEDIUM: Sequence[str] + COMPLEX: Sequence[str] + REASONING: Sequence[str] + + +class AutoRouterPresetConfig(BaseModel): + """The complexity_router_config a preset prefills. + + Only tiers is validated, because every dashboard consumer dereferences it; everything else + passes through verbatim with unknown fields kept (extra="allow"), so a catalog published after + this proxy shipped still serves its new fields intact. + """ + + model_config = ConfigDict(extra="allow") + + tiers: AutoRouterPresetTiers + + +class AutoRouterPresetRecord(BaseModel): + """One auto-router preset as served to the dashboard's template picker.""" + + model_config = ConfigDict(extra="allow") + + label: str + description: str + complexity_router_config: AutoRouterPresetConfig + + class ComplexityScorerDefaults(BaseModel): """The complexity router's shipped heuristic scorer defaults. diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 8006f64ba41..551e27a18f5 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -1077,3 +1077,243 @@ def test_public_mcp_hub_does_not_expose_upstream_url(): assert all("url" not in item for item in data) assert secret_url not in response.text app.dependency_overrides.clear() + + + +@pytest.fixture +def reset_autorouter_presets_cache(): + from litellm.proxy.public_endpoints.public_endpoints import _AutoRouterPresetsCache + + _AutoRouterPresetsCache.presets = None + _AutoRouterPresetsCache.lock = None + yield + _AutoRouterPresetsCache.presets = None + _AutoRouterPresetsCache.lock = None + + +def test_get_autorouter_presets_local_mode_serves_bundled_catalog( + monkeypatch, reset_autorouter_presets_cache +): + monkeypatch.setenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", "True") + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + response = client.get("/public/autorouter_presets") + + assert response.status_code == 200 + payload = response.json() + assert "anthropic_family" in payload + for preset in payload.values(): + assert isinstance(preset["label"], str) + assert isinstance(preset["description"], str) + assert "tiers" in preset["complexity_router_config"] + + +@pytest.mark.asyncio +async def test_get_autorouter_presets_fetches_once_per_process( + monkeypatch, reset_autorouter_presets_cache +): + from litellm.proxy.public_endpoints.public_endpoints import ( + _AUTOROUTER_PRESETS_ADAPTER, + get_autorouter_presets, + ) + + monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False) + remote = _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "remote_only": { + "label": "Remote Only", + "description": "from the remote catalog", + "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}}, + } + } + ) + calls = [] + + async def fake_fetch(url): + calls.append(url) + return remote + + first = await get_autorouter_presets(url="https://example.test/presets.json", fetch=fake_fetch) + second = await get_autorouter_presets(url="https://example.test/presets.json", fetch=fake_fetch) + + assert first == remote + assert second == remote + assert calls == ["https://example.test/presets.json"] + + +@pytest.mark.asyncio +async def test_get_autorouter_presets_single_flight_on_concurrent_cold_start( + monkeypatch, reset_autorouter_presets_cache +): + import asyncio + + from litellm.proxy.public_endpoints.public_endpoints import ( + _AUTOROUTER_PRESETS_ADAPTER, + get_autorouter_presets, + ) + + monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False) + remote = _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "remote_only": { + "label": "Remote Only", + "description": "from the remote catalog", + "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}}, + } + } + ) + calls = [] + + async def slow_fetch(url): + calls.append(url) + await asyncio.sleep(0.05) + return remote + + results = await asyncio.gather( + get_autorouter_presets(url="https://example.test/presets.json", fetch=slow_fetch), + get_autorouter_presets(url="https://example.test/presets.json", fetch=slow_fetch), + get_autorouter_presets(url="https://example.test/presets.json", fetch=slow_fetch), + ) + + assert all(result == remote for result in results) + assert len(calls) == 1 + + +@pytest.mark.asyncio +async def test_get_autorouter_presets_caches_bundled_fallback_on_remote_failure( + monkeypatch, reset_autorouter_presets_cache +): + from litellm.proxy.public_endpoints.public_endpoints import get_autorouter_presets + + monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False) + calls = [] + + async def broken_fetch(url): + calls.append(url) + raise ValueError("remote catalog unavailable") + + first = await get_autorouter_presets(url="https://example.test/presets.json", fetch=broken_fetch) + second = await get_autorouter_presets(url="https://example.test/presets.json", fetch=broken_fetch) + + assert "anthropic_family" in first + assert second == first + assert len(calls) == 1 + + +@pytest.mark.asyncio +async def test_autorouter_presets_adapter_rejects_wrong_shapes(): + from pydantic import ValidationError + + from litellm.proxy.public_endpoints.public_endpoints import _AUTOROUTER_PRESETS_ADAPTER + + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python({"bad": {"label": "no description or config"}}) + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python(["not", "a", "mapping"]) + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python( + {"no_tiers": {"label": "L", "description": "D", "complexity_router_config": {}}} + ) + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "missing_builtin_tier": { + "label": "L", + "description": "D", + "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"]}}, + } + } + ) + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "unknown_tier_name": { + "label": "L", + "description": "D", + "complexity_router_config": { + "tiers": { + "SIMPLE": ["m1"], + "MEDIUM": ["m2"], + "COMPLEX": ["m3"], + "REASONING": ["m4"], + "ULTRA": ["m5"], + } + }, + } + } + ) + with pytest.raises(ValidationError): + _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "bad_tiers": { + "label": "L", + "description": "D", + "complexity_router_config": {"tiers": "not-a-mapping"}, + } + } + ) + + +def test_get_autorouter_presets_passes_unknown_catalog_fields_through( + monkeypatch, reset_autorouter_presets_cache +): + from litellm.proxy.public_endpoints.public_endpoints import ( + _AUTOROUTER_PRESETS_ADAPTER, + _AutoRouterPresetsCache, + ) + + monkeypatch.delenv("LITELLM_LOCAL_AUTOROUTER_PRESETS", raising=False) + _AutoRouterPresetsCache.presets = _AUTOROUTER_PRESETS_ADAPTER.validate_python( + { + "future_preset": { + "label": "Future", + "description": "carries fields this proxy version does not know", + "complexity_router_config": { + "tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}, + "future_config_knob": 3, + }, + "icon": "sparkles", + } + } + ) + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + response = client.get("/public/autorouter_presets") + + assert response.status_code == 200 + served = response.json()["future_preset"] + assert served["icon"] == "sparkles" + assert served["complexity_router_config"]["future_config_knob"] == 3 + assert served["complexity_router_config"]["tiers"]["SIMPLE"] == ["m1"] + + +@pytest.mark.asyncio +async def test_fetch_remote_autorouter_presets_parses_and_rejects_empty(monkeypatch): + import litellm.llms.custom_httpx.http_handler as http_handler_module + from litellm.proxy.public_endpoints.public_endpoints import _fetch_remote_autorouter_presets + + catalog = { + "remote_only": { + "label": "Remote Only", + "description": "from the remote catalog", + "complexity_router_config": {"tiers": {"SIMPLE": ["m1"], "MEDIUM": ["m2"], "COMPLEX": ["m3"], "REASONING": ["m4"]}}, + } + } + response = MagicMock() + response.raise_for_status = MagicMock() + response.json = MagicMock(return_value=catalog) + client = MagicMock() + client.get = AsyncMock(return_value=response) + monkeypatch.setattr(http_handler_module, "get_async_httpx_client", lambda llm_provider: client) + + presets = await _fetch_remote_autorouter_presets("https://example.test/presets.json") + assert presets["remote_only"].label == "Remote Only" + response.raise_for_status.assert_called_once() + + response.json = MagicMock(return_value={}) + with pytest.raises(ValueError, match="empty"): + await _fetch_remote_autorouter_presets("https://example.test/presets.json") diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets.ts new file mode 100644 index 00000000000..82f9cc2da99 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets.ts @@ -0,0 +1,16 @@ +import { AutoRouterPreset, hydratePresets } from "@/lib/autorouter_presets"; +import { getAutoRouterPresets } from "@/components/networking"; +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +const presetKeys = createQueryKeys("autoRouterPresets"); + +export const useAutoRouterPresets = () => { + const options = { + queryKey: presetKeys.list({}), + queryFn: async () => hydratePresets(await getAutoRouterPresets()), + staleTime: 24 * 60 * 60 * 1000, + gcTime: 24 * 60 * 60 * 1000, + }; + return useQuery(options); +}; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index ba380662403..7769042a0c7 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -9,11 +9,19 @@ import { getSubmitBlockedReason } from "./add_auto_router_tab"; import { buildModelAvailability } from "@/lib/autorouter_presets"; import { testAutoRouterRouting } from "../networking"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; -import { getAllPresets, getPresetByKey, getRequiredModelsInPreset } from "@/lib/autorouter_presets"; +import { AutoRouterPreset, getRequiredModelsInPreset } from "@/lib/autorouter_presets"; +import { BUNDLED_PRESETS, LOADED_PRESETS_QUERY, useAutoRouterPresets } from "../../../tests/mocks/autoRouterPresets"; vi.mock( "@/app/(dashboard)/hooks/autoRouter/useComplexityScorerDefaults", async () => await import("../../../tests/mocks/complexityScorerDefaults"), ); +vi.mock( + "@/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets", + async () => await import("../../../tests/mocks/autoRouterPresets"), +); + +const getAllPresets = (): AutoRouterPreset[] => BUNDLED_PRESETS; +const getPresetByKey = (key: string): AutoRouterPreset | undefined => BUNDLED_PRESETS.find((p) => p.key === key); const ANTHROPIC_PRESET = getPresetByKey("anthropic_family")!; const ANTHROPIC_TIERS = ANTHROPIC_PRESET.complexity_router_config.tiers; @@ -1142,3 +1150,52 @@ describe("getSubmitBlockedReason", () => { ); }); }); + +describe("preset catalog fetch states", () => { + afterEach(() => vi.mocked(useAutoRouterPresets).mockReturnValue(LOADED_PRESETS_QUERY)); + + it("keeps showing cached presets without the error banner when only a refetch fails", () => { + vi.mocked(useAutoRouterPresets).mockReturnValue({ + ...LOADED_PRESETS_QUERY, + isError: true, + } as never); + renderWithProviders(); + + expect(screen.queryByText(/Could not load templates/)).not.toBeInTheDocument(); + + openTemplateDropdown(); + expect(screen.queryAllByRole("option").length).toBeGreaterThan(1); + }); + + it("shows a loading hint while the catalog fetch is pending", () => { + vi.mocked(useAutoRouterPresets).mockReturnValue({ + ...LOADED_PRESETS_QUERY, + data: undefined, + isPending: true, + } as never); + renderWithProviders(); + + expect(screen.getByText("Loading templates...")).toBeInTheDocument(); + }); + + it("degrades to Custom Configuration with a retry hint that refetches the catalog", async () => { + const refetch = vi.fn(); + vi.mocked(useAutoRouterPresets).mockReturnValue({ + ...LOADED_PRESETS_QUERY, + data: undefined, + isError: true, + refetch, + } as never); + renderWithProviders(); + + expect(await screen.findByText(/Could not load templates/)).toBeInTheDocument(); + + openTemplateDropdown(); + const options = screen.queryAllByRole("option"); + expect(options).toHaveLength(1); + expect(options[0]).toHaveTextContent("Custom Configuration"); + + fireEvent.click(screen.getByRole("button", { name: "Retry" })); + expect(refetch).toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 81acb7b7e50..ab5ea4cbfc9 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -50,8 +50,6 @@ import AutoRouterConnectionTest from "./auto_router_connection_test"; import AutoRouterRoutingTest from "./AutoRouterRoutingTest"; import { toast } from "@/lib/toast"; import { - getAllPresets, - getPresetByKey, getMissingModelsInPreset, getReferencedModelsError, buildEmptyPrefill, @@ -62,6 +60,7 @@ import { PresetPrefill, AutoRouterPreset, } from "@/lib/autorouter_presets"; +import { useAutoRouterPresets } from "@/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; interface AddAutoRouterTabProps { @@ -102,9 +101,7 @@ const presetDisabledHint = (availability: PresetAvailability): string | null => // caller-specific missing-model reason gets the alarming red treatment. const isPresetHintAlarming = (availability: PresetAvailability): boolean => availability.kind === "missing_models"; -// getAllPresets() already returns a stable, module-level array (see autorouter_presets.ts), so -// this is resolved once at import time rather than re-called from inside the component every render. -const presets = getAllPresets(); +const NO_PRESETS: AutoRouterPreset[] = []; // A one-line summary of what's configured, shown when the detailed section is collapsed so a // caller can see the shape of the config without opening it. @@ -229,6 +226,14 @@ const AddAutoRouterTab: React.FC = ({ }); const modelsLoading = groupsLoading || deploymentsLoading; const modelInfo = React.useMemo(() => data ?? [], [data]); + const { + data: presetsData, + isPending: presetsPending, + isError: presetsError, + refetch: refetchPresets, + } = useAutoRouterPresets(); + const presets = presetsData ?? NO_PRESETS; + const presetsUnavailable = presetsError && presetsData === undefined; // react-query keeps the last successful list around when a later refetch fails, so isError alone // can't tell "never loaded" apart from "loaded, then a background refetch errored" - only the // former leaves us with nothing trustworthy to verify a preset's models against. @@ -277,7 +282,7 @@ const AddAutoRouterTab: React.FC = ({ presets .map((preset) => ({ preset, availability: presetAvailability(preset) })) .sort((a, b) => Number(b.availability.kind === "available") - Number(a.availability.kind === "available")), - [presetAvailability], + [presets, presetAvailability], ); const templateItems = React.useMemo( @@ -307,7 +312,7 @@ const AddAutoRouterTab: React.FC = ({ return; } - const preset = getPresetByKey(presetKey); + const preset = presets.find((p) => p.key === presetKey); // Refuse to apply a preset whose models are not verified available. The dropdown disables // these options, so this is a guard against a stale click resolving after the list changed. if (!preset) return; @@ -538,6 +543,15 @@ const AddAutoRouterTab: React.FC = ({ )} + {presetsPending &&
Loading templates...
} + {presetsUnavailable && ( +
+ Could not load templates, so only Custom Configuration is shown.{" "} + +
+ )} {requiresTeamScope && ( diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 63232afc46b..eaffa2b4802 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -90,6 +90,7 @@ import type { } from "@/app/(dashboard)/caching/_components/coordination_redis_settings/types"; import { MCP_TOOLS_PREVIEW_FORBIDDEN_MESSAGE } from "./mcp_tools/constants"; import type { ComplexityRouterConfigPayload } from "./add_model/build_complexity_router_config"; +import type { AutoRouterPresetsResponse } from "@/lib/autorouter_presets"; import type { VectorStoreIndex } from "@/app/(dashboard)/vector-stores/_components/IndexesTab"; import type { RoutingDecision } from "./view_logs/LogDetailsDrawer/RoutingDecisionCard"; import { @@ -410,6 +411,15 @@ export const getComplexityScorerDefaults = async (): Promise => { + /** + * Fetch the auto-router preset catalog from the proxy's public endpoint. The template picker + * renders from this rather than from a copy in the dashboard, so a catalog update propagates + * without a dashboard release. + */ + return await apiClient.get(`/public/autorouter_presets`); +}; + export const getAgentCreateMetadata = async (): Promise => { /** * Fetch agent type metadata from the proxy's public endpoint. diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index b5802ffa2dc..e3132067f4c 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -1,7 +1,9 @@ import { describe, it, expect } from "vitest"; +import bundledPresets from "../../../../litellm/proxy/public_endpoints/autorouter_presets.json"; import { - getAllPresets, - getPresetByKey, + hydratePresets, + AutoRouterPreset, + AutoRouterPresetsResponse, getRequiredModelsInPreset, getMissingModelsInPreset, getRequiredModels, @@ -18,8 +20,13 @@ import { DEFAULT_ESCALATION_KEYWORDS } from "@/components/add_model/EscalationKe const groupsOnly = (models: Iterable) => buildModelAvailability(models, []); +// Hydrated from the real bundled catalog so a catalog edit flows into these expectations. +const PRESETS = hydratePresets(bundledPresets as AutoRouterPresetsResponse); +const getAllPresets = (): AutoRouterPreset[] => PRESETS; +const getPresetByKey = (key: string): AutoRouterPreset | undefined => PRESETS.find((p) => p.key === key); + describe("autorouter_presets", () => { - it("loads exactly the bundled presets", () => { + it("hydrates exactly the bundled presets", () => { const presets = getAllPresets(); expect(presets.map((p) => p.label).sort()).toEqual(["Anthropic Family", "Gemini Family", "Lite", "OpenAI Family"]); // Every preset carries all four fields the UI relies on; a JSON typo dropping one fails here. diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index 721bd6f2b2a..30f8ed99e74 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -20,7 +20,6 @@ import { } from "@/components/add_model/complexity_router_tiers"; import { DEFAULT_ESCALATION_KEYWORDS } from "@/components/add_model/EscalationKeywords"; import { DEFAULT_MATCH_THRESHOLD } from "@/components/add_model/SemanticKeywordMatching"; -import presetsRaw from "@/autorouter_presets.json"; // `key` is the stable JSON object key (e.g. "anthropic_family"); `label` is display text and // never an identity. @@ -31,16 +30,10 @@ export interface AutoRouterPreset { complexity_router_config: ComplexityRouterConfigPayload; } -// The bundled JSON is a developer-authored, build-time asset, so it is trusted at the import -// boundary rather than re-validated at runtime (resolveJsonModule widens its string literals, -// hence this one cast). autorouter_presets.test.ts pins the parsed shape, so a JSON typo fails CI. -const RAW = presetsRaw as Record>; +export type AutoRouterPresetsResponse = Record>; -const PRESETS: AutoRouterPreset[] = Object.entries(RAW).map(([key, preset]) => ({ key, ...preset })); - -export const getAllPresets = (): AutoRouterPreset[] => PRESETS; - -export const getPresetByKey = (key: string): AutoRouterPreset | undefined => PRESETS.find((p) => p.key === key); +export const hydratePresets = (raw: AutoRouterPresetsResponse): AutoRouterPreset[] => + Object.entries(raw).map(([key, preset]) => ({ key, ...preset })); // Generalized over ComplexityRouterConfigPayload so the same accessors check either a preset's own // bundled config or a caller's actually-built config - the two need to agree, since a preset only diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7583b6f6df1..90c4f03bf08 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -12129,6 +12129,31 @@ export interface paths { patch?: never; trace?: never; }; + "/public/autorouter_presets": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Public Autorouter Presets + * @description Return the auto-router preset catalog the dashboard's template picker renders. + * + * Resolved once per process, like the model cost map: fetched from ``litellm.autorouter_presets_url`` + * (override with ``LITELLM_AUTOROUTER_PRESETS_URL``) on the first request, falling back to the + * catalog bundled with the package on any failure. Set ``LITELLM_LOCAL_AUTOROUTER_PRESETS=True`` + * to serve the bundled catalog only. A restart picks up a newly published catalog. + */ + get: operations["get_public_autorouter_presets_public_autorouter_presets_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/public/complexity_router/scorer_defaults": { parameters: { query?: never; @@ -23391,6 +23416,49 @@ export interface components { /** Tier Definitions */ tier_definitions: components["schemas"]["TierDefinition"][]; }; + /** + * AutoRouterPresetConfig + * @description The complexity_router_config a preset prefills. + * + * Only tiers is validated, because every dashboard consumer dereferences it; everything else + * passes through verbatim with unknown fields kept (extra="allow"), so a catalog published after + * this proxy shipped still serves its new fields intact. + */ + AutoRouterPresetConfig: { + tiers: components["schemas"]["AutoRouterPresetTiers"]; + } & { + [key: string]: unknown; + }; + /** + * AutoRouterPresetRecord + * @description One auto-router preset as served to the dashboard's template picker. + */ + AutoRouterPresetRecord: { + complexity_router_config: components["schemas"]["AutoRouterPresetConfig"]; + /** Description */ + description: string; + /** Label */ + label: string; + } & { + [key: string]: unknown; + }; + /** + * AutoRouterPresetTiers + * @description Exactly the four built-in tiers the dashboard's preset prefill can apply. + * + * extra="forbid" on purpose: a tier name this dashboard cannot apply would grey out or crash the + * picker, so such a catalog is rejected wholesale and the bundled one serves instead. + */ + AutoRouterPresetTiers: { + /** Complex */ + COMPLEX: string[]; + /** Medium */ + MEDIUM: string[]; + /** Reasoning */ + REASONING: string[]; + /** Simple */ + SIMPLE: string[]; + }; /** * AutoRouterRoutingTestRequest * @description A single request to classify against a complexity-router config that need not be saved yet. @@ -54679,6 +54747,28 @@ export interface operations { }; }; }; + get_public_autorouter_presets_public_autorouter_presets_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: components["schemas"]["AutoRouterPresetRecord"]; + }; + }; + }; + }; + }; get_complexity_scorer_defaults_public_complexity_router_scorer_defaults_get: { parameters: { query?: never; diff --git a/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts new file mode 100644 index 00000000000..f73faaa70fe --- /dev/null +++ b/ui/litellm-dashboard/tests/mocks/autoRouterPresets.ts @@ -0,0 +1,16 @@ +import { vi } from "vitest"; +import bundledPresets from "../../../../litellm/proxy/public_endpoints/autorouter_presets.json"; +import { hydratePresets, type AutoRouterPresetsResponse } from "@/lib/autorouter_presets"; + +// Derived from the real bundled catalog so a preset edit there flows into test expectations +// instead of redding on a stale copy. Exported as vi.fn so a test can override the query state. +export const BUNDLED_PRESETS = hydratePresets(bundledPresets as AutoRouterPresetsResponse); + +export const LOADED_PRESETS_QUERY = { + data: BUNDLED_PRESETS, + isPending: false, + isError: false, + refetch: vi.fn(), +}; + +export const useAutoRouterPresets = vi.fn(() => LOADED_PRESETS_QUERY); From d33fe95d192a819eef63848468c3630903a135a2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:08:45 -0700 Subject: [PATCH 09/23] test(responses): expect the 404 OpenAI now returns for an unknown model --- .../test_e2e_openai_responses_api.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index 1a7fb1f3e41..abae26e02cd 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -1,5 +1,5 @@ import httpx -from openai import OpenAI, BadRequestError, APIStatusError +from openai import OpenAI, BadRequestError, NotFoundError, APIStatusError import pytest @@ -105,10 +105,9 @@ def test_streaming_response(): assert len(collected_chunks) > 0 -def test_bad_request_error(): +def test_model_not_found_error(): client = get_test_client() - with pytest.raises(BadRequestError): - # Trigger error with invalid model name + with pytest.raises(NotFoundError): client.responses.create(model="non-existent-model", input="This should fail") From 87618b7da8e755cb007b3bf2470dd05b93154f6d Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 2 Sep 2026 18:12:03 -0700 Subject: [PATCH 10/23] doc (#39456) --- tests/rust-python-harness/AGENTS.md | 44 +++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 tests/rust-python-harness/AGENTS.md diff --git a/tests/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md new file mode 100644 index 00000000000..e9b17027ddc --- /dev/null +++ b/tests/rust-python-harness/AGENTS.md @@ -0,0 +1,44 @@ +# Expected Structure + +```text +tests/rust-python-harness/ +├── __main__.py +│ +├── strategies/ +│ ├── e2e_parity/ +│ │ ├── runner.py +│ │ ├── sdk/ +│ │ │ ├── ocr/ +│ │ │ ├── messages/ +│ │ │ ├── chat_completions/ +│ │ │ └── responses/ +│ │ └── gateway/ +│ │ +│ ├── trace_parity/ +│ │ ├── runner.py +│ │ ├── sdk/ +│ │ └── gateway/ +│ │ +│ └── unit_tests/ +│ ├── runner.py +│ ├── mapping_validator.py +│ ├── python_runner.py +│ └── rust_runner.py +│ +└── shared/ + ├── parity/ + ├── tracing/ + └── reporting/ +``` + +- Run locally only; no CI integration +- `__main__.py` selects strategies and combines their reports; each strategy also runs independently +- `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses +- `trace_parity/` compares mapped operations, call counts, and required execution ordering +- E2E and trace runners share orchestration across `sdk/` and `gateway/`; surface-specific execution lives in those folders +- `unit_tests/runner.py` combines mapping validation, Python test runs, and native Rust test runs +- `mapping_validator.py` matches Python/Rust tests by agreed names or annotations and reports missing or ambiguous counterparts +- `python_runner.py` runs existing Python tests with Rust disabled and enabled in separate processes, verifies backend selection, and compares results +- `rust_runner.py` runs Cargo tests; native Rust unit tests stay beside their implementation +- `shared/` contains reusable parity, tracing, and reporting machinery +- Keep fixtures with their owning API and existing Python tests in their current locations From 4f7b20ec102bc4f01f32156bc9d0d8325af3a081 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:25:28 -0700 Subject: [PATCH 11/23] fix(guardrails): skip streaming guardrail rounds that re-scan cleared output (#39386) * fix(guardrails): skip streaming guardrail rounds that re-scan cleared output Streaming guardrails scanned the finished answer twice at end of stream whenever the chunk count landed on a multiple of the sampling rate, ran sampled rounds whose payload was identical to the previous one, and on /v1/messages could scan an empty text before the first content chunk. Every redundant round is a paid guardrail provider call. Each endpoint handler now exposes a scan key describing what a round would hand to apply_guardrail (the text so far, plus tool calls once the stream has ended), and the unified streaming hook skips a sampled or end-of-stream round whose key equals the last scanned one or carries nothing to scan yet. Rounds that carry tool calls are never skipped. * test(guardrails): expect one end-of-stream scan when the terminal chunk is sampled Update sampled cadence expectations and use tuple-backed scan state Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../a2a/chat/guardrail_translation/handler.py | 12 +- .../chat/guardrail_translation/handler.py | 25 +- .../guardrail_translation/base_translation.py | 19 ++ .../base_llm/guardrail_translation/utils.py | 12 + .../chat/guardrail_translation/handler.py | 37 ++- .../guardrail_translation/handler.py | 72 ++++- .../unified_guardrail/unified_guardrail.py | 40 ++- .../test_a2a_guardrail_handler.py | 35 +++ .../test_anthropic_guardrail_handler.py | 54 ++++ .../test_openai_guardrail_handler.py | 72 +++++ ...test_openai_responses_guardrail_handler.py | 90 ++++++ .../test_openai_moderation_streaming.py | 19 +- .../test_generic_guardrail_api.py | 18 +- .../test_unified_guardrail.py | 268 ++++++++++++++++++ 14 files changed, 733 insertions(+), 40 deletions(-) create mode 100644 tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index f1c7451796d..5c30ff4747a 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -17,7 +17,10 @@ from typing import TYPE_CHECKING, Any, Final, Optional from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -313,9 +316,14 @@ class A2AGuardrailHandler(BaseTranslation): return responses_so_far + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + _, valid_parsed = self._parse_streaming_responses(responses_so_far) + combined_text, _ = self._collect_text_from_parsed_chunks(valid_parsed) + return StreamingScanKey(texts=(combined_text,)) + def _parse_streaming_responses( self, - responses_so_far: list[object], + responses_so_far: Sequence[object], ) -> tuple[list[dict[str, object] | None], list[tuple[int, dict[str, object]]]]: """Parse JSON-RPC items, returning aligned parsed list and valid entries.""" parsed: Final[list[dict[str, object] | None]] = [None] * len(responses_so_far) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index c7d12e5cf3a..c23797f72af 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -26,7 +26,10 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im LiteLLMAnthropicMessagesAdapter, is_provider_native_tool_dict, ) -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.llms.base_llm.guardrail_translation.utils import ( anthropic_tool_name, anthropic_tool_names, @@ -36,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, scoped_structured_message_indices, + stream_item_fingerprint, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -1176,6 +1180,25 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["model"] = response_model return inputs + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + stream_ended: Final = self._check_streaming_has_ended(responses_so_far) + return StreamingScanKey( + texts=(self.get_streaming_string_so_far(responses_so_far),), + tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (), + stream_ended=stream_ended, + ) + + @classmethod + def _streamed_tool_use_fingerprints(cls, responses_so_far: Sequence[object]) -> tuple[str, ...]: + return tuple( + stream_item_fingerprint(block) + for item in responses_so_far + for event in cls._iter_sse_events(item) + if event.get("type") == "content_block_start" + and isinstance(block := event.get("content_block"), Mapping) + and block.get("type") == "tool_use" + ) + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Parse streaming responses and extract accumulated text content. diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 220fcedb0f8..b28daf73bc4 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -35,6 +35,22 @@ class StreamTransformSink: holdback_per_choice: dict[int, int] = field(default_factory=dict) +@dataclass(frozen=True, slots=True) +class StreamingScanKey: + """What a streaming guardrail round would hand to ``apply_guardrail``. Two keys + compare equal when the round would scan the same content again; ``stream_ended`` + stays out of the comparison and only says whether the handler is on its + end-of-stream path, where an empty payload is still scanned today.""" + + texts: tuple[str, ...] + tool_calls: tuple[str, ...] = () + stream_ended: bool = field(default=False, compare=False) + + @property + def has_nothing_to_scan(self) -> bool: + return not self.stream_ended and not any(self.texts) and not self.tool_calls + + class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( @@ -151,6 +167,9 @@ class BaseTranslation(ABC): """ return responses_so_far + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + return None + def build_block_sse_chunks( self, exc: "ModifyResponseException", diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 9b6f9c47105..8dee262001d 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -4,6 +4,8 @@ import json from collections.abc import Callable, Iterator, Sequence from typing import Any, Final, TypeVar +from pydantic import BaseModel + from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage @@ -130,6 +132,16 @@ def stream_item_field(item: object, field: str) -> object | None: return getattr(item, field, None) +def stream_item_fingerprint(item: object) -> str: + plain: Final = item.model_dump() if isinstance(item, BaseModel) else item + return json.dumps(plain, sort_keys=True, default=str) + + +def stream_item_items(item: object, field: str) -> tuple[object, ...]: + value: Final = stream_item_field(item, field) + return tuple(value) if isinstance(value, (list, tuple)) else () + + def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]: """ ``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 96a5ed663fc..d41c8557d72 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -26,6 +26,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, + StreamingScanKey, StreamTransformSink, ) from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -39,6 +40,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( role_out_of_guardrail_scope, scoped_structured_message_indices, stream_item_field, + stream_item_fingerprint, + stream_item_items, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -503,12 +506,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """Block-only streaming path: run the guardrail so an in-flight BLOCK can terminate the stream. Text rewrites are not propagated to the client here (see ``_process_streaming_transform`` for the incremental_diff path).""" - # check if the stream has ended - has_stream_ended = False - for chunk in responses_so_far: - if chunk.choices and chunk.choices[0].finish_reason is not None: - has_stream_ended = True - break + has_stream_ended: Final = self._first_choice_has_finished(responses_so_far) if has_stream_ended: # convert to model response @@ -706,8 +704,33 @@ class OpenAIChatCompletionsHandler(BaseTranslation): indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback) } + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream)) + stream_ended: Final = self._first_choice_has_finished(responses_so_far) + return StreamingScanKey( + texts=tuple(self._combine_streaming_texts(chunks).values()), + tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (), + stream_ended=stream_ended, + ) + + @staticmethod + def _streamed_tool_call_fingerprints(responses_so_far: Sequence[object]) -> tuple[str, ...]: + return tuple( + stream_item_fingerprint(tool_call) + for chunk in responses_so_far + for choice in _stream_chunk_choices(chunk) + for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls") + ) + + @staticmethod + def _first_choice_has_finished(responses_so_far: Sequence[object]) -> bool: + first_choices: Final = tuple( + choices[0] for choices in (_stream_chunk_choices(chunk) for chunk in responses_so_far) if choices + ) + return any(stream_item_field(choice, "finish_reason") is not None for choice in first_choices) + def _combine_streaming_texts( - self, responses_so_far: list["ModelResponseStream"] + self, responses_so_far: Sequence["ModelResponseStream"] ) -> dict[tuple[int, int | None], str]: """ Combine all streaming chunks into complete text per choice. diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 5a5970fb867..a0db7aadb9e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -44,10 +44,15 @@ from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( OpenAiResponsesToChatCompletionStreamIterator, ) -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, stream_item_field, + stream_item_fingerprint, + stream_item_items, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.transformation import ( @@ -593,18 +598,55 @@ class OpenAIResponsesHandler(BaseTranslation): ) return responses_so_far - def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool: + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if the streaming has ended. """ if not responses_so_far: return False - terminal_types: Final = { - ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, - ResponsesAPIStreamEvents.RESPONSE_FAILED.value, - ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, - } - return responses_so_far[-1].get("type") in terminal_types + terminal_types: Final = frozenset( + ( + ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, + ResponsesAPIStreamEvents.RESPONSE_FAILED.value, + ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, + ) + ) + return stream_item_field(responses_so_far[-1], "type") in terminal_types + + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + if not responses_so_far or not hasattr(responses_so_far[-1], "get"): + return None + last_event: Final = responses_so_far[-1] + last_event_type: Final = stream_item_field(last_event, "type") + if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value: + return None + if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value: + return self._completed_response_scan_key(stream_item_field(last_event, "response")) + return StreamingScanKey( + texts=(self.get_streaming_string_so_far(responses_so_far),), + stream_ended=self._check_streaming_has_ended(responses_so_far), + ) + + @staticmethod + def _completed_response_scan_key(response: object) -> StreamingScanKey: + output_items: Final = stream_item_items(response, "output") + message_items: Final = tuple( + item for item in output_items if stream_item_field(item, "type") != "function_call" + ) + return StreamingScanKey( + texts=tuple( + text + for item in message_items + for part in stream_item_items(item, "content") + if isinstance(text := stream_item_field(part, "text"), str) and text + ), + tool_calls=tuple( + stream_item_fingerprint(item) + for item in output_items + if stream_item_field(item, "type") == "function_call" + ), + stream_ended=True, + ) def build_stream_error_items( self, @@ -629,7 +671,7 @@ class OpenAIResponsesHandler(BaseTranslation): ), ) - def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str: + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Get the string so far from the responses so far. @@ -641,12 +683,16 @@ class OpenAIResponsesHandler(BaseTranslation): """ keyed_events: Final = tuple( ( - (event.get("item_id"), event.get("output_index"), event.get("content_index")), - event.get("text"), - event.get("delta"), + ( + stream_item_field(event, "item_id"), + stream_item_field(event, "output_index"), + stream_item_field(event, "content_index"), + ), + stream_item_field(event, "text"), + stream_item_field(event, "delta"), ) for event in responses_so_far - if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str) + if isinstance(stream_item_field(event, "text"), str) or isinstance(stream_item_field(event, "delta"), str) ) def part_text(part_key: tuple[object, object, object]) -> str: diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 46b00829b74..c6b8df1b493 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -36,6 +36,7 @@ if TYPE_CHECKING: from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, + StreamingScanKey, ) # Call types that stream JSON-RPC events (A2A); guardrail HTTPException is emitted as in-stream error @@ -54,6 +55,9 @@ class _EndpointTranslation(Protocol): @property def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ... + @property + def get_streaming_scan_key(self) -> "Callable[[Sequence[object]], StreamingScanKey | None]": ... + @property def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... @@ -70,6 +74,12 @@ def _chunk_choices(item: object) -> Sequence[object]: return choices +def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool: + if scan_key is None: + return False + return scan_key == last_scan_key or scan_key.has_nothing_to_scan + + class _StreamTerminated(Exception): """Internal signal that the incremental transform stream has already emitted its terminal chunks (block message or in-stream error) and must stop.""" @@ -1011,6 +1021,7 @@ class UnifiedLLMGuardrails(CustomLogger): # Drives how a block terminates the stream: continue the in-progress # message (True) vs emit a standalone block message (False, buffered). chunks_yielded = False + last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round async for item in response: chunk_counter += 1 @@ -1052,6 +1063,19 @@ class UnifiedLLMGuardrails(CustomLogger): # Process chunk based on sampling rate if chunk_counter % sampling_rate == 0: + endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far) + if _is_redundant_scan(scan_key, last_scan_key): + verbose_proxy_logger.debug( + "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round", + chunk_counter, + guardrail_to_apply.guardrail_name, + ) + chunks_yielded = True + responses_yielded.append(item) + yield item + continue + verbose_proxy_logger.debug( "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s", chunk_counter, @@ -1067,8 +1091,6 @@ class UnifiedLLMGuardrails(CustomLogger): # string, permanently losing this chunk's content. original_item = copy.deepcopy(item) - endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() - try: await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, @@ -1110,6 +1132,8 @@ class UnifiedLLMGuardrails(CustomLogger): ): yield error_item return + if scan_key is not None: + last_scan_key = scan_key chunks_yielded = True responses_yielded.append(original_item) yield original_item @@ -1136,6 +1160,18 @@ class UnifiedLLMGuardrails(CustomLogger): # preserve the list, not clone every chunk (deepcopy would double # peak memory for large responses). buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None + end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far) + if _is_redundant_scan(end_scan_key, last_scan_key): + verbose_proxy_logger.debug( + "Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all", + guardrail_to_apply.guardrail_name, + ) + for buffered_item in buffered_items or (): + yield buffered_item + for pending_item in pending_end_of_stream_items: + responses_yielded.append(pending_item) + yield pending_item + return try: await endpoint_translation.process_output_streaming_response( diff --git a/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py new file mode 100644 index 00000000000..dd7e8fadcd8 --- /dev/null +++ b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py @@ -0,0 +1,35 @@ +"""Tests for litellm/llms/a2a/chat/guardrail_translation/handler.py.""" + +import json + +from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey + + +def _text_event(text: str) -> str: + return json.dumps( + { + "jsonrpc": "2.0", + "id": "req-1", + "result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": text}]}, + } + ) + + +def _status_event() -> str: + return json.dumps({"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "status-update", "status": {}}}) + + +class TestA2AGuardrailHandlerStreamingScanKey: + def test_key_joins_the_text_of_every_message_event(self): + key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hello "), _text_event("world")]) + assert key == StreamingScanKey(texts=("hello world",)) + + def test_events_without_text_leave_the_key_unchanged(self): + handler = A2AGuardrailHandler() + events = [_text_event("hello")] + assert handler.get_streaming_scan_key(events + [_status_event()]) == handler.get_streaming_scan_key(events) + + def test_unparseable_items_are_ignored(self): + key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "not json", b"bytes"]) + assert key.texts == ("hi",) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 0fe7730e91e..3044a321aa6 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -13,6 +13,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.anthropic.chat.guardrail_translation.handler import ( AnthropicMessagesHandler, ) @@ -1991,3 +1992,56 @@ class TestStructuredWriteBackKeepsToolResults: } later_blocks = [b for m in messages[tool_use_index + 1 :] for b in self._blocks(m)] assert {"type": "text", "text": "Now fetch the page."} in later_blocks + + +class TestAnthropicMessagesHandlerStreamingScanKey: + """get_streaming_scan_key mirrors what process_output_streaming_response would scan""" + + @staticmethod + def _sse(event_type, data): + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() + + def _text_delta(self, text): + return self._sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + + def test_key_is_empty_before_any_text_arrives(self): + head = self._sse("message_start", {"type": "message_start", "message": {"stop_reason": None}}) + key = AnthropicMessagesHandler().get_streaming_scan_key([head]) + assert key == StreamingScanKey(texts=("",)) + + def test_key_accumulates_text_deltas(self): + key = AnthropicMessagesHandler().get_streaming_scan_key([self._text_delta("hello "), self._text_delta("world")]) + assert key.texts == ("hello world",) + assert key.stream_ended is False + + def _stop(self, stop_reason): + return self._sse( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": stop_reason, "stop_sequence": None}, "usage": {}}, + ) + + def test_stop_without_tool_use_scans_the_same_payload(self): + handler = AnthropicMessagesHandler() + open_key = handler.get_streaming_scan_key([self._text_delta("hi")]) + ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), self._stop("end_turn")]) + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_tool_use_blocks_enter_the_key_once_the_stream_has_ended(self): + handler = AnthropicMessagesHandler() + tool_use = self._sse( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}, + }, + ) + open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use]) + ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")]) + assert open_key == StreamingScanKey(texts=("hi",)) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 7dd6065063a..cebab2512d0 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -12,6 +12,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.openai.chat.guardrail_translation.handler import ( OpenAIChatCompletionsHandler, ) @@ -1643,3 +1644,74 @@ class TestCheckStreamingHasEnded: ) ] assert handler._check_streaming_has_ended(chunks) is True + + +class TestStreamingScanKey: + """get_streaming_scan_key identifies what a sampled round would scan so the + unified hook can skip rounds that would re-scan already-cleared text""" + + @staticmethod + def _chunk(content, finish_reason=None, index=0): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + return ModelResponseStream( + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)] + ) + + def test_key_carries_accumulated_text_and_open_stream(self): + handler = OpenAIChatCompletionsHandler() + key = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")]) + assert key == StreamingScanKey(texts=("hello",)) + + def test_chunks_without_text_leave_the_key_unchanged(self): + handler = OpenAIChatCompletionsHandler() + before = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")]) + after = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo"), self._chunk(None)]) + assert after == before + + def test_finish_chunk_without_tool_calls_scans_the_same_payload(self): + handler = OpenAIChatCompletionsHandler() + open_key = handler.get_streaming_scan_key([self._chunk("hi")]) + ended_key = handler.get_streaming_scan_key([self._chunk("hi"), self._chunk(None, finish_reason="stop")]) + assert open_key.stream_ended is False + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_tool_calls_only_enter_the_key_once_the_stream_has_ended(self): + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + handler = OpenAIChatCompletionsHandler() + tool_call = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}') + ) + tool_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason=None)] + ) + open_key = handler.get_streaming_scan_key([self._chunk("hi"), tool_chunk]) + ended_key = handler.get_streaming_scan_key( + [self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")] + ) + assert open_key == StreamingScanKey(texts=("hi",)) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key + + def test_text_after_the_first_choice_finishes_still_changes_the_key(self): + handler = OpenAIChatCompletionsHandler() + first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)] + key_at_first_finish = handler.get_streaming_scan_key(first_done) + key_after_more_text = handler.get_streaming_scan_key(first_done + [self._chunk("y", index=1)]) + assert key_at_first_finish.stream_ended is True + assert key_after_more_text.stream_ended is True + assert key_after_more_text != key_at_first_finish + + def test_non_stream_items_are_ignored(self): + handler = OpenAIChatCompletionsHandler() + key = handler.get_streaming_scan_key([self._chunk("hi"), b"data: [DONE]"]) + assert key.texts == ("hi",) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index d071ef78c2d..295121167d6 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1731,3 +1731,93 @@ class TestBuildBlockSseChunks: dones = [payload for payload in payloads if payload["type"] == "response.output_item.done"] assert len(dones) == 1 assert dones[0]["item"]["content"][0]["text"] == "Blocked by policy." + + +class TestOpenAIResponsesHandlerStreamingScanKey: + """get_streaming_scan_key mirrors what process_output_streaming_response would scan""" + + @staticmethod + def _delta(sequence_number, text): + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + + def test_no_events_yields_no_key(self): + assert OpenAIResponsesHandler().get_streaming_scan_key([]) is None + + def test_key_accumulates_deltas_while_the_stream_is_open(self): + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey + + key = OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hel"), self._delta(1, "lo")]) + assert key == StreamingScanKey(texts=("hello",)) + + def test_typed_delta_events_accumulate_like_dicts(self): + from litellm.types.llms.openai import OutputTextDeltaEvent + + events = [ + OutputTextDeltaEvent( + type="response.output_text.delta", + item_id="msg_1", + output_index=0, + content_index=0, + delta=text, + sequence_number=i, + ) + for i, text in enumerate(("hel", "lo")) + ] + key = OpenAIResponsesHandler().get_streaming_scan_key(events) + assert key.texts == ("hello",) + assert key.stream_ended is False + + def test_events_without_text_leave_the_key_unchanged(self): + handler = OpenAIResponsesHandler() + events = [self._delta(0, "hi")] + quiet = events + [{"type": "response.in_progress", "sequence_number": 1}] + assert handler.get_streaming_scan_key(quiet) == handler.get_streaming_scan_key(events) + + @staticmethod + def _completed(sequence_number, output): + return {"type": "response.completed", "sequence_number": sequence_number, "response": {"output": output}} + + def test_completed_event_keys_on_the_final_output_text(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + open_key = handler.get_streaming_scan_key([self._delta(0, "hi")]) + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message])]) + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_completed_event_with_a_function_call_changes_the_key(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"} + open_key = handler.get_streaming_scan_key([self._delta(0, "hi")]) + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message, function_call])]) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key + + def test_completed_event_reads_every_output_text_part(self): + from litellm.types.responses.main import GenericResponseOutputItem, OutputText + + item = GenericResponseOutputItem( + type="message", + id="msg_1", + status="completed", + role="assistant", + content=[ + OutputText(type="output_text", text="one", annotations=[]), + OutputText(type="output_text", text="two", annotations=[]), + ], + ) + key = OpenAIResponsesHandler().get_streaming_scan_key([self._completed(0, [item])]) + assert key.texts == ("one", "two") + + def test_output_item_done_round_is_never_deduped(self): + done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}} + assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index 476d443d8d8..cb6772977ec 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -310,8 +310,8 @@ def _make_stream_chunk(content: str, finish_reason=None): @pytest.mark.asyncio async def test_openai_moderation_streaming_default_uses_sampled_cadence(): """Default config samples every 5th streamed chunk and runs a final aggregate - pass after the stream ends. 10 chunks → sampled at chunks 5 and 10 → 2 in-stream - calls, plus 1 final = 3 total. + pass after the stream ends. 10 chunks are sampled at 5 and 10; the end-of-stream + round is skipped because chunk 10 already scanned the full text, for 2 total calls """ import litellm @@ -370,8 +370,9 @@ async def test_openai_moderation_streaming_default_uses_sampled_cadence(): ): pass - assert patched_make_request.await_count == 3, ( - f"Expected 3 moderation calls (2 sampled at chunks 5 / 10 + 1 final), " + assert patched_make_request.await_count == 2, ( + f"Expected 2 moderation calls (2 sampled at chunks 5 / 10; " + f"the end-of-stream round is skipped because chunk 10 already scanned the full text), " f"got {patched_make_request.await_count}" ) @@ -448,7 +449,8 @@ async def test_openai_moderation_streaming_end_of_stream_only_opt_in_calls_moder @pytest.mark.asyncio async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disabled(): """With streaming_end_of_stream_only=False and streaming_sampling_rate=2, - moderation runs every 2nd chunk during the stream, plus once more at end. + moderation runs every 2nd chunk during the stream. The terminal chunk scan covers + the final aggregate, for 3 total calls """ import litellm @@ -509,9 +511,8 @@ async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disab ): pass - # 6 chunks, sampling_rate=2 → in-stream calls at chunks 2, 4, 6 (3 calls), - # plus the final aggregate pass after the stream ends (1 call) = 4 total. - assert patched_make_request.await_count == 4, ( - f"Expected 4 moderation calls (3 sampled + 1 final aggregate), " + assert patched_make_request.await_count == 3, ( + f"Expected 3 moderation calls (3 sampled; the end-of-stream round is skipped " + f"because chunk 6 already scanned the full text), " f"got {patched_make_request.await_count}" ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 523ec1a37b4..83cc9ae8bb9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -1517,7 +1517,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: @pytest.mark.asyncio async def test_streaming_default_uses_sampled_cadence(self): - """Default samples every 5th chunk + final pass: 10 chunks → calls at 5, 10, and final = 3.""" + """Default samples every 5th chunk. For 10 chunks, sampled scans at 5 and 10 + cover the full text, so the end-of-stream round is skipped and there are 2 calls + """ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -1566,8 +1568,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: ): pass - assert mock_post.await_count == 3, ( - f"Expected 3 guardrail calls (2 sampled at chunks 5 / 10 + 1 final), " + assert mock_post.await_count == 2, ( + f"Expected 2 guardrail calls (2 sampled at chunks 5 / 10; " + f"the end-of-stream round is skipped because chunk 10 already scanned the full text), " f"got {mock_post.await_count}" ) for call in mock_post.await_args_list: @@ -1631,7 +1634,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: @pytest.mark.asyncio async def test_streaming_sampling_rate_override(self): - """sampling_rate=2 on 6 chunks → in-stream at 2,4,6 plus final = 4 calls.""" + """sampling_rate=2 on 6 chunks. Scans at 2, 4, and 6 cover the full text, so + the end-of-stream round is skipped and there are 3 calls + """ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -1680,8 +1685,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: ): pass - assert mock_post.await_count == 4, ( - f"Expected 4 guardrail calls (3 sampled + 1 final aggregate), " + assert mock_post.await_count == 3, ( + f"Expected 3 guardrail calls (3 sampled; the end-of-stream round is skipped " + f"because chunk 6 already scanned the full text), " f"got {mock_post.await_count}" ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8cad1c634a9..a28a2a71613 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1971,3 +1971,271 @@ class TestStreamingGuardrailInformationBucket: assert recorded[0]["guardrail_name"] == "audit-recorder" assert recorded[0]["guardrail_status"] == "success" assert request_data["metadata"]["user_api_key_user_id"] == "user-1" + + +class _ScanCountingGuardrail(CustomGuardrail): + """Pass-through guardrail that records every response-side scan payload.""" + + def __init__(self, *, sampling_rate=5, end_of_stream_only=False, buffer_until_moderated=False): + super().__init__(guardrail_name="scan-counter") + self.streaming_sampling_rate = sampling_rate + self.streaming_end_of_stream_only = end_of_stream_only + self.streaming_buffer_until_moderated = buffer_until_moderated + self.guardrail_config = {} + self.scans: tuple[dict[str, object], ...] = () + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.scans = ( + *self.scans, + { + "texts": list(inputs.get("texts") or []), + "tool_calls": list(inputs.get("tool_calls") or []), + "model": inputs.get("model"), + }, + ) + return inputs + + +def _responses_delta(sequence_number, text): + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + + +def _responses_tail(sequence_number, text): + return [ + { + "type": "response.output_text.done", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "text": text, + }, + { + "type": "response.completed", + "sequence_number": sequence_number + 1, + "response": { + "model": "gpt-5.6", + "output": [{"type": "message", "content": [{"type": "output_text", "text": text}]}], + }, + }, + ] + + +class TestStreamingScanDedup: + """A sampled round whose scan payload matches the previous round (or carries + no text yet) is skipped, so a stream is never re-scanned for output the + guardrail already cleared. Regression for LIT-6692.""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self, monkeypatch): + monkeypatch.setattr( + unified_module, + "endpoint_guardrail_translation_mappings", + load_guardrail_translation_mappings(), + ) + + @pytest.mark.asyncio + async def test_chat_terminal_chunk_on_sampled_index_is_scanned_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 3 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + @pytest.mark.asyncio + async def test_chat_round_with_unchanged_text_is_skipped(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [ + _stream_chunk("a"), + _stream_chunk("b"), + _stream_chunk("c"), + _stream_chunk(None), + _stream_chunk(None), + _stream_chunk(None), + _stream_chunk("d", finish_reason="stop"), + ] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 7 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abcd"]] + + @pytest.mark.asyncio + async def test_chat_finish_chunk_right_after_a_sampled_round_is_not_rescanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), _stream_chunk(None, finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 4 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + @pytest.mark.asyncio + async def test_chat_finish_chunk_carrying_tool_calls_is_still_scanned(self): + from litellm.types.utils import ChatCompletionDeltaToolCall, Function + + guardrail = _ScanCountingGuardrail(sampling_rate=3) + tool_call = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}') + ) + finish = ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason="tool_calls") + ] + ) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), finish] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 4 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abc"]] + assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"] + + @pytest.mark.asyncio + async def test_chat_second_choice_finishing_later_still_gets_the_end_scan(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [ + _stream_chunk("a", index=0), + _stream_chunk("x", index=1), + _stream_chunk("b", finish_reason="stop", index=0), + _stream_chunk("y", index=1), + _stream_chunk("z", finish_reason="stop", index=1), + ] + + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(guardrail.scans) == 2 + assert any("yz" in text for text in guardrail.scans[-1]["texts"]) + + @pytest.mark.asyncio + async def test_responses_completed_event_on_sampled_index_is_scanned_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(8)] + full_text = "".join(f"t{i}" for i in range(8)) + chunks = deltas + _responses_tail(8, full_text) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 10 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], [full_text]] + assert guardrail.scans[-1]["model"] == "gpt-5.6" + + @pytest.mark.asyncio + async def test_responses_completed_right_after_a_sampled_round_is_not_rescanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + chunks = deltas + _responses_tail(5, "t0t1t2t3t4") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 7 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"]] + + @pytest.mark.asyncio + async def test_responses_completed_carrying_a_function_call_is_still_scanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + completed = { + "type": "response.completed", + "sequence_number": 5, + "response": { + "model": "gpt-5.6", + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "t0t1t2t3t4"}]}, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + ], + }, + } + chunks = deltas + [completed] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 6 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], ["t0t1t2t3t4"]] + assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"] + + @pytest.mark.asyncio + async def test_responses_round_with_unchanged_text_is_skipped(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + quiet = [{"type": "response.in_progress", "sequence_number": i} for i in range(5, 10)] + chunks = deltas + quiet + _responses_tail(10, "t0t1t2t3t4") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 12 + assert guardrail.scans == ({"texts": ["t0t1t2t3t4"], "tool_calls": [], "model": None},) + + @pytest.mark.asyncio + async def test_responses_tool_call_done_event_is_still_scanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2) + tool_call_done = { + "type": "response.output_item.done", + "sequence_number": 1, + "output_index": 1, + "item": { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + } + chunks = [_responses_delta(0, "hi"), tool_call_done] + _responses_tail(2, "hi") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 4 + assert len(guardrail.scans) == 2 + assert [call["function"]["name"] for call in guardrail.scans[0]["tool_calls"]] == ["get_weather"] + assert guardrail.scans[1]["texts"] == ["hi"] + + @pytest.mark.asyncio + async def test_anthropic_skips_empty_round_and_terminal_duplicate(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2) + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]] + + @pytest.mark.asyncio + async def test_end_of_stream_only_still_scans_exactly_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2, end_of_stream_only=True) + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]] + + @pytest.mark.asyncio + async def test_buffer_until_moderated_still_scans_exactly_once_and_releases_every_chunk(self): + guardrail = _ScanCountingGuardrail(sampling_rate=1, buffer_until_moderated=True) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] From 92edcb90dbb41175f526187826f70058e5a57f94 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:27:19 -0700 Subject: [PATCH 12/23] fix: keep litellm importable on Python 3.10 and guard 3.11-only typing imports in CI (#39448) * ci: guard against Python 3.10-incompatible typing imports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: address Python 3.10 typing guard review Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ci): honor version-guard direction and scan litellm-proxy-extras in py310 typing check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-code-quality.yml | 33 ++++ .../websearch_interception/handler.py | 4 +- .../_experimental/mcp_server/tool_search.py | 4 +- litellm/proxy/agent_endpoints/endpoints.py | 4 +- .../proxy/common_utils/reset_budget_job.py | 4 +- .../batch_file_validation.py | 4 +- .../general_upload_validation.py | 4 +- .../types/llms/gemini_audio_transcription.py | 4 +- .../check_py310_typing_imports.py | 150 ++++++++++++++++++ .../test_check_py310_typing_imports.py | 86 ++++++++++ 10 files changed, 286 insertions(+), 11 deletions(-) create mode 100644 tests/code_coverage_tests/check_py310_typing_imports.py create mode 100644 tests/test_litellm/test_check_py310_typing_imports.py diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index c112bf2bb22..d02f5878396 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -128,6 +128,9 @@ jobs: - name: check_fastuuid_usage run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py + - name: check_py310_typing_imports + run: uv run --no-sync python ./tests/code_coverage_tests/check_py310_typing_imports.py + - name: check_e2e_no_raw_requests run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -145,3 +148,33 @@ jobs: - name: documentation_test_api_docs run: uv run --no-sync python ./tests/documentation_tests/test_api_docs.py + + python-310-import-smoke: + runs-on: ubuntu-latest + timeout-minutes: 15 + + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.10" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Install dependencies + run: uv sync --frozen --extra proxy --python 3.10 + + - run: uv run --no-sync python --version + + - name: Import litellm + run: uv run --no-sync python -c "import litellm" + + - name: Check litellm CLI + run: uv run --no-sync litellm --version diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 2d737bc34e7..587da997f94 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -10,9 +10,9 @@ import asyncio import math import uuid from collections.abc import AsyncIterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast -from typing_extensions import ReadOnly +from typing_extensions import Never, ReadOnly import litellm from litellm._logging import verbose_logger diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 4f6305d88cf..af02c11ad86 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -5,10 +5,10 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, assert_never +from typing import TYPE_CHECKING, Any, Final, TypedDict from pydantic import ValidationError -from typing_extensions import ReadOnly, Required +from typing_extensions import ReadOnly, Required, assert_never import litellm from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 3e4dc07a521..3b8151d1064 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -13,10 +13,10 @@ import os import uuid from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Annotated, Final, TypedDict, assert_never +from typing import Annotated, Final, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query, Request -from typing_extensions import ReadOnly, Required +from typing_extensions import ReadOnly, Required, assert_never import litellm from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 1682cf12f4e..47f69732e95 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -7,7 +7,9 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum from types import MappingProxyType -from typing import Final, Literal, Protocol, TypeVar, assert_never +from typing import Final, Literal, Protocol, TypeVar + +from typing_extensions import assert_never import litellm from litellm._logging import verbose_proxy_logger diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index 0aee5e8cc54..a41bd36d510 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -2,7 +2,9 @@ import json from collections.abc import Iterator from dataclasses import dataclass from itertools import chain -from typing import BinaryIO, Final, NoReturn, assert_never +from typing import BinaryIO, Final, NoReturn + +from typing_extensions import assert_never from litellm.proxy._types import ProxyException diff --git a/litellm/proxy/openai_files_endpoints/general_upload_validation.py b/litellm/proxy/openai_files_endpoints/general_upload_validation.py index 9d450cb5b8d..8c59a520272 100644 --- a/litellm/proxy/openai_files_endpoints/general_upload_validation.py +++ b/litellm/proxy/openai_files_endpoints/general_upload_validation.py @@ -8,7 +8,9 @@ extensions, path-traversal filenames) regardless of purpose. from dataclasses import dataclass from pathlib import Path -from typing import BinaryIO, Final, NoReturn, assert_never +from typing import BinaryIO, Final, NoReturn + +from typing_extensions import assert_never from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.path_utils import safe_filename diff --git a/litellm/types/llms/gemini_audio_transcription.py b/litellm/types/llms/gemini_audio_transcription.py index cb12e0f45b8..f7e74ba4bf8 100644 --- a/litellm/types/llms/gemini_audio_transcription.py +++ b/litellm/types/llms/gemini_audio_transcription.py @@ -1,7 +1,7 @@ -from typing import Literal, Required +from typing import Literal from pydantic import BaseModel, ConfigDict -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict class GeminiTranscriptionAudioInput(TypedDict): diff --git a/tests/code_coverage_tests/check_py310_typing_imports.py b/tests/code_coverage_tests/check_py310_typing_imports.py new file mode 100644 index 00000000000..0cd4d089890 --- /dev/null +++ b/tests/code_coverage_tests/check_py310_typing_imports.py @@ -0,0 +1,150 @@ +import ast +import os +import sys +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +PY311_PLUS_TYPING_NAMES: Final[frozenset[str]] = frozenset( + { + "NotRequired", + "Required", + "Self", + "LiteralString", + "Never", + "assert_never", + "assert_type", + "reveal_type", + "TypeVarTuple", + "Unpack", + "dataclass_transform", + "override", + "TypeAliasType", + "get_original_bases", + "ReadOnly", + "TypeIs", + "NoDefault", + "get_protocol_members", + "is_protocol", + "evaluate_forward_ref", + "TypeForm", + } +) + + +@dataclass(frozen=True, slots=True) +class TypingImportViolation: + file: str + line: int + name: str + + +def _walk_with_ancestors( + node: ast.AST, ancestors: tuple[tuple[ast.AST, str], ...] = () +) -> Iterator[tuple[ast.AST, tuple[tuple[ast.AST, str], ...]]]: + yield node, ancestors + for field_name, field_value in ast.iter_fields(node): + if isinstance(field_value, ast.AST): + yield from _walk_with_ancestors(field_value, (*ancestors, (node, field_name))) + elif isinstance(field_value, list): + for child in field_value: + if isinstance(child, ast.AST): + yield from _walk_with_ancestors(child, (*ancestors, (node, field_name))) + + +def _is_sys_version_info(node: ast.AST) -> bool: + return ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "sys" + and node.attr == "version_info" + ) + + +def _is_version_guarded(ancestors: tuple[tuple[ast.AST, str], ...]) -> bool: + nearest_if: Final[tuple[ast.If, str] | None] = next( + ( + (ancestor, field_name) + for ancestor, field_name in reversed(ancestors) + if isinstance(ancestor, ast.If) + ), + None, + ) + if nearest_if is None: + return False + enclosing_if, branch = nearest_if + test: Final[ast.expr] = enclosing_if.test + if not isinstance(test, ast.Compare) or len(test.ops) != 1 or not _is_sys_version_info(test.left): + return False + operator: Final[ast.cmpop] = test.ops[0] + return (isinstance(operator, (ast.Gt, ast.GtE)) and branch == "body") or ( + isinstance(operator, (ast.Lt, ast.LtE)) and branch == "orelse" + ) + + +def scan_file(file_path: str | os.PathLike[str]) -> tuple[TypingImportViolation, ...]: + path: Final[Path] = Path(file_path) + tree: Final[ast.Module] = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + return tuple( + violation + for node, ancestors in _walk_with_ancestors(tree) + if not _is_version_guarded(ancestors) + for violation in _violations_for_node(node, path) + ) + + +def _violations_for_node( + node: ast.AST, path: Path +) -> tuple[TypingImportViolation, ...]: + if isinstance(node, ast.ImportFrom) and node.module == "typing": + return tuple( + TypingImportViolation(file=str(path), line=node.lineno, name=alias.name) + for alias in node.names + if alias.name in PY311_PLUS_TYPING_NAMES + ) + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "typing" + and node.attr in PY311_PLUS_TYPING_NAMES + ): + return (TypingImportViolation(file=str(path), line=node.lineno, name=node.attr),) + return () + + +def scan_directory(base_dir: str | os.PathLike[str] = ".") -> tuple[TypingImportViolation, ...]: + base_path: Final[Path] = Path(base_dir) + return tuple( + violation + for directory in ( + base_path / "litellm", + base_path / "enterprise", + base_path / "litellm-proxy-extras" / "litellm_proxy_extras", + ) + if directory.exists() + for path in directory.rglob("*.py") + for violation in scan_file(path) + ) + + +def main() -> None: + violations: Final[tuple[TypingImportViolation, ...]] = scan_directory() + if violations: + message: Final[str] = "\n".join( + ( + "Python 3.10-incompatible typing imports found:", + *( + f"{violation.file}:{violation.line}: {violation.name} is unavailable in Python 3.10; " + "import it from typing_extensions instead because litellm supports Python 3.10" + for violation in violations + ), + ) + ) + sys.stdout.write(f"{message}\n") + raise RuntimeError("Import Python 3.10-incompatible typing names from typing_extensions instead") + sys.stdout.write("No Python 3.10-incompatible typing imports found.\n") + + +if __name__ == "__main__": + main() diff --git a/tests/test_litellm/test_check_py310_typing_imports.py b/tests/test_litellm/test_check_py310_typing_imports.py new file mode 100644 index 00000000000..de370326091 --- /dev/null +++ b/tests/test_litellm/test_check_py310_typing_imports.py @@ -0,0 +1,86 @@ +import sys +from pathlib import Path +from typing import Final + +_CODE_COVERAGE_DIR: Final[Path] = Path(__file__).resolve().parents[1] / "code_coverage_tests" +sys.path.insert(0, str(_CODE_COVERAGE_DIR)) # test-quality-ok: required to import checker from its source directory +import check_py310_typing_imports as checker # noqa: E402 # load checker from its source directory + + +def _scan(tmp_path: Path, source: str) -> tuple[object, ...]: + file_path = tmp_path / "fixture.py" + file_path.write_text(source, encoding="utf-8") + return checker.scan_file(file_path) + + +def test_typing_import_flags_python_311_name(tmp_path: Path) -> None: + violations = _scan(tmp_path, "from typing import NotRequired, TypedDict\n") + assert tuple(violation.name for violation in violations) == ("NotRequired",) + + +def test_typing_extensions_import_passes(tmp_path: Path) -> None: + assert _scan(tmp_path, "from typing_extensions import NotRequired\n") == () + + +def test_typing_attribute_flags_python_311_name(tmp_path: Path) -> None: + violations = _scan(tmp_path, "import typing\nx: typing.Self\n") + assert tuple(violation.name for violation in violations) == ("Self",) + + +def test_version_guarded_typing_import_passes(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info >= (3, 11):\n" + " from typing import NotRequired\n" + "else:\n" + " from typing_extensions import NotRequired\n" + ) + assert _scan(tmp_path, source) == () + + +def test_python_310_branch_flags_typing_import(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info >= (3, 11):\n" + " from typing_extensions import NotRequired\n" + "else:\n" + " from typing import NotRequired\n" + ) + violations = _scan(tmp_path, source) + assert tuple(violation.name for violation in violations) == ("NotRequired",) + + +def test_python_310_branch_is_exempt_for_less_than_guard(tmp_path: Path) -> None: + source = ( + "import sys\n" + "if sys.version_info < (3, 11):\n" + " from typing_extensions import NotRequired\n" + "else:\n" + " from typing import NotRequired\n" + ) + assert _scan(tmp_path, source) == () + + +def test_nearest_if_controls_version_guard(tmp_path: Path) -> None: + source = ( + "if sys.version_info >= (3, 11):\n" + " from typing import Self\n" + " x = 1\n" + "if True:\n" + " from typing import Self\n" + ) + violations = _scan(tmp_path, source) + assert tuple((violation.name, violation.line) for violation in violations) == (("Self", 5),) + + +def test_scan_directory_includes_proxy_extras(tmp_path: Path) -> None: + file_path = tmp_path / "litellm-proxy-extras" / "litellm_proxy_extras" / "m.py" + file_path.parent.mkdir(parents=True) + file_path.write_text("from typing import NotRequired\n", encoding="utf-8") + + violations = checker.scan_directory(tmp_path) + assert tuple((violation.name, violation.file) for violation in violations) == (("NotRequired", str(file_path)),) + + +def test_python_310_typing_name_passes(tmp_path: Path) -> None: + assert _scan(tmp_path, "from typing import Optional\n") == () From 8441dd6e8ca16a480e0766651f7c4ab99d22d5d5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:28:06 -0700 Subject: [PATCH 13/23] fix(proxy): keep SpendLogs and callback session ids in sync when the request has none (#39450) * fix(proxy): keep SpendLogs and callback session ids in sync when the request has none Add general_settings.missing_session_id (generate | reject). In generate mode one id is stamped into litellm_session_id, litellm_trace_id and metadata.session_id before callbacks run, so LiteLLM_SpendLogs.session_id and the Langfuse session id match. In reject mode such requests get a 400. Unset keeps the legacy behavior. MCP routes are not affected Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(proxy): regenerate schema.d.ts and shorten mutable-ok comment for ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): mark generated session ids so affinity consumers do not pin on them Fireworks x-session-affinity, the router session_affinity pre-call check and the complexity router session pin all read metadata.session_id as a caller-chosen stable key. A missing_session_id: generate id is fresh per request, so it now carries metadata.litellm_session_id_generated and those consumers skip it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/llms/fireworks_ai/common_utils.py | 11 +- litellm/proxy/_types.py | 4 + litellm/proxy/litellm_pre_call_utils.py | 55 ++++++ .../complexity_router/complexity_router.py | 8 +- .../deployment_affinity_check.py | 4 +- .../test_fireworks_ai_chat_transformation.py | 16 ++ .../proxy/test_litellm_pre_call_utils.py | 176 ++++++++++++++++++ .../router_strategy/test_complexity_router.py | 22 ++- .../test_session_id_affinity.py | 43 ++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 11 files changed, 335 insertions(+), 10 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index c7b74e176db..ef9329b9dfc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1449,6 +1449,7 @@ RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" +SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 8c306faa036..ac934ad0cb5 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -2,6 +2,7 @@ from typing import Final from httpx import Headers +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -16,16 +17,18 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None: """ Session id to send as `x-session-affinity`, or None when the caller gave none. - Deliberately does not fall back to `litellm_trace_id`: that is generated per - request (`str(uuid.uuid4())` when absent), so using it pins every request to a - different Fireworks node and prompt caching never hits. + Deliberately does not fall back to `litellm_trace_id`, and ignores session ids the + proxy generated for a request that had none: both are per request, so using them + pins every request to a different Fireworks node and prompt caching never hits. """ params: Final = litellm_params + metadata: Final = params.get("metadata") + if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY): + return None for key in ("litellm_session_id", "session_id"): value = params.get(key) if value: return str(value) - metadata: Final = params.get("metadata") if isinstance(metadata, dict): value = metadata.get("session_id") if value: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2da7ceb2d50..849e54c65aa 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2594,6 +2594,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.", ) + missing_session_id: Literal["generate", "reject"] | None = Field( + None, + description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.", + ) enable_public_model_hub: bool = Field( default=False, description="Public model hub for users to see what models they have access to, supported openai params, etc.", diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 20f83085286..1d440448c2f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -16,6 +16,7 @@ from starlette.datastructures import Headers import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging +from litellm._uuid import uuid from litellm.constants import ( CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -23,6 +24,7 @@ from litellm.constants import ( OTEL_SERVICE_NAME_METADATA_KEYS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -40,6 +42,7 @@ from litellm.proxy._types import ( AddTeamCallback, CommonProxyErrors, LitellmDataForBackendLLMCall, + LiteLLMRoutes, LitellmUserRoles, ProxyErrorTypes, ProxyException, @@ -47,6 +50,8 @@ from litellm.proxy._types import ( TeamCallbackMetadata, UserAPIKeyAuth, ) +from litellm.proxy.auth.auth_utils import get_request_route +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, get_metadata_variable_name_from_kwargs, @@ -715,6 +720,50 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None: return session_id +def _is_llm_inference_route(request: Request) -> bool: + route: Final = get_request_route(request) + return RouteChecks.is_llm_api_route(route=route) and not RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value + ) + + +def apply_missing_session_id_policy( + data: dict[str, object], # mutable-ok: stamps session ids in place on the request body the pipeline threads through + _metadata_variable_name: str, + general_settings: Mapping[str, object] | None, + request: Request, +) -> None: + policy: Final = general_settings.get("missing_session_id") if general_settings else None + if policy is None or not _is_llm_inference_route(request): + return + metadata: Final = data.get(_metadata_variable_name) + if not isinstance(metadata, dict): + return + if data.get("litellm_session_id") or metadata.get("session_id"): + return + match policy: + case "generate": + session_id: Final = str(data.get("litellm_trace_id") or metadata.get("trace_id") or uuid.uuid4()) + data["litellm_session_id"] = session_id # rebind-ok: data is an out-param + data.setdefault("litellm_trace_id", session_id) + metadata["session_id"] = session_id + metadata[SESSION_ID_GENERATED_METADATA_KEY] = True + case "reject": + raise ProxyException( + message=( + "Request has no session id. Send an `x-litellm-session-id` header or `metadata.session_id`. " + "Required by `general_settings.missing_session_id: reject`." + ), + type=ProxyErrorTypes.bad_request_error, + param="session_id", + code=400, + ) + case _: + verbose_proxy_logger.warning( + "Ignoring unknown general_settings.missing_session_id=%r; expected 'generate' or 'reject'", policy + ) + + def is_claude_code_user_agent(user_agent: str) -> bool: """Claude Code identifies itself as ``claude-cli/ ...``; the IDE extensions and the Agent SDK run through the same CLI and share that prefix.""" @@ -1818,6 +1867,12 @@ async def add_litellm_data_to_request( data=data, _metadata_variable_name=_metadata_variable_name, ) + apply_missing_session_id_policy( + data=data, + _metadata_variable_name=_metadata_variable_name, + general_settings=general_settings, + request=request, + ) # Expose request headers under the metadata field for guardrails (fixes #17477) if _metadata_variable_name in data and isinstance(data[_metadata_variable_name], dict): diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index d205db90607..430efe339a2 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -26,7 +26,11 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from pydantic import BaseModel, create_model from litellm._logging import verbose_router_logger -from litellm.constants import EMPTY_MAPPING, RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.constants import ( + EMPTY_MAPPING, + RETURN_RAW_MODEL_NAME_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, +) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata @@ -2712,7 +2716,7 @@ class ComplexityRouter(CustomLogger): """Resolve a client-supplied session_id.""" for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs): session_id = metadata.get("session_id") - if session_id is not None: + if session_id is not None and not metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return str(session_id) return None diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index b1e9dbdefa8..39d3e25aacb 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -21,7 +21,7 @@ from typing_extensions import TypedDict from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues @@ -265,7 +265,7 @@ class DeploymentAffinityCheck(CustomLogger): @staticmethod def _get_session_id_from_metadata_dict(metadata: dict) -> str | None: session_id: Final = metadata.get("session_id") - if session_id is None: + if session_id is None or metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return None return str(session_id) diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index d7cc89868af..ec8725db5f7 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -8,6 +8,7 @@ import litellm from litellm import get_model_info, supports_reasoning, supports_vision from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import ( ChatCompletionMessageToolCall, @@ -235,6 +236,21 @@ def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id(): ) +def test_get_fireworks_session_id_ignores_proxy_generated_session_id(): + """general_settings.missing_session_id: generate stamps a fresh id per request; sending it + as x-session-affinity would pin every request to a different node.""" + assert ( + get_fireworks_session_id( + { + "litellm_session_id": "generated-1", + "litellm_trace_id": "generated-1", + "metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}, + } + ) + is None + ) + + def test_handle_message_content_with_tool_calls(): config = FireworksAIConfig() message = Message( diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index ee0e2014951..8366e5546a9 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -41,7 +41,9 @@ from litellm.litellm_core_utils.get_provider_specific_headers import ( from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( TRUSTED_CALLBACK_VARS_FIELD, ) +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import CredentialItem @@ -7719,3 +7721,177 @@ def test_stamped_model_access_groups_survive_the_litellm_metadata_merge(): } assert get_litellm_metadata_from_kwargs(kwargs)[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] + + +def _request_for(path: str) -> MagicMock: + request = MagicMock(spec=Request) + request.scope = {"path": path} + request.url = MagicMock() + request.url.path = path + request.url.__str__.return_value = f"http://localhost{path}" + request.method = "POST" + request.query_params = {} + request.headers = {"Content-Type": "application/json"} + request.client = MagicMock() + request.client.host = "127.0.0.1" + return request + + +def _spend_log_session_id(data: dict[str, object]) -> str: + """Resolve session_id the way LiteLLM_SpendLogs does: standard_logging_payload.trace_id.""" + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_session_id_for_spend_log + + metadata = data["metadata"] + assert isinstance(metadata, dict) + litellm_params = get_litellm_params( + litellm_session_id=str(data["litellm_session_id"]) if "litellm_session_id" in data else None, + litellm_trace_id=str(data["litellm_trace_id"]) if "litellm_trace_id" in data else None, + metadata=metadata, + ) + trace_id = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=SimpleNamespace(litellm_trace_id="per-call-random-trace-id"), + litellm_params=litellm_params, + ) + return _get_session_id_for_spend_log(kwargs={}, standard_logging_payload={"trace_id": trace_id}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_correlation_in_logs", [False, True]) +async def test_missing_session_id_generate_makes_spend_log_and_callback_session_ids_agree( + monkeypatch: pytest.MonkeyPatch, request_correlation_in_logs: bool +): + """Without a session header, SpendLogs.session_id and the metadata.session_id that Langfuse logs + must be the same generated id, so cross-referencing the two by session_id works. The id is marked + as generated so affinity consumers (Fireworks x-session-affinity, router session pins) skip it.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", request_correlation_in_logs) + data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]} + + updated = await add_litellm_data_to_request( + data=data, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "generate"}, + ) + + callback_session_id = updated["metadata"]["session_id"] + assert isinstance(callback_session_id, str) and len(callback_session_id) == 36 + assert _spend_log_session_id(updated) == callback_session_id + assert updated["metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True + assert get_fireworks_session_id( + {"litellm_session_id": updated["litellm_session_id"], "metadata": updated["metadata"]} + ) is None + + +@pytest.mark.asyncio +async def test_missing_session_id_unset_keeps_legacy_divergence(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + ) + + assert "session_id" not in updated["metadata"] + assert "litellm_session_id" not in updated + assert _spend_log_session_id(updated) == "per-call-random-trace-id" + + +@pytest.mark.asyncio +async def test_missing_session_id_generate_reuses_traceparent_trace_id(): + """A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it.""" + request = _request_for("/v1/chat/completions") + request.headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "generate"}, + ) + + assert updated["metadata"]["session_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert _spend_log_session_id(updated) == "4bf92f3577b34da6a3ce929d0e0e4736" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", ["generate", "reject"]) +async def test_missing_session_id_policy_keeps_client_supplied_session_id(policy: str): + request = _request_for("/v1/chat/completions") + request.headers = {"x-litellm-session-id": "client-session-1"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": policy}, + ) + + assert updated["litellm_session_id"] == "client-session-1" + assert updated["metadata"]["session_id"] == "client-session-1" + assert _spend_log_session_id(updated) == "client-session-1" + assert SESSION_ID_GENERATED_METADATA_KEY not in updated["metadata"] + assert ( + get_fireworks_session_id({"litellm_session_id": "client-session-1", "metadata": updated["metadata"]}) + == "client-session-1" + ) + + +@pytest.mark.asyncio +async def test_missing_session_id_reject_accepts_body_metadata_session_id(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [], "metadata": {"session_id": "body-session-1"}}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert updated["metadata"]["session_id"] == "body-session-1" + + +@pytest.mark.asyncio +async def test_missing_session_id_reject_returns_400_without_session_id(): + with pytest.raises(ProxyException) as exc_info: + await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert exc_info.value.code == "400" + assert exc_info.value.param == "session_id" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/mcp/", "/mcp/tools", "/key/health"]) +async def test_missing_session_id_policy_skips_non_inference_routes(path: str): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o"}, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "reject"}, + ) + + assert "session_id" not in updated["metadata"] + + +@pytest.mark.asyncio +async def test_missing_session_id_unknown_value_is_ignored(): + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=_request_for("/v1/chat/completions"), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "typo"}, + ) + + assert "session_id" not in updated["metadata"] diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 941d78085e4..d7d02544efb 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -16,7 +16,7 @@ import litellm from litellm import Router from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -4274,6 +4274,26 @@ class TestSessionAffinity: assert first.model == "o1-preview" assert second.model == "gpt-4o-mini" + @pytest.mark.asyncio + async def test_proxy_generated_session_id_never_pins(self, mock_router_instance, session_affinity_config): + """A session id the proxy generated for a request that had none is per request, so + it must not create a pin even with session_affinity enabled.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + request_kwargs = {"metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}} + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.model == "o1-preview" + assert second.model == "gpt-4o-mini" + @pytest.mark.asyncio async def test_can_be_enabled_to_pin_every_later_turn(self, mock_router_instance, session_affinity_config): """Regression: session_affinity=True is the opt-in, so a shared session_id reuses the diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py index a3772a276fa..cf48888600e 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py @@ -7,7 +7,7 @@ import json import litellm from litellm.caching.dual_cache import DualCache -from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -180,6 +180,47 @@ async def test_async_session_id_affinity_priority_over_user_key(): assert filtered[0]["model_info"]["id"] == "deployment-2" +@pytest.mark.asyncio +async def test_proxy_generated_session_id_does_not_pin_a_deployment(): + """A session id the proxy generated for a request that had none is per request, so a + pin stored under it must be ignored and none must be written.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=123, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + ) + healthy_deployments = [ + {"model_name": "model_group", "litellm_params": {"model": "model_1"}, "model_info": {"id": "deployment-1"}}, + {"model_name": "model_group", "litellm_params": {"model": "model_2"}, "model_info": {"id": "deployment-2"}}, + ] + await cache.async_set_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1"), + {"model_id": "deployment-2"}, + ) + request_kwargs = { + "metadata": {"user_api_key_hash": "user1", "session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True} + } + + filtered = await callback.async_filter_deployments( + model="model_group", healthy_deployments=healthy_deployments, messages=[], request_kwargs=request_kwargs + ) + await callback.async_pre_call_deployment_hook( + kwargs={ + "metadata": {**request_kwargs["metadata"], "deployment_model_name": "model_group"}, + "model_info": {"id": "deployment-1"}, + }, + call_type=None, + ) + + assert len(filtered) == 2 + assert await cache.async_get_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1") + ) == {"model_id": "deployment-2"} + + MOCK_RESPONSES_API_RESPONSE = { "id": "resp_mock-resp-456", "object": "response", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 90c4f03bf08..491c4fb6a44 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25772,6 +25772,11 @@ export interface components { * @description Number of trusted reverse proxies/load balancers in front of the gateway that append to X-Forwarded-For. When set (and mcp_trusted_proxy_ranges validates the direct peer), the client IP for MCP access control is read this many entries from the right of the chain instead of the spoofable leftmost value, defeating append-style X-Forwarded-For forgery. */ mcp_xff_num_trusted_hops?: number | null; + /** + * Missing Session Id + * @description What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id. + */ + missing_session_id?: ("generate" | "reject") | null; /** * Model List Healthy Only * @description When true, `/models`, `/v1/models/{id}` and `/model/info` hide models whose backing deployments are all unhealthy, for every caller, without needing `healthy_only=true` per request. Requires `background_health_checks: true`, and keeps deployment health state cached without turning on `enable_health_check_routing`, so routing is unaffected. With no health state nothing is hidden. Hiding is presentation-only, a hidden model can still be called. From 78ff5ac9cd144770f9e53931d190ca2d480d9d98 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 2 Sep 2026 18:31:25 -0700 Subject: [PATCH 14/23] feat(router): arm safeguard-refusal fallback on generic chains when no content-policy list exists (#39274) --- litellm/router.py | 30 ++++- .../router_utils/fallback_event_handlers.py | 32 ++++++ ...test_router_anthropic_messages_fallback.py | 106 ++++++++++++++++++ 3 files changed, 166 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 0af514fe8a2..303b22c9484 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -150,8 +150,10 @@ from litellm.router_utils.fallback_event_handlers import ( _check_non_standard_fallback_format, clear_pre_routing_selection, fallback_lookup_groups, + fallbacks_disabled_for_request, get_fallback_model_group_for_lookup_groups, get_pre_routing_selection, + record_disable_fallbacks, record_pre_routing_selection, run_async_fallback, ) @@ -5193,7 +5195,7 @@ class Router: if not has_generated_content and error_event is None else None ) - if refusal_stop_details is not None and self._has_content_policy_fallback(model, initial_kwargs): + if refusal_stop_details is not None and self._refusal_fallback_available(model, initial_kwargs): refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details) raise MidStreamFallbackError( message=refusal_error.message, @@ -7266,6 +7268,7 @@ class Router: _fallback_metadata["original_model_group"] = model_group include_fallback_errors: Final = kwargs.get("include_fallback_errors", False) is True disable_fallbacks: Final[bool | None] = kwargs.pop("disable_fallbacks", False) + record_disable_fallbacks(kwargs, disable_fallbacks is True) fallbacks: Final[list | None] = kwargs.get("fallbacks", self.fallbacks) context_window_fallbacks: list | None = kwargs.get("context_window_fallbacks", self.context_window_fallbacks) content_policy_fallbacks: list | None = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks) @@ -8131,6 +8134,29 @@ class Router: ) return False + def _refusal_fallback_available(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: + """ + Whether a safeguard refusal can actually be recovered by the dispatcher. A configured + content-policy list is authoritative; with none configured at all, the dispatcher falls + through to the generic fallbacks lookup, so the gate mirrors that reachability and arms + on a resolving generic chain (tier first, then the requested group, then "*"). + """ + if fallbacks_disabled_for_request(kwargs): + return False + content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks) + if content_policy_fallbacks is not None: + return self._has_content_policy_fallback(model_group, kwargs) + if self._has_default_fallbacks(): + return True + fallbacks: Final = kwargs.get("fallbacks", self.fallbacks) + if fallbacks is None: + return False + resolved, _ = get_fallback_model_group_for_lookup_groups( + fallbacks=fallbacks, + lookup_groups=fallback_lookup_groups(kwargs, model_group), + ) + return resolved is not None + def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool: """ Determines if a content policy error should be raised. @@ -8162,7 +8188,7 @@ class Router: return False if get_safeguard_refusal_stop_details(response) is None: return False - return self._has_content_policy_fallback(model, kwargs) + return self._refusal_fallback_available(model, kwargs) def _get_healthy_deployments(self, model: str, parent_otel_span: Span | None): _all_deployments: list = [] diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index f7855cb38ff..601b32c4386 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -263,6 +263,38 @@ def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None: return next((selected for selected in selections if isinstance(selected, str) and selected), None) +DISABLE_FALLBACKS_METADATA_KEY: Final = "_disable_fallbacks" + + +def record_disable_fallbacks(request_kwargs: Mapping[str, Any] | None, disabled: bool) -> None: + """ + Write-or-clear the request's disable_fallbacks verdict into the router-internal metadata + bucket. The wrapper pops the raw kwarg before any downstream frame runs, so the refusal + gate (which decides whether to convert a refusal into a recoverable error) needs this + carrier to know recovery is impossible. + """ + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + + if request_kwargs is None: + return + bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)) + if not isinstance(bucket, dict): + return + if disabled: + bucket[DISABLE_FALLBACKS_METADATA_KEY] = True + else: + bucket.pop(DISABLE_FALLBACKS_METADATA_KEY, None) + + +def fallbacks_disabled_for_request(kwargs: Mapping[str, Any]) -> bool: + """True when this request opted out of fallbacks, read from the raw kwarg (pre-pop + snapshots keep it) or the router-internal bucket the wrapper stamps after popping it.""" + if kwargs.get("disable_fallbacks") is True: + return True + buckets: Final = (kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS) + return any(isinstance(bucket, dict) and bucket.get(DISABLE_FALLBACKS_METADATA_KEY) is True for bucket in buckets) + + def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]: """ Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins, diff --git a/tests/router_unit_tests/test_router_anthropic_messages_fallback.py b/tests/router_unit_tests/test_router_anthropic_messages_fallback.py index 0c4d1dfc21e..4812d199c06 100644 --- a/tests/router_unit_tests/test_router_anthropic_messages_fallback.py +++ b/tests/router_unit_tests/test_router_anthropic_messages_fallback.py @@ -338,6 +338,112 @@ def test_record_pre_routing_selection_writes_only_the_internal_bucket(): assert kwargs["metadata"] == {"user_id": "u1"} +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_generic_only_row_recovers_safeguard_refusal(stream): + """With no content-policy list configured, a generic fallback row covers safeguard refusals, + so the dashboard's generic fallbacks work without config-only content_policy rows.""" + fake = FakeAnthropicUpstream() + router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + response = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, stream=stream, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(response) if stream else response + + if stream: + assert b'"refusal"' not in body + assert b"text_delta" in body + else: + assert body["stop_reason"] == "end_turn" + assert len(fake.calls) == 2 + assert "claude-opus-5" in fake.calls[1] + + +@pytest.mark.asyncio +async def test_configured_content_policy_list_stays_authoritative_over_generic_rows(): + fake = FakeAnthropicUpstream() + router = Router( + model_list=[FABLE_TIER, OPUS_TARGET], + fallbacks=[{"fable-tier": ["opus-target"]}], + content_policy_fallbacks=[{"unrelated-group": ["opus-target"]}], + ) + + with fake.install(): + response = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}] + ) + + assert response["stop_reason"] == "refusal" + assert len(fake.calls) == 1 + + +def test_refusal_fallback_available_arms_on_generic_rows_only_without_content_policy(): + router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"tier-group": ["opus-target"]}]) + stamped = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}} + + assert router._refusal_fallback_available("router-group", stamped) is True + assert router._refusal_fallback_available("router-group", {}) is False + assert router._refusal_fallback_available("router-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False + + +def test_chat_content_filter_gate_unchanged_by_generic_rows(): + """The generic-row arming is scoped to /v1/messages safeguard refusals; the chat surface's + content_filter gate keeps its long-standing content-policy-only semantics.""" + from litellm.types.utils import Choices, ModelResponse + + router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}]) + response = ModelResponse(choices=[Choices(finish_reason="content_filter")]) + + assert router._should_raise_content_policy_error(model="fable-tier", response=response, kwargs={}) is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_disable_fallbacks_returns_the_refusal_instead_of_raising(stream): + """A request that opted out of fallbacks must receive the provider's refusal response, + never a ContentPolicyViolationError the dispatcher refuses to recover.""" + fake = FakeAnthropicUpstream() + router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + response = await router.aanthropic_messages( + model="fable-tier", + max_tokens=16, + stream=stream, + disable_fallbacks=True, + messages=[{"role": "user", "content": "hi"}], + ) + body = await _collect(response) if stream else response + + if stream: + assert b'"stop_reason": "refusal"' in body + else: + assert body["stop_reason"] == "refusal" + assert len(fake.calls) == 1 + + +@pytest.mark.asyncio +async def test_disable_fallbacks_beats_a_content_policy_row_too(): + fake = FakeAnthropicUpstream() + router = Router( + model_list=[FABLE_TIER, OPUS_TARGET], + content_policy_fallbacks=[{"fable-tier": ["opus-target"]}], + ) + + with fake.install(): + response = await router.aanthropic_messages( + model="fable-tier", + max_tokens=16, + disable_fallbacks=True, + messages=[{"role": "user", "content": "hi"}], + ) + + assert response["stop_reason"] == "refusal" + assert len(fake.calls) == 1 + + def test_refusal_gate_keys_on_pre_routing_tier_stamp(): router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}]) From e0e249225b1fd2d6620eaad90a13a5a8e62f713d Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 2 Sep 2026 18:55:22 -0700 Subject: [PATCH 15/23] feat(azure): support credential chain for storage (#39229) * feat(azure): support credential chain for storage * test(azure): clarify credential seam suppressions * fix(azure): read chain tokens in a worker thread The credential chain walk (IMDS probe, CLI subprocess) is blocking I/O, so reading the provider inline in async set_valid_azure_ad_token stalls every request on the worker's event loop --- .../azure_storage/azure_storage.py | 49 +++- .../azure_storage/test_azure_storage.py | 241 +++++++++++++++--- .../files/test_azure_blob_storage_backend.py | 47 +++- 3 files changed, 282 insertions(+), 55 deletions(-) diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index cb7175691df..16ef6920114 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -1,7 +1,9 @@ import asyncio import os import time +from collections.abc import Callable from datetime import datetime, timedelta +from functools import cache from typing import Final from litellm._logging import verbose_logger @@ -19,21 +21,40 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.secret_managers.get_azure_ad_token_provider import ( + get_azure_ad_token_provider, +) +from litellm.types.secret_managers.get_azure_ad_token_provider import ( + AzureCredentialType, +) from litellm.types.utils import StandardLoggingPayload +AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default" + + +@cache +def _cached_credential_chain_token_provider() -> Callable[[], str]: + return get_azure_ad_token_provider( + azure_scope=AZURE_STORAGE_TOKEN_SCOPE, + azure_credential=AzureCredentialType.DefaultAzureCredential, + ) + class AzureBlobStorageLogger(CustomBatchLogger): def __init__( self, + build_credential_chain_token_provider: Callable[ + [], Callable[[], str] + ] = _cached_credential_chain_token_provider, **kwargs, ): try: verbose_logger.debug("AzureBlobStorageLogger: in init azure blob storage logger") # Env Variables used for Azure Storage Authentication - self.tenant_id = os.getenv("AZURE_STORAGE_TENANT_ID") - self.client_id = os.getenv("AZURE_STORAGE_CLIENT_ID") - self.client_secret = os.getenv("AZURE_STORAGE_CLIENT_SECRET") + self.tenant_id = os.getenv("AZURE_STORAGE_TENANT_ID") or None + self.client_id = os.getenv("AZURE_STORAGE_CLIENT_ID") or None + self.client_secret = os.getenv("AZURE_STORAGE_CLIENT_SECRET") or None self.azure_storage_account_key: str | None = os.getenv("AZURE_STORAGE_ACCOUNT_KEY") # Required Env Variables for Azure Storage @@ -55,6 +76,9 @@ class AzureBlobStorageLogger(CustomBatchLogger): # Internal variables used for Token based authentication self.azure_auth_token: str | None = None # the Azure AD token to use for Azure Storage API requests self.token_expiry: datetime | None = None # the expiry time of the currentAzure AD token + self._build_credential_chain_token_provider: Callable[[], Callable[[], str]] = ( + build_credential_chain_token_provider + ) asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() @@ -231,10 +255,15 @@ class AzureBlobStorageLogger(CustomBatchLogger): """ Wrapper to set self.azure_auth_token to a valid Azure AD token, refreshing if necessary - Refreshes the token when: - - Token is expired - - Token is not set + Without a service principal configured, the credential chain provider is read every + time; it caches internally and refreshes against the token's real expiry. The read runs + in a worker thread because the chain walk (IMDS probe, CLI subprocess) is blocking """ + if self.tenant_id is None and self.client_id is None and self.client_secret is None: + token_provider: Final = self._build_credential_chain_token_provider() + self.azure_auth_token = await asyncio.to_thread(token_provider) + return + # Check if token needs refresh if self._azure_ad_token_is_expired() or self.azure_auth_token is None: verbose_logger.debug("Azure AD token needs refresh") @@ -273,13 +302,9 @@ class AzureBlobStorageLogger(CustomBatchLogger): tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, - scope="https://storage.azure.com/.default", + scope=AZURE_STORAGE_TOKEN_SCOPE, ) - token: Final = token_provider() - - verbose_logger.debug("azure auth token %s", token) - - return token + return token_provider() def _azure_ad_token_is_expired(self): """ diff --git a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py index 16c518ff412..a96eae0f9c3 100644 --- a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py +++ b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py @@ -1,10 +1,15 @@ +import asyncio import sys +import threading from unittest.mock import AsyncMock, MagicMock, patch import pytest - -from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger +from litellm.integrations.azure_storage.azure_storage import ( + AzureBlobStorageLogger, + _cached_credential_chain_token_provider, +) +from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType from litellm.types.utils import StandardLoggingPayload @@ -25,6 +30,26 @@ def mock_gov_env_vars(mock_env_vars, monkeypatch): monkeypatch.setenv("AZURE_STORAGE_ENDPOINT_SUFFIX", "core.usgovcloudapi.net") +@pytest.fixture +def workload_identity_env_vars(monkeypatch): + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_NAME", "test-account") + monkeypatch.setenv("AZURE_STORAGE_FILE_SYSTEM", "test-container") + for unset in ( + "AZURE_STORAGE_TENANT_ID", + "AZURE_STORAGE_CLIENT_ID", + "AZURE_STORAGE_CLIENT_SECRET", + "AZURE_STORAGE_ACCOUNT_KEY", + "AZURE_STORAGE_ENDPOINT_SUFFIX", + "AZURE_CLIENT_SECRET", + "AZURE_CREDENTIAL", + "AZURE_SCOPE", + ): + monkeypatch.delenv(unset, raising=False) + monkeypatch.setenv("AZURE_CLIENT_ID", "workload-identity-client-id") + monkeypatch.setenv("AZURE_TENANT_ID", "workload-identity-tenant-id") + monkeypatch.setenv("AZURE_FEDERATED_TOKEN_FILE", "/var/run/secrets/azure/tokens/azure-identity-token") + + @pytest.mark.asyncio async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars): """ @@ -32,17 +57,12 @@ async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars): a payload to Azure Blob Storage using the 3-step process (create, append, flush). """ with ( - patch( - "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client" - ) as mock_get_client, - patch( - "litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id" - ) as mock_get_token, + patch("litellm.integrations.azure_storage.azure_storage.get_async_httpx_client") as mock_get_client, + patch("litellm.integrations.azure_storage.azure_storage.get_azure_ad_token_from_entra_id") as mock_get_token, ): # Create mock HTTP client mock_http_client = AsyncMock() - mock_response = AsyncMock() - mock_response.raise_for_status = AsyncMock() + mock_response = MagicMock() mock_http_client.put.return_value = mock_response mock_http_client.patch.return_value = mock_response mock_get_client.return_value = mock_http_client @@ -79,9 +99,7 @@ async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars): put_call_args = mock_http_client.put.call_args assert put_call_args[0][0] == f"{expected_base_url}?resource=file" assert put_call_args[1]["headers"]["x-ms-version"] is not None - assert ( - put_call_args[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token" - ) + assert put_call_args[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token" # Step 2: Append data assert mock_http_client.patch.call_count == 2 # Called for append and flush @@ -89,9 +107,7 @@ async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars): assert append_call[0][0] == f"{expected_base_url}?action=append&position=0" assert append_call[1]["headers"]["x-ms-version"] is not None assert append_call[1]["headers"]["Content-Type"] == "application/json" - assert ( - append_call[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token" - ) + assert append_call[1]["headers"]["Authorization"] == "Bearer mock-azure-ad-token" assert "test-log-id-123" in append_call[1]["data"] # Step 3: Flush data @@ -110,9 +126,7 @@ async def test_async_upload_payload_uses_configured_endpoint_suffix(mock_gov_env AZURE_STORAGE_ENDPOINT_SUFFIX must reach the Entra-ID REST upload path so a sovereign-cloud account is addressed instead of the commercial dfs host. """ - with patch( - "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client" - ) as mock_get_client: + with patch("litellm.integrations.azure_storage.azure_storage.get_async_httpx_client") as mock_get_client: mock_http_client = AsyncMock() mock_response = MagicMock() mock_http_client.put.return_value = mock_response @@ -127,17 +141,10 @@ async def test_async_upload_payload_uses_configured_endpoint_suffix(mock_gov_env await logger.async_upload_payload_to_azure_blob_storage(test_payload) - expected_base_url = ( - "https://test-account.dfs.core.usgovcloudapi.net/test-container/gov-log-id.json" - ) + expected_base_url = "https://test-account.dfs.core.usgovcloudapi.net/test-container/gov-log-id.json" assert mock_http_client.put.call_args[0][0] == f"{expected_base_url}?resource=file" - assert ( - mock_http_client.patch.call_args_list[0][0][0] - == f"{expected_base_url}?action=append&position=0" - ) - assert mock_http_client.patch.call_args_list[1][0][0].startswith( - f"{expected_base_url}?action=flush" - ) + assert mock_http_client.patch.call_args_list[0][0][0] == f"{expected_base_url}?action=append&position=0" + assert mock_http_client.patch.call_args_list[1][0][0].startswith(f"{expected_base_url}?action=flush") @pytest.mark.asyncio @@ -148,9 +155,7 @@ async def test_service_client_uses_configured_endpoint_suffix(mock_gov_env_vars) """ fake_aio_module = MagicMock() - with patch.dict( - sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module} - ): + with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}): logger = AzureBlobStorageLogger() await logger.get_service_client() @@ -160,14 +165,180 @@ async def test_service_client_uses_configured_endpoint_suffix(mock_gov_env_vars) ) +@pytest.mark.asyncio +async def test_upload_authenticates_through_the_credential_chain_under_workload_identity( + workload_identity_env_vars, +): + build_provider = MagicMock(return_value=lambda: "workload-identity-token") + with patch( # test-quality-ok: REST client is created inside the method; assert emitted request headers + "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client" + ) as mock_get_client: + mock_http_client = AsyncMock() + mock_http_client.put.return_value = MagicMock() + mock_http_client.patch.return_value = MagicMock() + mock_get_client.return_value = mock_http_client + + logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider) + await logger.async_upload_payload_to_azure_blob_storage({"id": "wif-log-id"}) + + build_provider.assert_called_once_with() + assert logger.azure_auth_token == "workload-identity-token" + sent_headers = [mock_http_client.put.call_args[1]["headers"]] + [ + call[1]["headers"] for call in mock_http_client.patch.call_args_list + ] + assert len(sent_headers) == 3 + assert all(headers["Authorization"] == "Bearer workload-identity-token" for headers in sent_headers) + + +def test_default_chain_provider_is_storage_scoped_and_built_once_per_process(): + _cached_credential_chain_token_provider.cache_clear() + with ( + patch( # test-quality-ok: assert the default factory's fixed scope and credential type without constructing Azure SDK credentials + "litellm.integrations.azure_storage.azure_storage.get_azure_ad_token_provider", + return_value=lambda: "chain-token", + ) as mock_builder + ): + first = _cached_credential_chain_token_provider() + second = _cached_credential_chain_token_provider() + _cached_credential_chain_token_provider.cache_clear() + + assert first is second + assert first() == "chain-token" + mock_builder.assert_called_once_with( + azure_scope="https://storage.azure.com/.default", + azure_credential=AzureCredentialType.DefaultAzureCredential, + ) + + +@pytest.mark.asyncio +async def test_chain_tokens_are_read_from_the_provider_on_every_refresh( + workload_identity_env_vars, +): + provider = MagicMock(side_effect=["chain-token-1", "chain-token-2"]) + logger = AzureBlobStorageLogger(build_credential_chain_token_provider=MagicMock(return_value=provider)) + await logger.set_valid_azure_ad_token() + first_token = logger.azure_auth_token + await logger.set_valid_azure_ad_token() + + assert first_token == "chain-token-1" + assert logger.azure_auth_token == "chain-token-2" + assert provider.call_count == 2 + + +@pytest.mark.asyncio +async def test_chain_token_read_yields_to_the_event_loop(workload_identity_env_vars): + """ + The chain walk is blocking I/O (IMDS probe, CLI subprocess), so reading the provider + inline would stall every request on the worker. Prove other coroutines run during the read. + """ + loop_was_free = threading.Event() + + def provider() -> str: + if not loop_was_free.wait(timeout=5): + raise TimeoutError("the event loop never ran the observer while the token was being read") + return "chain-token" + + async def observer(): + loop_was_free.set() + + logger = AzureBlobStorageLogger(build_credential_chain_token_provider=MagicMock(return_value=provider)) + observer_task = asyncio.create_task(observer()) + await logger.set_valid_azure_ad_token() + await observer_task + + assert logger.azure_auth_token == "chain-token" + + +@pytest.mark.asyncio +async def test_empty_string_service_principal_vars_still_use_the_credential_chain( + workload_identity_env_vars, monkeypatch +): + for name in ("AZURE_STORAGE_TENANT_ID", "AZURE_STORAGE_CLIENT_ID", "AZURE_STORAGE_CLIENT_SECRET"): + monkeypatch.setenv(name, "") + + logger = AzureBlobStorageLogger( + build_credential_chain_token_provider=MagicMock(return_value=lambda: "workload-identity-token") + ) + await logger.set_valid_azure_ad_token() + + assert logger.azure_auth_token == "workload-identity-token" + + +@pytest.mark.asyncio +async def test_client_secret_auth_still_uses_the_storage_scoped_service_principal(mock_env_vars): + build_provider = MagicMock() + with ( + patch( # test-quality-ok: assert the storage scope passed to the shared token factory without making an external auth call + "litellm.integrations.azure_storage.azure_storage.get_azure_ad_token_from_entra_id", + return_value=lambda: "client-secret-token", + ) as mock_entra_id + ): + logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider) + await logger.set_valid_azure_ad_token() + + assert logger.azure_auth_token == "client-secret-token" + build_provider.assert_not_called() + assert mock_entra_id.call_args.kwargs == { + "tenant_id": "test-tenant-id", + "client_id": "test-client-id", + "client_secret": "test-client-secret", + "scope": "https://storage.azure.com/.default", + } + + +@pytest.mark.parametrize( + "missing_var", + ["AZURE_STORAGE_TENANT_ID", "AZURE_STORAGE_CLIENT_ID", "AZURE_STORAGE_CLIENT_SECRET"], +) +@pytest.mark.asyncio +async def test_partially_configured_service_principal_still_names_the_missing_variable( + mock_env_vars, monkeypatch, missing_var +): + monkeypatch.delenv(missing_var) + + build_provider = MagicMock() + logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider) + with pytest.raises(ValueError, match=f"Missing required environment variable: {missing_var}"): + await logger.set_valid_azure_ad_token() + + build_provider.assert_not_called() + + +@pytest.mark.asyncio +async def test_account_key_auth_never_requests_a_token(workload_identity_env_vars, monkeypatch): + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_KEY", "dGVzdC1rZXk=") + + file_client = MagicMock() + file_client.create_file = AsyncMock() + file_client.append_data = AsyncMock() + file_client.flush_data = AsyncMock() + directory_client = MagicMock() + directory_client.exists = AsyncMock(return_value=True) + directory_client.get_file_client = MagicMock(return_value=file_client) + file_system_client = MagicMock() + file_system_client.get_directory_client = MagicMock(return_value=directory_client) + service_client = MagicMock() + service_client.get_file_system_client = MagicMock(return_value=file_system_client) + fake_aio_module = MagicMock() + fake_aio_module.DataLakeServiceClient = MagicMock(return_value=service_client) + + build_provider = MagicMock() + with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}): + logger = AzureBlobStorageLogger(build_credential_chain_token_provider=build_provider) + await logger.async_upload_payload_to_azure_blob_storage({"id": "account-key-log-id"}) + + build_provider.assert_not_called() + assert logger.azure_auth_token is None + file_client.flush_data.assert_awaited_once() + assert fake_aio_module.DataLakeServiceClient.call_args.kwargs["credential"] == "dGVzdC1rZXk=" + + @pytest.mark.asyncio async def test_service_client_defaults_to_commercial_endpoint(mock_env_vars): """Unset AZURE_STORAGE_ENDPOINT_SUFFIX keeps the pre-existing commercial host""" fake_aio_module = MagicMock() - with patch.dict( - sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module} - ): + with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}): logger = AzureBlobStorageLogger() await logger.get_service_client() diff --git a/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py b/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py index 7e1139ca79a..222ead92b54 100644 --- a/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py +++ b/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py @@ -27,6 +27,20 @@ def mock_gov_env_vars(mock_env_vars, monkeypatch): monkeypatch.setenv("AZURE_STORAGE_ENDPOINT_SUFFIX", GOV_SUFFIX) +@pytest.fixture +def credential_chain_env_vars(monkeypatch): + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_NAME", "test-account") + monkeypatch.setenv("AZURE_STORAGE_FILE_SYSTEM", "test-container") + for name in ( + "AZURE_STORAGE_TENANT_ID", + "AZURE_STORAGE_CLIENT_ID", + "AZURE_STORAGE_CLIENT_SECRET", + "AZURE_STORAGE_ACCOUNT_KEY", + "AZURE_STORAGE_ENDPOINT_SUFFIX", + ): + monkeypatch.delenv(name, raising=False) + + def _make_backend() -> AzureBlobStorageBackend: backend = AzureBlobStorageBackend() backend.azure_auth_token = "mock-azure-ad-token" @@ -42,6 +56,29 @@ def _mock_upload_client() -> AsyncMock: return client +@pytest.mark.asyncio +async def test_upload_file_with_credential_chain(credential_chain_env_vars): + client = _mock_upload_client() + build_provider = MagicMock(return_value=lambda: "workload-identity-token") + + with patch( # test-quality-ok: the backend creates its REST client internally; assert the emitted authorization header + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=client + ): + backend = AzureBlobStorageBackend(build_credential_chain_token_provider=build_provider) + storage_url = await backend.upload_file( + file_content=b"hello", + filename="report.json", + content_type="application/json", + path_prefix="logs", + file_naming_strategy="original_filename", + ) + + build_provider.assert_called_once_with() + assert storage_url == "https://test-account.blob.core.windows.net/test-container/logs/report.json" + assert client.put.call_args[1]["headers"]["Authorization"] == "Bearer workload-identity-token" + assert client.patch.call_count == 2 + + @pytest.mark.parametrize( "env_fixture, expected_suffix", [("mock_env_vars", "core.windows.net"), ("mock_gov_env_vars", GOV_SUFFIX)], @@ -125,10 +162,7 @@ async def test_download_file_accepts_url_persisted_before_the_suffix_was_set(moc ) assert content == b"file-bytes" - assert ( - client.get.call_args[0][0] - == f"https://test-account.blob.{GOV_SUFFIX}/test-container/logs/report.json" - ) + assert client.get.call_args[0][0] == f"https://test-account.blob.{GOV_SUFFIX}/test-container/logs/report.json" @pytest.mark.parametrize( @@ -178,10 +212,7 @@ async def test_download_file_drops_query_string_from_the_stored_url(mock_env_var "https://test-account.blob.core.windows.net/test-container/logs/report.json?sig=redacted&se=2026" ) - assert ( - client.get.call_args[0][0] - == "https://test-account.blob.core.windows.net/test-container/logs/report.json" - ) + assert client.get.call_args[0][0] == "https://test-account.blob.core.windows.net/test-container/logs/report.json" @pytest.mark.parametrize( From 8065ede40bb380d8947939120f95a740f970614f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:11:08 +0000 Subject: [PATCH 16/23] test(guardrails): expect the deduped end-of-stream scan in crowdstrike cadence test (#39467) Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/guardrail_hooks/test_crowdstrike_aidr.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index 1b50ea53db2..a07157396df 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -1703,8 +1703,8 @@ async def _guard_calls_for_stream(handler: CrowdStrikeAIDRHandler, chunk_texts: @pytest.mark.parametrize( ("configured", "expected_calls"), [ - ({}, 3), - ({"streaming_sampling_rate": 2}, 6), + ({}, 2), + ({"streaming_sampling_rate": 2}, 5), ({"streaming_end_of_stream_only": True}, 1), ({"streaming_end_of_stream_only": True, "streaming_sampling_rate": 2}, 1), ], @@ -1712,7 +1712,10 @@ async def _guard_calls_for_stream(handler: CrowdStrikeAIDRHandler, chunk_texts: async def test_streaming_params_from_config_control_output_scan_cadence( configured: dict[str, object], expected_calls: int ) -> None: - """10 chunks: default samples at 5 and 10 plus the final pass, rate 2 samples 5 times plus final, end-of-stream scans once.""" + """10 chunks: default samples at 5 and 10, rate 2 samples 5 times, end-of-stream scans once. + + The final pass is skipped because chunk 10 already scanned the complete output. + """ handler = _initialize_from_config(mode="post_call", **configured) assert await _guard_calls_for_stream(handler, list("ABCDEFGHIJ")) == expected_calls From 7a81ae98e68746aab8aadad58ee6706a2b04de00 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 2 Sep 2026 19:15:14 -0700 Subject: [PATCH 17/23] fix(model_armor): handle Anthropic Messages and Responses streams in post_call (#39181) * fix(model_armor): handle Anthropic Messages and Responses streams in post_call The post_call streaming hook buffered every chunk and fed it to stream_chunk_builder, which only understands chat-completion deltas. /v1/messages streams raw Anthropic SSE bytes and /v1/responses streams typed Responses events, so both raised litellm.APIError and surfaced to the client as a 500 on every streamed request. Assemble each surface with its own reader, frame guardrail failures as terminal items in that surface's wire format, and pass the stream through unscanned when it cannot be assembled instead of raising. * fix(model_armor): classify the stream surface and fail closed when it cannot be assembled Decide the wire format explicitly instead of inferring it from a boolean pair, so an opaque raw SSE stream (the Google :streamGenerateContent route) is never refused in Anthropic framing, and a stream that cannot be assembled is blocked rather than released unscanned unless fail_on_error is disabled. Also scan Responses tool-call arguments, read the body only off a terminal Responses event, and record the applied guardrail on the fail-closed path. * test(model_armor): pin the error-only stream predicate against content-carrying streams is_sse_error_stream decides whether a buffered stream is forwarded to the client untouched, so a stream that still carries content must not qualify: the frames-only join drops typed chunks, an empty stream is not a refusal, and a content event may carry an empty error field. * fix(model_armor): let a streamed de-identify match mask instead of blocking A de-identify template reports MATCH_FOUND for every redaction it makes. The streaming block check omitted allow_sanitization, so with mask_response_content enabled that match read as a refusal and the client got a 400 where the non-streaming sibling returned the redacted text. Pass the flag through, as the non-streaming hook already does, and stamp the logged status from the same decision so the spend row agrees with what the client received. Also drop Any from the chat-completion assembler's parameter; stream_chunk_builder takes a bare list, so list[object] carries the mutability requirement without erasing the element type. * fix(model_armor): fail closed when a streamed de-identify match cannot be applied Allowing sanitization past the streaming block check is a promise to apply the redaction Model Armor asked for. Two paths broke that promise and released the buffered original instead: a match that comes back with no sanitized text, and a surface with no assembled body to rewrite. The outcome is now resolved once, before it is recorded, so the status stamped on request metadata agrees with what the client receives rather than reporting the success the block check alone would have implied. * fix: scan the deltas when a Responses stream ends without a body response.failed and response.incomplete are terminal events like response.completed, but a turn that broke mid-generation reports an empty output while the deltas ahead of it already spelled the answer out to the client. Reading only the terminal body found nothing to scan there, and the empty-content shortcut then forwarded every buffered delta past the guardrail. Fall back to the text the delta events carry whenever a Responses stream assembles to nothing. * fix: read the Responses delta event types off the event enum The hand-listed set left out response.mcp_call_arguments.delta, so a turn that streamed only MCP tool arguments and then reported an empty body still took the no-content shortcut and forwarded those chunks unscanned. Deriving the set from ResponsesAPIStreamEvents keeps it complete as the enum grows, and the str guard in the reader already covers any event whose delta is not text. * fix(model_armor): scan responses deltas alongside the terminal body A /v1/responses stream spells out reasoning summaries and tool-call arguments in delta events that its terminal body never repeats, so scanning the body alone handed every summary delta to the client unscanned whenever the body carried text. * fix(model_armor): scan responses delta fields apart from each other A Responses turn spells out its reasoning summary, its visible answer and its tool-call arguments in separate delta events. Joining every delta into one string let a finding form across the boundary between two fields that each carry nothing to find, so a safe stream could be blocked. Group the deltas by the field they belong to, join a field's own deltas as they streamed, and keep the fields apart. * fix(model_armor): scan each responses field once, not twice Separating delta fields stopped the terminal body from matching the delta text, so a turn with two visible fields sent Model Armor both copies. Only the delta fields the body does not already carry are appended now. --------- Co-authored-by: yassin --- litellm/proxy/guardrails/anthropic_sse.py | 64 +- .../model_armor/model_armor.py | 479 +++++-- .../guardrail_hooks/test_model_armor.py | 1151 +++++++++++++++++ 3 files changed, 1589 insertions(+), 105 deletions(-) diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index 50c05daee11..28220f09f00 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -13,6 +13,19 @@ from typing import Final from litellm.types.utils import Choices, ModelResponse +_ANTHROPIC_EVENT_TYPES: Final = frozenset( + { + "message_start", + "message_delta", + "message_stop", + "content_block_start", + "content_block_delta", + "content_block_stop", + "ping", + "error", + } +) + def is_raw_sse_stream(all_chunks: Sequence[object]) -> bool: return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks) @@ -30,23 +43,43 @@ def _joined_sse_stream(all_chunks: Sequence[object]) -> str | None: return None -def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None: +def _parsed_sse_events(sse_stream: str) -> tuple[Mapping[str, object], ...]: from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, ) + return tuple( + event_data + for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream) # pyright: ignore[reportPrivateUsage] # same parser the assembler uses + if (event_data := AnthropicPassthroughLoggingHandler._extract_sse_data(event)) is not None # pyright: ignore[reportPrivateUsage] # same parser the assembler uses; a private import beats forking SSE parsing + ) + + +def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None: return next( ( message - for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream) # pyright: ignore[reportPrivateUsage] # same parser the assembler uses - if (event_data := AnthropicPassthroughLoggingHandler._extract_sse_data(event)) is not None # pyright: ignore[reportPrivateUsage] # same parser the assembler uses; a private import beats forking SSE parsing - and event_data.get("type") == "message_start" - and isinstance(message := event_data.get("message"), dict) + for event_data in _parsed_sse_events(sse_stream) + if event_data.get("type") == "message_start" and isinstance(message := event_data.get("message"), dict) ), None, ) +def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool: + """Whether raw SSE frames are Anthropic Messages events. + + ``is_raw_sse_stream`` only says the chunks are unparsed bytes, and ``/v1/messages`` is not the + only endpoint that streams those: the Google ``:streamGenerateContent`` route marks its own + stream raw too. Reading its frames as Anthropic ones would refuse the response in a wire format + its client cannot parse, so the surface is decided on the event types actually present. + """ + sse_stream: Final = _joined_sse_stream(all_chunks) + if sse_stream is None: + return False + return any(event.get("type") in _ANTHROPIC_EVENT_TYPES for event in _parsed_sse_events(sse_stream)) + + def assemble_anthropic_sse_stream( all_chunks: Sequence[object], *, restore_identity: bool = False ) -> ModelResponse | None: @@ -111,6 +144,27 @@ def anthropic_sse_error_frames(message: str) -> tuple[bytes, ...]: ) +def is_sse_error_stream(all_chunks: Sequence[object]) -> bool: + """Whether the buffered stream carries nothing but error frames. + + post_call guardrails run in a chain, so a hook can be handed the terminal error frames an + earlier guardrail emitted when it blocked. Those carry no message to assemble, and replacing + them would hide the refusal the client is owed. Covers both wire forms a guardrail emits: the + Anthropic ``error`` event and the chat-completions ``{"error": ...}`` payload. + """ + if not all(isinstance(chunk, (str, bytes)) for chunk in all_chunks): + # A stream mixing typed chunks with an error frame still carries content to scan, and the + # frames-only join below would drop exactly the part that has to be scanned + return False + sse_stream: Final = _joined_sse_stream(all_chunks) + if sse_stream is None: + return False + events: Final = _parsed_sse_events(sse_stream) + return len(events) > 0 and all( + event.get("type") == "error" or isinstance(event.get("error"), Mapping) for event in events + ) + + def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]: from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index d187b5b12e9..7d88a037f4f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,4 +1,5 @@ from collections.abc import AsyncGenerator, Mapping, Sequence +from enum import Enum, auto from typing import TYPE_CHECKING, Any, Final, Literal import httpx @@ -27,12 +28,25 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.anthropic_sse import ( + anthropic_sse_chunks_from_response, + anthropic_sse_error_frames, + assemble_anthropic_sse_stream, + is_anthropic_sse_stream, + is_raw_sse_stream, + is_sse_error_stream, +) from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import ( MODEL_ARMOR_MAX_FILE_SIZE_BYTES, plan_file_scans, ) from litellm.types.guardrails import GuardrailEventHooks, LitellmParams -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionToolCallChunk, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, +) from litellm.types.utils import ( CallTypes, CallTypesLiteral, @@ -41,10 +55,33 @@ from litellm.types.utils import ( ModelResponse, ModelResponseStream, StandardLoggingGuardrailInformation, + TextCompletionResponse, ) GUARDRAIL_NAME: Final = "model_armor" +# Only these carry the finished output; response.created carries an empty body +_RESPONSES_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete", "response.failed"}) + +# Every event whose ``delta`` is model output already on its way to the client. Read off the event +# enum rather than listed, so an event added there cannot quietly fall out of the scan +_RESPONSES_DELTA_EVENT_TYPES: Final = frozenset( + event.value for event in ResponsesAPIStreamEvents if event.value.endswith(".delta") +) + +# What makes two delta events part of the same field of the turn, rather than two fields that merely +# streamed next to each other +_RESPONSES_DELTA_FIELD_ATTRS: Final = ("type", "item_id", "output_index", "content_index", "summary_index") + + +class _StreamSurface(Enum): + """Wire format of a buffered streaming response, which decides how it is read and how it is refused.""" + + CHAT_COMPLETIONS = auto() + ANTHROPIC_MESSAGES = auto() + RESPONSES = auto() + OPAQUE_SSE = auto() + class ModelArmorAPIError(Exception): """Model Armor API failure (non-2xx), distinct from a content-block decision so @@ -322,19 +359,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): else: return {"modelResponseData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}} - def _should_block_content(self, armor_response: dict, allow_sanitization: bool = False) -> bool: + def _should_block_content(self, armor_response: Mapping[str, Any], allow_sanitization: bool = False) -> bool: """Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult.""" - sanitization_result: Final = armor_response.get("sanitizationResult", {}) - filter_results: Final = sanitization_result.get("filterResults", {}) - - # filterResults can be a dict (named keys) or a list (array of filter result dicts) - filter_result_items = [] - if isinstance(filter_results, dict): - filter_result_items = list(filter_results.values()) - elif isinstance(filter_results, list): - filter_result_items = filter_results - - for filt in filter_result_items: + for filt in self._filter_result_items(armor_response): # Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before if filt.get("raiFilterResult", {}).get("matchState") == "MATCH_FOUND": return True @@ -358,22 +385,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Fallback dict code removed; all cases handled above return False - def _get_sanitized_content(self, armor_response: dict) -> str | None: + def _get_sanitized_content(self, armor_response: Mapping[str, Any]) -> str | None: """ Get the sanitized content from a Model Armor response, if available. Looks for sanitized text in deidentifyResult, and falls back to root-level fields if not found. """ - result: Final = armor_response.get("sanitizationResult", {}) - filter_results: Final = result.get("filterResults", {}) - - # filterResults can be a dict (single filter) or a list (multiple filters) - filters: Final = ( - list(filter_results.values()) - if isinstance(filter_results, dict) - else filter_results - if isinstance(filter_results, list) - else [] - ) + filters: Final = self._filter_result_items(armor_response) # Prefer sanitized text from deidentifyResult if present for filter_entry in filters: @@ -397,6 +414,61 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Fallback: if Model Armor put sanitized text at the root, use it return armor_response.get("sanitizedText") or armor_response.get("text") + @staticmethod + def _filter_result_items(armor_response: Mapping[str, Any]) -> Sequence[Any]: + """Every filter result in a scan response. + + filterResults is a dict of named filters on most templates and a list on some, so both + shapes are flattened to the same list of filter entries. + """ + filter_results: Final = armor_response.get("sanitizationResult", {}).get("filterResults", {}) + if isinstance(filter_results, dict): + return list(filter_results.values()) + if isinstance(filter_results, list): + return filter_results + return [] + + def _has_deidentify_match(self, armor_response: Mapping[str, Any]) -> bool: + """Whether an SDP de-identify filter matched, i.e. Model Armor owes this response a redaction.""" + for filter_entry in self._filter_result_items(armor_response): + sdp = filter_entry.get("sdpFilterResult") + if sdp and sdp.get("deidentifyResult", {}).get("matchState") == "MATCH_FOUND": + return True + return False + + def _resolve_streaming_outcome( + self, + armor_response: Mapping[str, Any], + assembled_response: object, + content: str, + ) -> tuple[bool, str | None]: + """Whether to block the buffered stream, and the rewrite to emit when it is not blocked. + + A de-identify match only reaches here unblocked because masking is on, so the redaction it + stands for has to be both resolvable and emittable. Where it is neither, the buffered + original still carries what Model Armor matched on, so this fails closed instead of + releasing it. + """ + if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content): + return True, None + if not self.mask_response_content: + return False, None + + sanitized_content: Final = self._get_sanitized_content(armor_response) + if not sanitized_content: + # No rewrite to apply. Harmless unless a match is outstanding, in which case applying + # nothing would hand back the very content that matched + return self._has_deidentify_match(armor_response), None + if sanitized_content == content: + return False, None + if not isinstance(assembled_response, ModelResponse): + verbose_proxy_logger.warning( + "Model Armor: sanitized content cannot be re-emitted on this streaming endpoint, " + "blocking the response instead" + ) + return True, None + return False, sanitized_content + @staticmethod def _append_armor_response(existing: object, armor_response: Mapping[str, object]) -> object: """Accumulate scan responses so a later text scan does not drop an earlier file scan. @@ -831,6 +903,185 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return response + @staticmethod + def _is_terminal_error_stream(all_chunks: Sequence[object]) -> bool: + """Whether the buffered stream is only the refusal an earlier guardrail in the chain emitted. + + post_call guardrails are composed, so this hook can be handed the terminal error items a + preceding one produced. They carry no message to scan, and replacing them would hide the + refusal the client is owed. + """ + if all(getattr(chunk, "type", None) == "error" for chunk in all_chunks): + return True + return is_sse_error_stream(all_chunks) + + @staticmethod + def _classify_stream(all_chunks: Sequence[object]) -> _StreamSurface: + """Wire format the buffered chunks belong to.""" + if is_raw_sse_stream(all_chunks): + return ( + _StreamSurface.ANTHROPIC_MESSAGES if is_anthropic_sse_stream(all_chunks) else _StreamSurface.OPAQUE_SSE + ) + if any( + isinstance(event_type := getattr(chunk, "type", None), str) and event_type.startswith("response.") + for chunk in all_chunks + ): + return _StreamSurface.RESPONSES + return _StreamSurface.CHAT_COMPLETIONS + + @staticmethod + def _final_responses_api_response(all_chunks: Sequence[object]) -> ResponsesAPIResponse | None: + """Response body carried by a terminal ``/v1/responses`` event. + + A stream cut short before it completes has to read as unassembled rather than as a clean + empty response: ``response.created`` also carries a body, but an empty one, and scanning + that would release every buffered delta unscanned. + """ + return next( + ( + body + for chunk in reversed(all_chunks) + if getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES + and isinstance(body := getattr(chunk, "response", None), ResponsesAPIResponse) + ), + None, + ) + + @staticmethod + def _responses_api_response_text(response: ResponsesAPIResponse) -> str: + """Text to scan in a Responses API response, tool-call arguments included. + + Tool calls are folded in because ``get_content_from_model_response`` folds them into what + the chat surface scans, and a Responses turn can carry its whole payload in them. + """ + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + + texts: Final[list[str]] = [] # mutable-ok: the shared extractor below appends into caller-owned lists + tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] # mutable-ok: the same extractor's tool-call sink + handler: Final = OpenAIResponsesHandler() + for output_idx, output_item in enumerate(response.output or ()): + handler._extract_output_text_and_images( # pyright: ignore[reportPrivateUsage] # the shared Responses output extractor; forking it would duplicate per-item parsing + output_item=output_item, + output_idx=output_idx, + texts_to_check=texts, + images_to_check=[], # mutable-ok: the extractor's images sink, unused here + task_mappings=[], # mutable-ok: the extractor's task-mapping sink, unused here + tool_calls_to_check=tool_calls, + ) + return "".join((*texts, *(json.dumps(tool_call) for tool_call in tool_calls))) + + def _extract_streaming_content(self, assembled_response: object) -> str: + """Text to scan from an assembled stream, for every endpoint shape this hook serves.""" + if isinstance(assembled_response, ResponsesAPIResponse): + return self._responses_api_response_text(assembled_response) + return self._extract_content_from_response(assembled_response) + + @staticmethod + def _responses_delta_field(chunk: object) -> tuple[str, ...]: + """Which field of the turn a delta event belongs to.""" + return tuple(str(getattr(chunk, attr, None)) for attr in _RESPONSES_DELTA_FIELD_ATTRS) + + @staticmethod + def _responses_delta_field_texts(all_chunks: Sequence[object]) -> tuple[str, ...]: + """Text each field of a ``/v1/responses`` turn has already spelled out in its delta events. + + One field's deltas are joined as they streamed, since a finding can be split across them, + and separate fields stay apart, so a reasoning summary running into the visible answer + cannot spell out a finding that neither of them carries. + """ + deltas: Final = tuple( + (ModelArmorGuardrail._responses_delta_field(chunk), delta) + for chunk in all_chunks + if getattr(chunk, "type", None) in _RESPONSES_DELTA_EVENT_TYPES + and isinstance(delta := getattr(chunk, "delta", None), str) + ) + return tuple( + "".join(delta for field, delta in deltas if field == streamed_field) + for streamed_field in dict.fromkeys(field for field, _ in deltas) + ) + + def _streaming_content_to_scan( + self, + assembled_response: object, + all_chunks: Sequence[object], + surface: _StreamSurface, + ) -> str: + """Text to scan for a buffered stream, which is everything the client is about to receive. + + A ``/v1/responses`` stream also spells out reasoning summaries and tool-call arguments in + delta events that its terminal body never repeats, so every delta field the body does not + already carry is scanned after it. + """ + content: Final = self._extract_streaming_content(assembled_response) + if surface is not _StreamSurface.RESPONSES: + return content + unscanned: Final = tuple(text for text in self._responses_delta_field_texts(all_chunks) if text not in content) + return "\n".join(part for part in (content, *unscanned) if part) + + @staticmethod + def _apply_sanitized_content(assembled_response: ModelResponse, sanitized_content: str) -> None: + """Replace every non-empty choice message with the Model Armor sanitized text.""" + for choice in assembled_response.choices: + if isinstance(choice, Choices) and choice.message.content: + choice.message.content = sanitized_content + + @staticmethod + def _assemble_chat_completion_stream( + all_chunks: list[object], # mutable-ok: stream_chunk_builder only accepts a mutable list + ) -> ModelResponse | TextCompletionResponse | None: + """Assemble chat-completion chunks, returning ``None`` when they cannot be assembled.""" + from litellm.main import stream_chunk_builder + + try: + return stream_chunk_builder(chunks=all_chunks) + except Exception as exc: + verbose_proxy_logger.warning("Model Armor: chat-completion stream assembly failed (%s)", exc) + return None + + def _assemble_stream( + self, all_chunks: Sequence[object], surface: _StreamSurface + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | None: + """Assemble the buffered stream into the scannable response its surface produces.""" + if surface is _StreamSurface.ANTHROPIC_MESSAGES: + return assemble_anthropic_sse_stream(all_chunks, restore_identity=True) + if surface is _StreamSurface.RESPONSES: + return self._final_responses_api_response(all_chunks) + if surface is _StreamSurface.OPAQUE_SSE: + return None + return self._assemble_chat_completion_stream(list(all_chunks)) + + @staticmethod + def _error_payload(exc: HTTPException) -> Mapping[str, object]: + """Error object for a terminal stream item, carrying the status the frame would otherwise lose.""" + detail: Final = exc.detail if isinstance(exc.detail, Mapping) else {"message": str(exc.detail)} + error_value: Final = detail.get("error", detail) + return { + **(dict(error_value) if isinstance(error_value, Mapping) else {"message": str(error_value)}), + "code": str(exc.status_code), + } + + @staticmethod + def _build_responses_error_items(exc: HTTPException) -> Sequence[object] | None: + """Responses API error events for a failure discovered after the stream started.""" + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + + return OpenAIResponsesHandler().build_stream_error_items(exc, responses_so_far=None) + + def _stream_error_items(self, exc: HTTPException, *, surface: _StreamSurface) -> Sequence[object]: + """Frame a guardrail failure as terminal stream items in this endpoint's wire format.""" + payload: Final = self._error_payload(exc) + if surface is _StreamSurface.ANTHROPIC_MESSAGES: + return anthropic_sse_error_frames(str(payload.get("message", ""))) + if surface is _StreamSurface.RESPONSES and (responses_items := self._build_responses_error_items(exc)): + return responses_items + # Also the fallback when a surface cannot frame its own error: create_response() reads the + # status back out of this form, so the refusal keeps its code instead of arriving as a 200 + return (f"data: {json.dumps({'error': payload})}\n\n",) + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -840,97 +1091,125 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): """Process streaming response chunks.""" from litellm.llms.base_llm.base_model_iterator import MockResponseIterator - from litellm.main import stream_chunk_builder + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) # Collect all chunks - all_chunks: Final[list[ModelResponseStream]] = [] + all_chunks: Final[list[Any]] = [] async for chunk in response: all_chunks.append(chunk) + if not all_chunks or self._is_terminal_error_stream(all_chunks): + for chunk in all_chunks: + yield chunk + return + + surface: Final = self._classify_stream(all_chunks) + # Build complete response - assembled_response: Final = stream_chunk_builder(chunks=all_chunks) + assembled_response: Final = self._assemble_stream(all_chunks, surface) - if isinstance(assembled_response, ModelResponse): - # Extract content - content: Final = self._extract_content_from_response(assembled_response) + if assembled_response is None: + if not self.optional_params.get("fail_on_error", True): + verbose_proxy_logger.warning( + "Model Armor: streamed response could not be assembled for scanning, " + "forwarding it unscanned because fail_on_error is disabled" + ) + for chunk in all_chunks: + yield chunk + return - if content: - try: - # Check with Model Armor - armor_response: Final = await self.make_model_armor_request( - content=content, - source="model_response", - request_data=request_data, - ) + # Forwarding an unscannable stream would silently disable the guardrail, so fail closed + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) + for error_item in self._stream_error_items( + HTTPException( + status_code=500, + detail=f"{self.guardrail_name}: streamed response could not be assembled for scanning, blocking it", + ), + surface=surface, + ): + yield error_item + return - # Attach Model Armor response & status to this request's metadata to avoid race conditions - if isinstance(request_data, dict): - _, metadata = get_or_create_metadata_bucket(request_data) - metadata["_model_armor_response"] = self._build_logging_response(armor_response) - metadata["_model_armor_status"] = ( - "blocked" if self._should_block_content(armor_response) else "success" - ) + # Extract content + content: Final = self._streaming_content_to_scan( + assembled_response=assembled_response, all_chunks=all_chunks, surface=surface + ) - # Add guardrail to applied_guardrails BEFORE potential blocking - # This ensures guardrail is recorded even when it blocks the request - from litellm.proxy.common_utils.callback_utils import ( - add_guardrail_to_applied_guardrails_header, - ) + if not content: + verbose_proxy_logger.debug("Model Armor: No text content in streaming response, skipping guardrail") + for chunk in all_chunks: + yield chunk + return - add_guardrail_to_applied_guardrails_header( - request_data=request_data, guardrail_name=self.guardrail_name - ) + try: + # Check with Model Armor + armor_response: Final = await self.make_model_armor_request( + content=content, + source="model_response", + request_data=request_data, + ) - # Check if blocked - if self._should_block_content(armor_response): - raise HTTPException( - status_code=400, - detail=self._build_block_error_detail( - "Streaming response blocked by Model Armor", - armor_response, - ), - ) + # Decide the outcome before recording it. Mirrors the non-streaming sibling: with + # masking on, a de-identify match is a redaction to apply rather than a refusal, but + # that only holds while the redaction can actually be delivered + blocked, sanitized_content = self._resolve_streaming_outcome( + armor_response=armor_response, + assembled_response=assembled_response, + content=content, + ) - # Apply sanitization if enabled - if self.mask_response_content: - sanitized_content: Final = self._get_sanitized_content(armor_response) - if sanitized_content and sanitized_content != content: - # Update assembled response - for choice in assembled_response.choices: - if isinstance(choice, Choices): - if choice.message.content: - choice.message.content = sanitized_content + # Attach Model Armor response & status to this request's metadata to avoid race conditions + if isinstance(request_data, dict): + _, metadata = get_or_create_metadata_bucket(request_data) + metadata["_model_armor_response"] = self._build_logging_response(armor_response) + metadata["_model_armor_status"] = "blocked" if blocked else "success" - # Return sanitized stream - mock_response: Final = MockResponseIterator(model_response=assembled_response) - async for chunk in mock_response: - yield chunk - return + # Add guardrail to applied_guardrails BEFORE potential blocking + # This ensures guardrail is recorded even when it blocks the request + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) - except ModelArmorAPIError as e: - if self.optional_params.get("fail_on_error", True): - error_obj = {"message": e.detail, "code": "500"} - yield f"data: {json.dumps({'error': error_obj})}\n\n" - return - except HTTPException as e: - # Yield error as SSE event so create_response() detects it and - # returns a proper JSON error response with the correct status code. - # (Raising from a generator hits create_response's generic except → 500.) - detail: Final = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)} - error_value: Final = detail.get("error", detail) - if isinstance(error_value, dict): - error_obj = dict(error_value) - else: - error_obj = {"message": str(error_value)} - error_obj["code"] = str(e.status_code) - yield f"data: {json.dumps({'error': error_obj})}\n\n" + if blocked: + raise HTTPException( + status_code=400, + detail=self._build_block_error_detail( + "Streaming response blocked by Model Armor", + armor_response, + ), + ) + + if sanitized_content is not None and isinstance(assembled_response, ModelResponse): + self._apply_sanitized_content(assembled_response, sanitized_content) + + # Return sanitized stream + if surface is _StreamSurface.ANTHROPIC_MESSAGES: + for sse_chunk in anthropic_sse_chunks_from_response(assembled_response): + yield sse_chunk return - except Exception as e: - verbose_proxy_logger.error("Model Armor streaming error: %s", str(e), exc_info=True) - if self.optional_params.get("fail_on_error", True): - raise - else: - verbose_proxy_logger.debug("Model Armor: No text content in streaming response, skipping guardrail") + mock_response: Final = MockResponseIterator(model_response=assembled_response) + async for chunk in mock_response: + yield chunk + return + + except ModelArmorAPIError as e: + if self.optional_params.get("fail_on_error", True): + for error_item in self._stream_error_items( + HTTPException(status_code=500, detail=e.detail), surface=surface + ): + yield error_item + return + except HTTPException as e: + # Yield the error as a terminal stream item so create_response() detects it and returns + # a proper JSON error response with the correct status code. Raising from a generator + # instead hits create_response's generic except and becomes a 500. + for error_item in self._stream_error_items(e, surface=surface): + yield error_item + return + except Exception as e: + verbose_proxy_logger.error("Model Armor streaming error: %s", str(e), exc_info=True) + if self.optional_params.get("fail_on_error", True): + raise # Return original chunks if no sanitization needed for chunk in all_chunks: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index da66c36328e..47089b7b1b1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -15,6 +15,7 @@ import litellm.types.utils from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError +from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ( @@ -3778,3 +3779,1153 @@ async def test_moderation_hook_skips_chat_traffic_when_configured_for_during_mcp assert result == data mock_post.assert_not_called() + + +_ANTHROPIC_SSE_CHUNKS = ( + b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message",' + b'"role":"assistant","model":"claude","content":[],"usage":{"input_tokens":5,"output_tokens":0}}}\n\n', + b'event: content_block_start\ndata: {"type":"content_block_start","index":0,' + b'"content_block":{"type":"text","text":""}}\n\n', + b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,' + b'"delta":{"type":"text_delta","text":"my card is 4111-1111-1111-1111"}}\n\n', + b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n', + b'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' + b'"usage":{"output_tokens":9}}\n\n', + b'event: message_stop\ndata: {"type":"message_stop"}\n\n', +) + +_MODEL_ARMOR_CLEAN = {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND"}} + +_MODEL_ARMOR_BLOCK = { + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "inspectResult": { + "matchState": "MATCH_FOUND", + "findings": [ + {"infoType": "CREDIT_CARD_NUMBER", "likelihood": "VERY_LIKELY"} + ], + } + } + } + }, + } +} + +# The root-level sanitizedText fallback in _get_sanitized_content, i.e. a rewrite that trips no +# named filter +_MODEL_ARMOR_SANITIZED = { + "sanitizedText": "my card is [REDACTED]", + "sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND"}, +} + +# The shape a real de-identify template returns: the SDP filter both matches and hands back the +# rewritten text, so whether it blocks or masks is decided by allow_sanitization alone +_MODEL_ARMOR_DEIDENTIFIED = { + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "deidentifyResult": { + "matchState": "MATCH_FOUND", + "data": {"text": "my card is [REDACTED]"}, + } + } + } + }, + } +} + + +def _chat_completion_chunks(): + """The chat-completions surface: typed ModelResponseStream chunks.""" + return ( + litellm.types.utils.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + index=0, + delta=litellm.types.utils.Delta(content="my card is 4111-1111-1111-1111"), + ) + ] + ), + litellm.types.utils.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + index=0, + delta=litellm.types.utils.Delta(content=""), + finish_reason="stop", + ) + ] + ), + ) + + +def _surface_guardrail(**kwargs): + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + **kwargs, + ) + guardrail._ensure_access_token_async = AsyncMock( + return_value=("test-token", "test-project") + ) + return guardrail + + +def _armor_post_mock(payload): + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json = AsyncMock(return_value=payload) + return AsyncMock(return_value=mock_response) + + +async def _anthropic_sse_stream(): + for chunk in _ANTHROPIC_SSE_CHUNKS: + yield chunk + + +def _responses_api_events(): + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + completed = ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "my card is 4111-1111-1111-1111"}], + } + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + return ( + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta="my card is 4111-1111-1111-1111", + ), + ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=completed, + ), + ) + + +async def _drain_surface_hook(guardrail, chunks, request_data=None): + async def _stream(): + for chunk in chunks: + yield chunk + + return [ + item + async for item in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=_stream(), + request_data=request_data + if request_data is not None + else { + "model": "claude-haiku", + "messages": [{"role": "user", "content": "show me a card"}], + "metadata": {"guardrails": ["model-armor-test"]}, + }, + ) + ] + + +@pytest.mark.asyncio +async def test_streaming_hook_scans_raw_anthropic_sse_instead_of_crashing(): + """A /v1/messages stream arrives as raw SSE bytes and must be assembled, then scanned. + + Regression for the 500 `Error building chunks for logging/streaming usage calculation`: + stream_chunk_builder calls .get() on each chunk, which raises on bytes. + """ + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, _ANTHROPIC_SSE_CHUNKS) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert "my card is 4111-1111-1111-1111" in scanned + assert tuple(delivered) == _ANTHROPIC_SSE_CHUNKS + + +@pytest.mark.asyncio +async def test_streaming_hook_scans_responses_api_events_instead_of_crashing(): + """A /v1/responses stream arrives as typed Responses events, which stream_chunk_builder + cannot subscript. The final response.completed event carries the text to scan.""" + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + events = _responses_api_events() + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, events) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert scanned == "my card is 4111-1111-1111-1111" + assert tuple(delivered) == events + + +@pytest.mark.asyncio +async def test_streaming_block_emits_anthropic_error_frame(): + """A block on /v1/messages must terminate the stream in Anthropic's error format. + + The OpenAI-shaped `data: {"error": ...}` frame the chat surface uses is rejected by + Anthropic clients. + """ + guardrail = _surface_guardrail() + + with patch.object( + guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_BLOCK) + ): + delivered = await _drain_surface_hook(guardrail, _ANTHROPIC_SSE_CHUNKS) + + body = b"".join(delivered) + assert b"event: error" in body + assert b'"type": "error"' in body + assert b"guardrail_error" in body + assert b"Streaming response blocked by Model Armor" in body + assert b"4111-1111-1111-1111" not in body + + +@pytest.mark.asyncio +async def test_streaming_block_emits_responses_api_error_event(): + """A block on /v1/responses must terminate the stream with a Responses ErrorEvent.""" + from litellm.types.llms.openai import ErrorEvent + + guardrail = _surface_guardrail() + + with patch.object( + guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_BLOCK) + ): + delivered = await _drain_surface_hook(guardrail, _responses_api_events()) + + assert len(delivered) == 1 + error_event = delivered[0] + assert isinstance(error_event, ErrorEvent) + assert error_event.error.type == "guardrail_error" + assert error_event.error.code == "400" + assert error_event.error.message == "Streaming response blocked by Model Armor" + + +@pytest.mark.asyncio +async def test_streaming_masking_re_emits_anthropic_sse_with_sanitized_text(): + """mask_response_content on /v1/messages must ship the sanitized text, not the original.""" + guardrail = _surface_guardrail(mask_response_content=True) + + with patch.object( + guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_SANITIZED) + ): + delivered = await _drain_surface_hook(guardrail, _ANTHROPIC_SSE_CHUNKS) + + body = b"".join(delivered) + assert b"[REDACTED]" in body + assert b"4111-1111-1111-1111" not in body + + +@pytest.mark.asyncio +async def test_streaming_masking_blocks_responses_api_stream(): + """A Responses event stream cannot be rebuilt from sanitized text, so releasing it would + ship the content the guardrail just rewrote. It is blocked instead.""" + from litellm.types.llms.openai import ErrorEvent + + guardrail = _surface_guardrail(mask_response_content=True) + + with patch.object( + guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_SANITIZED) + ): + delivered = await _drain_surface_hook(guardrail, _responses_api_events()) + + assert len(delivered) == 1 + assert isinstance(delivered[0], ErrorEvent) + assert delivered[0].error.code == "400" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface", ["anthropic_sse", "responses"]) +async def test_streaming_api_failure_frames_error_per_surface(surface): + """A Model Armor outage with fail_on_error must terminate the stream in the endpoint's + own error format rather than leaking an OpenAI SSE frame onto it.""" + from litellm.types.llms.openai import ErrorEvent + + guardrail = _surface_guardrail(fail_on_error=True) + chunks = _ANTHROPIC_SSE_CHUNKS if surface == "anthropic_sse" else _responses_api_events() + + mock_response = AsyncMock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + + with patch.object( + guardrail.async_handler, "post", AsyncMock(return_value=mock_response) + ): + delivered = await _drain_surface_hook(guardrail, chunks) + + assert len(delivered) >= 1 + if surface == "anthropic_sse": + assert b"event: error" in b"".join(delivered) + else: + assert isinstance(delivered[0], ErrorEvent) + assert delivered[0].error.code == "500" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chunks", + [ + pytest.param( + (b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,' + b'"delta":{"type":"text_delta","text":"hi"}}\n\n',), + id="anthropic-sse-without-message-start", + ), + pytest.param(None, id="responses-stream-without-completed-event"), + pytest.param("created", id="responses-stream-cut-off-after-response-created"), + ], +) +async def test_streaming_hook_fails_closed_when_a_surface_stream_cannot_be_assembled(chunks): + """Forwarding an unscannable /v1/messages or /v1/responses stream would silently disable the + guardrail, so the stream is refused in its own wire format instead of released unscanned.""" + from litellm.types.llms.openai import ( + ErrorEvent, + OutputTextDeltaEvent, + ResponsesAPIStreamEvents, + ) + + if chunks is None or chunks == "created": + delta = OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta="my card is 4111-1111-1111-1111", + ) + # response.created carries a ResponsesAPIResponse too, but an empty one: reading the body + # off it would scan "" and release every buffered delta unscanned + chunks = (delta,) if chunks is None else (_responses_created_event(), delta) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + assert tuple(delivered) != tuple(chunks) + if isinstance(chunks[0], bytes): + joined = b"".join(item.encode() if isinstance(item, str) else item for item in delivered).decode() + assert "event: error" in joined + assert "could not be assembled for scanning" in joined + return + assert len(delivered) == 1 + assert isinstance(delivered[0], ErrorEvent) + assert "could not be assembled for scanning" in delivered[0].error.message + + +@pytest.mark.asyncio +async def test_streaming_hook_forwards_a_preceding_guardrails_error_item(): + """A guardrail earlier in the post_call chain replaces the stream with its own terminal + error item. That item is not a chat delta, and feeding it to stream_chunk_builder is what + surfaced the ticket's 500, so it has to be forwarded untouched instead.""" + from litellm.types.llms.openai import ( + ErrorEvent, + ErrorEventError, + ResponsesAPIStreamEvents, + ) + + chunks = ( + ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, + sequence_number=1, + error=ErrorEventError( + type="guardrail_error", + code="400", + message="Streaming response blocked by Model Armor", + param=None, + ), + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + assert tuple(delivered) == chunks + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chunks", + [ + pytest.param(None, id="anthropic-error-event"), + pytest.param( + ('data: {"error": {"message": "Streaming response blocked by the first guardrail", "code": "400"}}\n\n',), + id="chat-completions-error-payload", + ), + ], +) +async def test_streaming_hook_forwards_a_preceding_guardrails_error_frame(chunks): + """Chained post_call guardrails hand each other their output. An earlier guardrail's error + frame carries no message to assemble, and replacing it would hide the real refusal.""" + if chunks is None: + chunks = anthropic_sse_error_frames("Streaming response blocked by the first guardrail") + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + assert tuple(delivered) == chunks + + +@pytest.mark.asyncio +async def test_streaming_responses_error_falls_back_to_sse_when_the_handler_declines(): + """build_stream_error_items may return None, which must not swallow the block into a clean + 200: the refusal falls back to the chat-completions SSE form that still carries the status.""" + from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ( + _StreamSurface, + ) + + class _DecliningGuardrail(ModelArmorGuardrail): + @staticmethod + def _build_responses_error_items(exc): + return None + + guardrail = _DecliningGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + exc = HTTPException(status_code=400, detail={"message": "blocked"}) + + items = guardrail._stream_error_items(exc, surface=_StreamSurface.RESPONSES) + + assert len(items) == 1 + assert '"code": "400"' in items[0] + assert "blocked" in items[0] + + +def _responses_created_event(): + from litellm.types.llms.openai import ( + ResponseCreatedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + return ResponseCreatedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_CREATED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + + +@pytest.mark.asyncio +async def test_streaming_hook_refuses_an_opaque_sse_stream_without_anthropic_framing(): + """/v1/messages is not the only endpoint that streams raw SSE: the Google generateContent + route marks its own stream raw too. Its frames carry no Anthropic event types, so refusing + them in Anthropic's format would hand a Google client a body it cannot parse.""" + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + chunks = (b'data: {"candidates":[{"content":{"parts":[{"text":"my card is 4111"}]}}]}\n\n',) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + assert tuple(delivered) != chunks + body = "".join(item.decode() if isinstance(item, bytes) else item for item in delivered) + assert "could not be assembled for scanning" in body + assert "event: error" not in body + assert '"code": "500"' in body + + +@pytest.mark.asyncio +async def test_streaming_unassemblable_stream_is_forwarded_when_fail_on_error_is_disabled(): + """fail_on_error: false is a deliberate choice to degrade open, and it governs every other + path in this hook. The fail-closed refusal has to honour it too.""" + guardrail = _surface_guardrail(fail_on_error=False) + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + chunks = ( + b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,' + b'"delta":{"type":"text_delta","text":"hi"}}\n\n', + ) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + assert tuple(delivered) == chunks + + +@pytest.mark.asyncio +async def test_streaming_fail_closed_records_the_applied_guardrail(): + """A refusal that no header or log attributes to the guardrail leaves on-call unable to tell + a guardrail block apart from a provider failure.""" + guardrail = _surface_guardrail() + request_data = { + "model": "claude-haiku", + "messages": [{"role": "user", "content": "show me a card"}], + "metadata": {"guardrails": ["model-armor-test"]}, + } + chunks = ( + b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,' + b'"delta":{"type":"text_delta","text":"hi"}}\n\n', + ) + + with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_CLEAN)): + await _drain_surface_hook(guardrail, chunks, request_data=request_data) + + assert request_data["metadata"]["applied_guardrails"] == ["model-armor-test"] + + +@pytest.mark.asyncio +async def test_streaming_responses_tool_call_output_is_scanned(): + """An agentic /v1/responses turn can carry its whole payload in tool-call arguments, which + is what the chat surface already folds into the scanned text.""" + from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[ + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "send_email", + "arguments": '{"body": "my card is 4111-1111-1111-1111"}', + } + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, (completed,)) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert "4111-1111-1111-1111" in scanned + assert tuple(delivered) == (completed,) + + +@pytest.mark.asyncio +async def test_streaming_hook_refuses_a_content_stream_that_ends_with_an_error_frame(): + """The chain-aware passthrough must stay narrow. A stream carrying real content plus a + trailing error frame is not a bare refusal to forward: the assembler cannot read it, and + releasing it would ship the buffered content unscanned.""" + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + chunks = (*_ANTHROPIC_SSE_CHUNKS, *anthropic_sse_error_frames("upstream gave up")) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + body = b"".join(delivered) + assert b"4111-1111-1111-1111" not in body + assert b"could not be assembled for scanning" in body + + +@pytest.mark.parametrize( + "chunks, expected, case", + [ + (anthropic_sse_error_frames("blocked upstream"), True, "anthropic-error-frames-only"), + ((f"data: {json.dumps({'error': {'message': 'blocked'}})}\n\n",), True, "chat-error-payload-only"), + ((), False, "empty-stream"), + ( + (b'event: message_delta\ndata: {"type":"message_delta","error":null}\n\n',), + False, + "content-event-carrying-a-null-error-field", + ), + ( + ( + litellm.types.utils.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + index=0, + delta=litellm.types.utils.Delta(content="my card is 4111-1111-1111-1111"), + ) + ] + ), + *anthropic_sse_error_frames("upstream gave up"), + ), + False, + "typed-content-chunks-plus-a-trailing-error-frame", + ), + ], +) +def test_is_sse_error_stream_only_matches_a_stream_that_is_nothing_but_refusals(chunks, expected, case): + """The chain-aware passthrough turns on this predicate, so anything it calls error-only is + forwarded to the client untouched. A stream that still carries content must not qualify: the + frames-only join drops typed chunks, and a content event may carry an empty ``error`` field.""" + from litellm.proxy.guardrails.anthropic_sse import is_sse_error_stream + + assert is_sse_error_stream(chunks) is expected, case + + +@pytest.mark.asyncio +async def test_streaming_hook_does_not_forward_typed_chunks_that_end_with_an_error_frame(): + """A stream mixing buffered content with a trailing refusal is not the bare refusal the chain + passthrough exists for. Forwarding it would release the content no scanner ever saw.""" + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + chunks = ( + litellm.types.utils.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + index=0, + delta=litellm.types.utils.Delta(content="my card is 4111-1111-1111-1111"), + ) + ] + ), + *anthropic_sse_error_frames("upstream gave up"), + ) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_not_called() + body = b"".join(item if isinstance(item, bytes) else str(item).encode() for item in delivered) + assert b"4111-1111-1111-1111" not in body + assert b"could not be assembled for scanning" in body + + +def _delivered_bytes(delivered): + return b"".join( + item + if isinstance(item, bytes) + else item.encode() + if isinstance(item, str) + else str(item.model_dump() if hasattr(item, "model_dump") else item).encode() + for item in delivered + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunks, case", [(None, "chat_completions"), (_ANTHROPIC_SSE_CHUNKS, "anthropic_sse")]) +async def test_streaming_deidentify_match_masks_when_masking_is_enabled(chunks, case): + """A de-identify template reports MATCH_FOUND for every redaction it makes, so reading that + match as a refusal makes mask_response_content unusable on a stream: the client gets an error + where its non-streaming sibling gets redacted text. The block check has to allow sanitization + exactly as the non-streaming hook does.""" + guardrail = _surface_guardrail(mask_response_content=True) + post = _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook( + guardrail, _chat_completion_chunks() if chunks is None else chunks + ) + + body = _delivered_bytes(delivered) + assert b"[REDACTED]" in body, case + assert b"4111-1111-1111-1111" not in body, case + assert b"blocked by Model Armor" not in body, case + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunks, case", [(None, "chat_completions"), (_ANTHROPIC_SSE_CHUNKS, "anthropic_sse")]) +async def test_streaming_deidentify_match_still_blocks_when_masking_is_disabled(chunks, case): + """Without mask_response_content there is nowhere to put the rewritten text, so the same + de-identify match must still end the stream rather than release the original.""" + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook( + guardrail, _chat_completion_chunks() if chunks is None else chunks + ) + + body = _delivered_bytes(delivered) + assert b"Streaming response blocked by Model Armor" in body, case + assert b"4111-1111-1111-1111" not in body, case + + +@pytest.mark.asyncio +async def test_streaming_deidentify_match_logs_masked_run_as_success_not_blocked(): + """The status stamped on request metadata feeds the spend log, so it has to agree with what + the client actually received: a masked stream is a success, not a block.""" + guardrail = _surface_guardrail(mask_response_content=True) + request_data = { + "model": "claude-haiku", + "messages": [{"role": "user", "content": "show me a card"}], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED)): + await _drain_surface_hook(guardrail, _chat_completion_chunks(), request_data=request_data) + + assert request_data["metadata"]["_model_armor_status"] == "success" + + +# A de-identify template that matched but handed back no rewrite, e.g. because the transformation +# itself failed. The match still says the buffered original carries what it matched on +_MODEL_ARMOR_DEIDENTIFIED_NO_TEXT = { + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": {"sdpFilterResult": {"deidentifyResult": {"matchState": "MATCH_FOUND"}}} + }, + } +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunks, case", [(None, "chat_completions"), (_ANTHROPIC_SSE_CHUNKS, "anthropic_sse")]) +async def test_streaming_deidentify_match_without_a_rewrite_fails_closed(chunks, case): + """Allowing sanitization past the block check is a promise to apply the redaction. When Model + Armor matches but returns no sanitized text there is nothing to apply, and yielding the + buffered chunks would hand back exactly what it matched on.""" + guardrail = _surface_guardrail(mask_response_content=True) + post = _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED_NO_TEXT) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook( + guardrail, _chat_completion_chunks() if chunks is None else chunks + ) + + body = _delivered_bytes(delivered) + assert b"4111-1111-1111-1111" not in body, case + assert b"Streaming response blocked by Model Armor" in body, case + + +@pytest.mark.asyncio +async def test_streaming_status_records_a_surface_that_cannot_carry_the_rewrite_as_blocked(): + """The Responses surface has no assembled body to rewrite, so a de-identify match ends as a + refusal. The status stamped on metadata feeds the spend log and has to say so rather than + reporting the success the block check alone would have implied.""" + guardrail = _surface_guardrail(mask_response_content=True) + request_data = { + "model": "gpt-4o-mini", + "input": "show me a card", + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED)): + delivered = await _drain_surface_hook( + guardrail, _responses_api_events(), request_data=request_data + ) + + body = _delivered_bytes(delivered) + assert b"4111-1111-1111-1111" not in body + assert b"Streaming response blocked by Model Armor" in body + assert request_data["metadata"]["_model_armor_status"] == "blocked" + + +def _responses_api_events_truncated(terminal: str): + """A /v1/responses stream whose text went out as deltas and whose terminal event reports no body. + + ``response.failed`` and ``response.incomplete`` are terminal like ``response.completed``, but a + turn that broke mid-generation reports an empty ``output`` while the deltas ahead of it already + spelled the answer out to the client. + """ + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + empty_body = ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) + terminal_event = ( + ResponseFailedEvent(type=ResponsesAPIStreamEvents.RESPONSE_FAILED, response=empty_body) + if terminal == "failed" + else ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=empty_body) + ) + return ( + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta="my card is 4111-1111-1111-1111", + ), + terminal_event, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("terminal", ["failed", "incomplete"]) +async def test_streaming_responses_terminal_event_without_a_body_still_scans_the_deltas(terminal): + """A /v1/responses turn that broke mid-generation has still delivered its deltas. + + Reading only the terminal body would find nothing to scan and hand every buffered delta to the + client untouched, so the deltas themselves are what gets scanned. + """ + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_BLOCK) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, _responses_api_events_truncated(terminal)) + + post.assert_called_once() + assert "4111-1111-1111-1111" in post.call_args.kwargs["json"]["modelResponseData"]["text"] + rendered = "".join(str(item) for item in delivered) + assert "4111-1111-1111-1111" not in rendered + assert "Streaming response blocked by Model Armor" in rendered + + +@pytest.mark.asyncio +async def test_streaming_responses_mcp_argument_deltas_are_scanned_when_the_body_is_empty(): + """A turn that only streamed MCP tool arguments still handed the client a payload. + + The delta fallback is read off the event enum rather than listed by hand, so an argument event + that carries no `output_text` cannot fall out of the scan. + """ + from litellm.types.llms.openai import ( + MCPCallArgumentsDeltaEvent, + ResponseIncompleteEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + chunks = ( + MCPCallArgumentsDeltaEvent( + type=ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA, + output_index=0, + item_id="mcp_1", + delta='{"note": "my card is 4111-1111-1111-1111"}', + sequence_number=0, + ), + ResponseIncompleteEvent( + type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_BLOCK) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, chunks) + + post.assert_called_once() + assert "4111-1111-1111-1111" in post.call_args.kwargs["json"]["modelResponseData"]["text"] + rendered = "".join(str(item) for item in delivered) + assert "4111-1111-1111-1111" not in rendered + assert "Streaming response blocked by Model Armor" in rendered + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "with_output_text_delta", + [True, False], + ids=["summary-and-text-deltas", "summary-delta-only"], +) +async def test_streaming_responses_reasoning_summary_deltas_are_scanned_alongside_the_body(with_output_text_delta): + """A reasoning turn streams its summary in deltas the terminal body never repeats. + + Reading only the body scans the visible answer and hands the client every summary delta + unscanned, so the body and the deltas are scanned together. + """ + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ReasoningSummaryTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + answer = "the weather is fine" + summary_delta = ReasoningSummaryTextDeltaEvent( + type=ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA, + item_id="rs_1", + output_index=0, + delta="the user said my card is 4111-1111-1111-1111", + ) + text_deltas = ( + ( + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=1, + content_index=0, + delta=answer, + ), + ) + if with_output_text_delta + else () + ) + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-5-mini", + object="response", + output=[ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": answer, "annotations": []}], + } + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_BLOCK) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, (summary_delta, *text_deltas, completed)) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert "4111-1111-1111-1111" in scanned + assert answer in scanned + assert scanned.count(answer) == 1 + rendered = "".join(str(item) for item in delivered) + assert "4111-1111-1111-1111" not in rendered + assert "Streaming response blocked by Model Armor" in rendered + + +@pytest.mark.asyncio +async def test_streaming_responses_deltas_of_separate_fields_do_not_form_a_finding_across_their_boundary(): + """Two fields of a turn are separate text, so what runs across their boundary is not model output. + + A reasoning summary ending in half a card number and an answer opening with the other half + each carry nothing to find, and joining them without a break would invent one. + """ + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ReasoningSummaryTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + answer = "1111-1111 is not a full card" + summary_delta = ReasoningSummaryTextDeltaEvent( + type=ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA, + item_id="rs_1", + output_index=0, + delta="the prefix they gave me is 4111-1111-", + ) + text_delta = OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=1, + content_index=0, + delta=answer, + ) + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-5-mini", + object="response", + output=[ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": answer, "annotations": []}], + } + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, (summary_delta, text_delta, completed)) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert "4111-1111-" in scanned + assert answer in scanned + assert "4111-1111-1111-1111" not in scanned + rendered = "".join(str(item) for item in delivered) + assert "Streaming response blocked by Model Armor" not in rendered + + +@pytest.mark.asyncio +async def test_streaming_responses_one_fields_deltas_still_join_into_a_single_finding(): + """A card number split across two deltas of one field is still one card number to scan.""" + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + halves = ("my card is 4111-1111-", "1111-1111") + text_deltas = tuple( + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta=half, + ) + for half in halves + ) + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-5-mini", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_BLOCK) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, (*text_deltas, completed)) + + post.assert_called_once() + assert "4111-1111-1111-1111" in post.call_args.kwargs["json"]["modelResponseData"]["text"] + rendered = "".join(str(item) for item in delivered) + assert "4111-1111-1111-1111" not in rendered + assert "Streaming response blocked by Model Armor" in rendered + + +@pytest.mark.asyncio +async def test_streaming_responses_fields_the_body_repeats_are_not_scanned_a_second_time(): + """A turn whose visible fields all reach the terminal body is scanned once, not twice. + + Two output_text fields stream as deltas and come back in the completed body, so scanning the + deltas on top of the body would send Model Armor two copies of everything the client sees. + """ + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + paragraphs = ("the first thing to know", "a second and separate point") + text_deltas = tuple( + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id=f"msg_{index}", + output_index=index, + content_index=0, + delta=paragraph, + ) + for index, paragraph in enumerate(paragraphs) + ) + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-5-mini", + object="response", + output=[ + { + "type": "message", + "id": f"msg_{index}", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": paragraph, "annotations": []}], + } + for index, paragraph in enumerate(paragraphs) + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, (*text_deltas, completed)) + + post.assert_called_once() + scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"] + assert [scanned.count(paragraph) for paragraph in paragraphs] == [1, 1] + rendered = "".join(str(item) for item in delivered) + assert all(paragraph in rendered for paragraph in paragraphs) + + +def test_every_responses_delta_event_is_in_the_scanned_set(): + """Every ``.delta`` the Responses event enum defines is model output on its way to the client.""" + from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ( + _RESPONSES_DELTA_EVENT_TYPES, + ) + from litellm.types.llms.openai import ResponsesAPIStreamEvents + + missing = { + event.value + for event in ResponsesAPIStreamEvents + if event.value.endswith(".delta") and event.value not in _RESPONSES_DELTA_EVENT_TYPES + } + assert not missing + assert "response.mcp_call_arguments.delta" in _RESPONSES_DELTA_EVENT_TYPES From a5639b8e2a9870fb0aeb7d2365a8a386ed4e0fe6 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:22:52 -0700 Subject: [PATCH 18/23] test: add OCR python-to-rust test parity ledger (WIP) (#39434) * test: add OCR python-to-rust test parity ledger * feat: add ledger loader for OCR test parity data * feat: add drift audit for OCR test parity ledger (WIP, untested) * fix: correct drift in OCR test-parity ledger Two entries referenced a typo'd Python test name, four duplicated entries already tracked under TestProxySecurityGuard, five real Python tests in test_rust_bridge.py were untracked, and three real Rust custom_logger tests were missing from rust_only_tests. Found by running validate_ledger.py's audit against the live repo. * test: add regression coverage for the OCR ledger and audit script Covers schema validation, AST/regex test enumeration, drift detection on both the Python and Rust sides, LedgerDriftError content, and a live-repo clean-audit guard against future drift. * test: simplify ledger test to one drift-guard assertion Replace the ledger-internals unit tests with a single test that runs the real audit against the live repo and asserts every OCR test is accounted for (mapped, unmapped-with-reason, or rust_only), printing the exact diff on failure. * chore: move OCR test-parity ledger to core/ocr validate_sub_methods/ mixes strategy-catalog metadata with the ledger. Ledger data belongs under a per-function core// path instead. * fix: point LEDGER_PATH at the new core/ocr location --- tests/rust-python-harness/cli.py | 48 +++- tests/rust-python-harness/shared/__init__.py | 0 .../shared/parity/__init__.py | 0 .../shared/parity/ledger.py | 136 ++++++++++++ .../strategies/__init__.py | 0 .../strategies/unit_tests/__init__.py | 0 .../ledgers/ocr/ocr_test_ledger.json | 208 ++++++++++++++++++ .../unit_tests/mapping_validator.py | 101 +++++++++ .../strategies/unit_tests/python_runner.py | 22 ++ .../strategies/unit_tests/rust_runner.py | 13 ++ tests/test_rust_python_harness.py | 55 +++++ 11 files changed, 581 insertions(+), 2 deletions(-) create mode 100644 tests/rust-python-harness/shared/__init__.py create mode 100644 tests/rust-python-harness/shared/parity/__init__.py create mode 100644 tests/rust-python-harness/shared/parity/ledger.py create mode 100644 tests/rust-python-harness/strategies/__init__.py create mode 100644 tests/rust-python-harness/strategies/unit_tests/__init__.py create mode 100644 tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json create mode 100644 tests/rust-python-harness/strategies/unit_tests/mapping_validator.py create mode 100644 tests/rust-python-harness/strategies/unit_tests/python_runner.py create mode 100644 tests/rust-python-harness/strategies/unit_tests/rust_runner.py diff --git a/tests/rust-python-harness/cli.py b/tests/rust-python-harness/cli.py index 41c46f9613a..f9e97d7ad43 100644 --- a/tests/rust-python-harness/cli.py +++ b/tests/rust-python-harness/cli.py @@ -9,9 +9,11 @@ from .catalog import load_catalog from .models import HarnessCase, Strategy from .runner import run_pytest from .ui import make_dashboard +from .strategies.unit_tests.mapping_validator import FunctionReport, build_function_report REPO_ROOT = Path(__file__).resolve().parents[2] COVERAGE_ROOT = REPO_ROOT / "target" / "rust-python-harness" +SDK_FUNCTION_CHOICES = ("ocr", "messages", "responses", "count_tokens") def _parser() -> argparse.ArgumentParser: @@ -40,9 +42,17 @@ def _parser() -> argparse.ArgumentParser: action="append", default=[], dest="sdk_functions", - choices=("ocr", "messages", "responses", "count_tokens"), + choices=SDK_FUNCTION_CHOICES, help="run only this SDK function", ) + parser.add_argument( + "--validate-ledger", + action="store_true", + help=( + "report Python<->Rust test-parity ledger gaps and drift instead of " + "running the dashboard; narrow with --function" + ), + ) parser.add_argument( "--plain", action="store_true", @@ -100,7 +110,7 @@ def _interactive_filters(strategies: Sequence[Strategy]) -> tuple[set[str], set[ ) sdk_functions = _pick_values( "SDK functions", - [(name, name) for name in ("ocr", "messages", "responses", "count_tokens")], + [(name, name) for name in SDK_FUNCTION_CHOICES], ) return strategy_ids, sdk_functions @@ -131,6 +141,38 @@ def _print_catalog(strategies: Sequence[Strategy]) -> None: print(f" {case.sdk_function:12} {case.coverage.value:14} {selectors}") +def _print_function_report(report: FunctionReport) -> None: + print(f"\n{report.sdk_function}") + if report.ledger is None or report.audit is None: + print(" no ledger yet") + return + ledger, audit = report.ledger, report.audit + print( + f" {ledger.mapped_count}/{ledger.total_count} python tests mapped to rust " + f"({ledger.percentage}%)" + ) + print(f" {len(ledger.rust_only_tests)} rust-only tests with no python counterpart") + if audit.is_clean: + print(" ledger is in sync with the live test files") + return + for label, items in ( + ("ledger references a python test that no longer exists", audit.missing_python_tests), + ("python test exists but is not tracked in the ledger", audit.stale_python_tests), + ("ledger references a rust test that no longer exists", audit.missing_rust_tests), + ("rust test exists but is not tracked in the ledger", audit.stale_rust_tests), + ): + for item in items: + print(f" {label}: {item}") + + +def _validate_ledger(sdk_functions: set[str]) -> int: + functions = sdk_functions or set(SDK_FUNCTION_CHOICES) + reports = tuple(build_function_report(function) for function in sorted(functions)) + for report in reports: + _print_function_report(report) + return 0 if all(report.is_clean for report in reports) else 1 + + def main(argv: Sequence[str] | None = None) -> int: args = _parser().parse_args(argv) if args.coverage and importlib.util.find_spec("pytest_cov") is None: @@ -138,6 +180,8 @@ def main(argv: Sequence[str] | None = None) -> int: "--coverage requires the project's pytest-cov dependency; run with " "`poetry run python -m tests.rust-python-harness --coverage`" ) + if args.validate_ledger: + return _validate_ledger(set(args.sdk_functions)) strategies = load_catalog() if args.list: _print_catalog(strategies) diff --git a/tests/rust-python-harness/shared/__init__.py b/tests/rust-python-harness/shared/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/rust-python-harness/shared/parity/__init__.py b/tests/rust-python-harness/shared/parity/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/rust-python-harness/shared/parity/ledger.py b/tests/rust-python-harness/shared/parity/ledger.py new file mode 100644 index 00000000000..40dfed583ae --- /dev/null +++ b/tests/rust-python-harness/shared/parity/ledger.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +@dataclass(frozen=True, slots=True) +class LedgerEntry: + python_file: str + python_test: str + status: str + rust_file: str + rust_test: str + justification: str + reason: str + + +@dataclass(frozen=True, slots=True) +class RustOnlyEntry: + rust_file: str + rust_test: str + reason: str + + +@dataclass(frozen=True, slots=True) +class TestLedger: + sdk_function: str + python_scope: tuple[str, ...] + rust_scope: tuple[str, ...] + entries: tuple[LedgerEntry, ...] + rust_only_tests: tuple[RustOnlyEntry, ...] + + @property + def mapped_count(self) -> int: + return sum(1 for entry in self.entries if entry.status == "mapped") + + @property + def total_count(self) -> int: + return len(self.entries) + + @property + def percentage(self) -> float: + if self.total_count == 0: + return 0.0 + return round(100.0 * self.mapped_count / self.total_count, 1) + + +def _require_string(value: Any, field: str, source: Path) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{source}: {field} must be a non-empty string") + return value + + +def _require_string_list(value: Any, field: str, source: Path) -> tuple[str, ...]: + if not isinstance(value, list) or not all(isinstance(item, str) and item for item in value): + raise ValueError(f"{source}: {field} must be a list of non-empty strings") + return tuple(value) + + +def _load_entry(data: Any, index: int, source: Path) -> LedgerEntry: + if not isinstance(data, dict): + raise ValueError(f"{source}: entries[{index}] must be an object") + python_file = _require_string(data.get("python_file"), f"entries[{index}].python_file", source) + python_test = _require_string(data.get("python_test"), f"entries[{index}].python_test", source) + status = data.get("status") + if status not in ("mapped", "unmapped"): + raise ValueError(f"{source}: entries[{index}].status must be 'mapped' or 'unmapped'") + + if status == "mapped": + rust_file = _require_string(data.get("rust_file"), f"entries[{index}].rust_file", source) + rust_test = _require_string(data.get("rust_test"), f"entries[{index}].rust_test", source) + justification = _require_string( + data.get("justification"), f"entries[{index}].justification", source + ) + return LedgerEntry( + python_file=python_file, + python_test=python_test, + status=status, + rust_file=rust_file, + rust_test=rust_test, + justification=justification, + reason="", + ) + + reason = _require_string(data.get("reason"), f"entries[{index}].reason", source) + return LedgerEntry( + python_file=python_file, + python_test=python_test, + status=status, + rust_file="", + rust_test="", + justification="", + reason=reason, + ) + + +def _load_rust_only_entry(data: Any, index: int, source: Path) -> RustOnlyEntry: + if not isinstance(data, dict): + raise ValueError(f"{source}: rust_only_tests[{index}] must be an object") + return RustOnlyEntry( + rust_file=_require_string(data.get("rust_file"), f"rust_only_tests[{index}].rust_file", source), + rust_test=_require_string(data.get("rust_test"), f"rust_only_tests[{index}].rust_test", source), + reason=_require_string(data.get("reason"), f"rust_only_tests[{index}].reason", source), + ) + + +def load_ledger(path: Path) -> TestLedger: + with path.open(encoding="utf-8") as stream: + data = json.load(stream) + + sdk_function = _require_string(data.get("sdk_function"), "sdk_function", path) + python_scope = _require_string_list(data.get("python_scope"), "python_scope", path) + rust_scope = _require_string_list(data.get("rust_scope"), "rust_scope", path) + + entries_data = data.get("entries") + if not isinstance(entries_data, list): + raise ValueError(f"{path}: entries must be a list") + entries = tuple( + _load_entry(entry, index, path) for index, entry in enumerate(entries_data) + ) + + rust_only_data = data.get("rust_only_tests") + if not isinstance(rust_only_data, list): + raise ValueError(f"{path}: rust_only_tests must be a list") + rust_only_tests = tuple( + _load_rust_only_entry(entry, index, path) for index, entry in enumerate(rust_only_data) + ) + + return TestLedger( + sdk_function=sdk_function, + python_scope=python_scope, + rust_scope=rust_scope, + entries=entries, + rust_only_tests=rust_only_tests, + ) diff --git a/tests/rust-python-harness/strategies/__init__.py b/tests/rust-python-harness/strategies/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/rust-python-harness/strategies/unit_tests/__init__.py b/tests/rust-python-harness/strategies/unit_tests/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json b/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json new file mode 100644 index 00000000000..799a1320463 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json @@ -0,0 +1,208 @@ +{ + "sdk_function": "ocr", + "python_scope": [ + "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "tests/test_litellm/ocr/test_rust_bridge.py", + "tests/test_litellm/ocr/test_ocr_file_input.py", + "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "tests/test_litellm/ocr/test_ocr_native_format.py", + "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py" + ], + "rust_scope": [ + "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", + "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs" + ], + "entries": [ + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_should_encode_azure_document_intelligence_model_id", "status": "unmapped", "reason": "model-id URL percent-encoding has no Rust test; Rust only tests pages/features query building"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_should_reject_dot_segment_azure_document_intelligence_model_id", "status": "unmapped", "reason": "model-id dot-segment validation has no Rust test"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_preserves_azure_native_fields", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_response_normalizes_pages", "justification": "both assert page markdown, dimension (inch-to-pixel) normalization, and usage_info.pages_processed from the same Azure succeeded response shape"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_async_transform_ocr_response_preserves_azure_native_fields", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_response_normalizes_pages", "justification": "async twin of the sync case above, same underlying transform is exercised on the Rust side"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_tolerates_missing_native_fields", "status": "unmapped", "reason": "tables/keyValuePairs absence tolerance is not asserted by the Rust response test"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_non_succeeded_status_raises", "status": "unmapped", "reason": "no Rust test asserts on a non-succeeded Azure DI status"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_supported_ocr_params_includes_features", "status": "unmapped", "reason": "supported-params list content has no Rust equivalent for Azure"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_native_format_carries_raw_operation", "status": "unmapped", "reason": "native req_format raw-operation passthrough is not tested in Rust"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_async_transform_ocr_response_native_format_carries_raw_operation", "status": "unmapped", "reason": "native req_format raw-operation passthrough is not tested in Rust"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_default_format_omits_raw_operation", "status": "unmapped", "reason": "req_format gating of raw-operation output has no Rust test"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_passes_through_req_format", "status": "unmapped", "reason": "req_format passthrough in map_ocr_params has no Rust test"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_rejects_unknown_req_format_as_bad_request", "status": "unmapped", "reason": "req_format validation error path has no Rust test"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_omits_req_format_query_param", "status": "unmapped", "reason": "no Rust test asserts req_format is excluded from the built URL"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_features", "status": "unmapped", "reason": "features-string normalization in map_ocr_params has no Rust test"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_empty_features_list_omitted", "status": "unmapped", "reason": "empty-features omission has no Rust test"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_invalid_features_raises", "status": "unmapped", "reason": "features validation error path has no Rust test"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_appends_features_query", "status": "unmapped", "reason": "features query-param construction has no Rust test"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_combines_pages_and_features", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_url_normalizes_zero_based_pages", "justification": "both assert 0-based, duplicate page indices are deduped, sorted, and rewritten 1-based into the request URL"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_validate_environment_uses_subscription_key", "status": "unmapped", "reason": "Python-side header derivation from litellm_params; Rust's poll test only checks the header is present, not how it was resolved"}, + {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_validate_environment_falls_back_to_entra_token", "status": "unmapped", "reason": "Entra bearer-token fallback logic has no Rust test"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_doc_intelligence_route", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_documentintelligence_and_is_case_insensitive", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_does_not_match_mistral_route", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_does_not_hijack_doc_intelligence", "status": "unmapped", "reason": "api_base resolution from the secret manager runs before the Rust bridge is called, no Rust test exists for it"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_explicit_api_base_is_honoured_for_doc_intelligence", "status": "unmapped", "reason": "api_base precedence resolution is Python-only"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_still_applies_to_mistral_ocr", "status": "unmapped", "reason": "api_base precedence resolution is Python-only"}, + + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_use_litellm_rust_toggles_flag", "status": "unmapped", "reason": "bridge-plumbing: Python-side feature-flag toggle, no Rust equivalent"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_env_var_enables_rust_ocr", "status": "unmapped", "reason": "bridge-plumbing: Python-side env-var flag gating"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_returns_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: dependency-injection test hook, not provider behavior"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_loader_returns_none_when_extension_absent", "status": "unmapped", "reason": "bridge-plumbing: native-extension import/loader fallback"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_loader_caches_absent_extension", "status": "unmapped", "reason": "bridge-plumbing: loader caching behavior"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_available_reflects_loader", "status": "unmapped", "reason": "bridge-plumbing: loader availability check"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_aocr_returns_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: dependency-injection test hook"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_toggle_without_ocr_arg_preserves_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: injected-impl state retention regression"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_explicit_ocr_none_clears_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: injected-impl clearing behavior"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_none_when_extension_absent", "status": "unmapped", "reason": "bridge-plumbing: degrade path when the native extension is missing"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_uses_compiled_extension", "status": "unmapped", "reason": "bridge-plumbing: native module resolution"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_timeout_to_seconds_handles_float_timeout_and_none", "status": "unmapped", "reason": "bridge-plumbing: Python-side timeout normalization helper"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_bridge_wrapper_forwards_prepared_args_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: wrapper argument forwarding, asserted against a fake bridge not the real Rust code"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: async wrapper argument forwarding"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_prepares_request_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: request preparation and response wrapping in Python"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_resolves_key_via_secret_manager_when_missing", "status": "unmapped", "reason": "secret-manager: API key resolution happens in Python before the bridge is invoked"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_prefers_explicit_key_over_resolver", "status": "unmapped", "reason": "secret-manager: key precedence resolution"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_uses_provider_api_key_env_var", "status": "unmapped", "reason": "secret-manager: provider-specific env var name resolution"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_forwards_vertex_routing_metadata", "status": "unmapped", "reason": "secret-manager: vertex routing metadata merge happens in Python"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager", "status": "unmapped", "reason": "secret-manager: vertex project/location resolution"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager", "status": "unmapped", "reason": "secret-manager: azure_ai api_base resolution"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint", "status": "unmapped", "reason": "secret-manager: doc-intelligence endpoint resolution"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_runs_pre_call_logging", "status": "unmapped", "reason": "bridge-plumbing: Python logging-object pre_call invocation"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_routes_to_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: routing to a fake bridge, not the real Rust transform"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_routes_azure_ai_to_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: provider-prefix stripping before routing"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_rust_path_converts_file_document_before_bridge", "status": "unmapped", "reason": "file-normalization: raw-bytes-to-data-URI conversion happens in Python before the bridge call"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_exception_type_uses_resolved_provider_context", "status": "unmapped", "reason": "bridge-plumbing: Python exception-type mapping on bridge failure"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_aocr_routes_to_async_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: async routing to a fake bridge"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_aocr_exception_type_uses_resolved_provider_context", "status": "unmapped", "reason": "bridge-plumbing: async exception-type mapping on bridge failure"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_does_not_route_to_rust_when_disabled", "status": "unmapped", "reason": "bridge-plumbing: Python control flow for the toggle-disabled branch, no Rust-owned behavior runs"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_falls_back_to_python_when_bridge_unavailable", "status": "unmapped", "reason": "bridge-plumbing: Python-only fallback when the compiled Rust extension is absent, Rust cannot test its own absence"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_forwards_timeout_to_rust", "status": "unmapped", "reason": "bridge-plumbing: asserts the Python call site forwards a timeout kwarg, Rust receives an already-constructed request"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_passes_default_request_timeout_to_rust", "status": "unmapped", "reason": "bridge-plumbing: asserts the Python call site supplies a default timeout kwarg, no Rust equivalent"}, + {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_provider_configs_expose_api_key_env_vars", "status": "unmapped", "reason": "asserts per-provider get_api_key_env_var() strings; the closest Rust test (ocr_dispatch_supports_migrated_providers) asserts provider dispatch/param resolution instead, not API key env var names"}, + + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_pdf_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection is Python-only preprocessing before the bridge call"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_png_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_jpg_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_jpeg_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_gif_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_webp_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_tiff_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_tif_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_bmp_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_be_case_insensitive", "status": "unmapped", "reason": "file-normalization: MIME detection case handling"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_fallback_for_unknown_extension", "status": "unmapped", "reason": "file-normalization: MIME detection fallback"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pdf_pathlib_path_to_document_url", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion happens in Python"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_image_pathlib_path_to_image_url", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_reject_bare_str_path", "status": "unmapped", "reason": "file-normalization: arbitrary-file-read guard on bare str paths"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pathlib_path", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes", "status": "unmapped", "reason": "file-normalization: raw-bytes-to-data-URI conversion"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_explicit_mime_type", "status": "unmapped", "reason": "file-normalization: explicit MIME override on raw bytes"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_image_mime_type", "status": "unmapped", "reason": "file-normalization: explicit MIME override on raw bytes"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object", "status": "unmapped", "reason": "file-normalization: file-like-object conversion"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object_with_name", "status": "unmapped", "reason": "file-normalization: file-like-object name-based MIME detection"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_missing_file_field", "status": "unmapped", "reason": "file-normalization: missing-field validation"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_nonexistent_pathlib_path", "status": "unmapped", "reason": "file-normalization: missing-file validation"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_empty_file", "status": "unmapped", "reason": "file-normalization: empty-file validation"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_unsupported_type", "status": "unmapped", "reason": "file-normalization: unsupported input type validation"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_invalid_mime_type", "status": "unmapped", "reason": "file-normalization: MIME-type injection validation"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_override_mime_type_for_pathlib_path", "status": "unmapped", "reason": "file-normalization: explicit MIME override precedence"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_document_url_for_pdf", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_png", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_jpeg", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_octet_stream", "status": "unmapped", "reason": "file-normalization: filename-based MIME fallback"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_none", "status": "unmapped", "reason": "file-normalization: filename-based MIME fallback"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_fallback_to_octet_stream_for_unknown", "status": "unmapped", "reason": "file-normalization: default MIME fallback"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_preserve_base64_content_correctly", "status": "unmapped", "reason": "file-normalization: binary round-trip through base64"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_from_content_type", "status": "unmapped", "reason": "file-normalization: content-type parameter stripping"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_with_multiple_params", "status": "unmapped", "reason": "file-normalization: content-type parameter stripping"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_reject_file_type_document_in_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body file-type guard, a different mechanism than Rust's URL-fetch SSRF guard"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_accept_document_url_type_in_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body parsing"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_raise_on_invalid_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body parsing error path"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_ignore_document_form_field_injection", "status": "unmapped", "reason": "proxy-layer multipart form-field injection guard, a different mechanism than Rust's URL-fetch SSRF guard"}, + + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_extract_header_in_supported_params", "status": "unmapped", "reason": "Rust's fixed-list test checks the full list as one assertion, not this individual param"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_extract_footer_in_supported_params", "status": "unmapped", "reason": "Rust's fixed-list test checks the full list as one assertion, not this individual param"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_existing_params_still_present", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "supported_params_match_python_mistral_ocr_config", "justification": "both assert the full supported_ocr_params list matches the same fixed set of param names"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_header_passed_through", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "map_ocr_params_drops_unknown_params", "justification": "both assert extract_header survives map_ocr_params filtering unchanged"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_footer_passed_through", "status": "unmapped", "reason": "Rust's map_ocr_params test does not assert on extract_footer specifically"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_header_and_footer_together", "status": "unmapped", "reason": "combined extract_header+extract_footer passthrough is not asserted together in Rust"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_unknown_param_is_dropped", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "map_ocr_params_drops_unknown_params", "justification": "both assert an unrecognized param key is dropped while a known one is kept"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestNewSupportedParams::test_new_param_in_supported_list", "status": "unmapped", "reason": "OCR4-specific new params (table_format etc) are not individually verified against the Rust fixed-list test"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestNewParamsMapOcr::test_new_param_passed_through", "status": "unmapped", "reason": "OCR4-specific new params are not individually asserted in the Rust map_ocr_params test"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrRequest::test_param_included_in_request_body", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_builds_mistral_body", "justification": "both assert an optional param value ends up in the built request body alongside model/document"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrRequest::test_multiple_new_params_together", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_builds_mistral_body", "justification": "both assert multiple optional params (table_format/include_image_base64) land correctly in the same request body"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrResponseOcr4Fields::test_blocks_and_confidence_scores_preserved", "status": "unmapped", "reason": "OCR4 blocks/confidence_scores fields are not asserted by the Rust response test"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrResponseOcr4Fields::test_ocr4_fields_survive_model_dump", "status": "unmapped", "reason": "OCR4 tables/hyperlinks/header/footer fields are not asserted by the Rust response test"}, + + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_model_info_ocr4_price", "status": "unmapped", "reason": "cost-calc: pricing/model-info lookup is Python-only"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr4_cost_scales_with_pages", "status": "unmapped", "reason": "cost-calc: per-page pricing math is Python-only"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_pricing_entry", "status": "unmapped", "reason": "cost-calc: cost-map JSON entry validation is Python-only"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_model_info_price", "status": "unmapped", "reason": "cost-calc: pricing/model-info lookup is Python-only"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_cost_scales_with_pages", "status": "unmapped", "reason": "cost-calc: per-page pricing math is Python-only"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_ocr_and_annotation_pages_at_their_own_rates", "status": "unmapped", "reason": "cost-calc: mixed-rate billing math is Python-only"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_annotation_only_response", "status": "unmapped", "reason": "cost-calc: annotation-only billing math is Python-only"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_annotation_pages_when_pages_processed_missing", "status": "unmapped", "reason": "cost-calc: fallback billing math is Python-only"}, + {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_azure_doc_ai_annotation_pages_fall_back_to_ocr_rate", "status": "unmapped", "reason": "cost-calc: fallback billing math is Python-only"}, + + {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_rust_ocr_serves_default_format", "status": "unmapped", "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_rust_ocr_skipped_for_native_format", "status": "unmapped", "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_native_format_rejected_for_provider_without_support_as_bad_request", "status": "unmapped", "reason": "provider-support validation for req_format happens in Python"}, + {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_unknown_format_rejected_for_provider_without_support_as_bad_request", "status": "unmapped", "reason": "req_format validation error path is Python-only"}, + + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestHandlerDiscovery::test_handler_discovered_for_ocr", "status": "unmapped", "reason": "guardrail-translation handler discovery is a Python proxy-layer concern"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestHandlerDiscovery::test_handler_discovered_for_aocr", "status": "unmapped", "reason": "guardrail-translation handler discovery is a Python proxy-layer concern"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_document_url", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_image_url", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_no_document", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_invalid_document", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_input_blocking_guardrail", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_single_page", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_multiple_pages", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_empty_pages", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_page_with_empty_markdown", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_preserves_page_metadata", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_output_blocking_guardrail", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, + {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestPIIMaskingScenario::test_pii_masking_in_ocr_pages", "status": "unmapped", "reason": "PII redaction in the translation handler has no Rust equivalent"}, + + {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_read_req_format_from_header", "status": "unmapped", "reason": "proxy-layer header parsing has no Rust equivalent"}, + {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_prefer_body_req_format_over_header", "status": "unmapped", "reason": "proxy-layer body-vs-header precedence has no Rust equivalent"}, + {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_omit_req_format_when_header_absent", "status": "unmapped", "reason": "proxy-layer parsing has no Rust equivalent"}, + {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_reject_unknown_req_format", "status": "unmapped", "reason": "proxy-layer validation has no Rust equivalent"}, + {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_return_native_payload_with_litellm_response_headers", "status": "unmapped", "reason": "proxy-layer response construction has no Rust equivalent"}, + {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_return_normalized_response_when_no_native_payload", "status": "unmapped", "reason": "proxy-layer response construction has no Rust equivalent"} + ], + "rust_only_tests": [ + {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "rust_custom_logger_reads_failure_payload_for_non_ocr_call_type", "reason": "exercises the non-OCR (acompletion) call-type branch of the logger; the OCR branch is covered separately by rust_custom_logger_reads_success_payload_for_ocr"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "no_callback_fast_path_dispatches_nothing", "reason": "Rust-only fast-path optimization test for when zero callbacks are registered; Python has no equivalent no-op dispatch path"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "with_standard_logging_payload_keeps_top_level_fields_in_sync", "reason": "Rust-internal builder-method invariant, Python has no equivalent internal builder"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "blocks_private_and_metadata_ips", "reason": "SSRF IP-blocking helper has no Python unit test; Python relies on the proxy-layer JSON/form guards instead"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "convert_document_url_rejects_loopback_fetch", "reason": "URL-fetch SSRF protection is Rust-gateway-only"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "convert_document_url_leaves_data_uri_untouched", "reason": "URL-fetch SSRF protection is Rust-gateway-only"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_passes_short_strings_through", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_caps_long_payloads", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_does_not_split_multibyte_chars", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_dispatch_supports_migrated_providers", "reason": "Rust-internal provider-config dispatch table has no equivalent Python unit test"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "string_headers_accepts_string_values", "reason": "Rust-gateway header-coercion helper has no Python counterpart"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "auth_header_detection_is_case_insensitive", "reason": "Rust-gateway header-detection helper has no Python counterpart"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_runs_pre_during_and_success_hooks", "reason": "full gateway-level guardrail-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_runs_failure_hook_on_provider_error", "reason": "full gateway-level failure-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_pre_call_block_skips_provider_socket", "reason": "full gateway-level pre-call-block-plus-socket-skip test with no Python equivalent at this integration scope"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_does_not_duplicate_authorization_header_when_header_is_supplied", "reason": "outgoing HTTP header dedup at the Rust gateway has no Python counterpart"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "document_intelligence_poll_uses_resolved_subscription_key", "reason": "full Azure DI poll-loop integration test with no Python equivalent at this scope"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "string_headers_rejects_non_string_values", "reason": "Rust-gateway header-coercion error path has no Python counterpart"}, + {"rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "azure_ai_reuses_mistral_body_transform", "reason": "Rust-internal delegation-to-Mistral-transform implementation detail, no Python test asserts this delegation"}, + {"rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_request_uses_base64_source_for_data_uri", "reason": "no Python test asserts on the base64Source request body shape"}, + {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_mistral_url_uses_project_location_and_model", "reason": "vertex OCR support has no Python unit test coverage yet"}, + {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_mistral_reuses_mistral_body_transform", "reason": "vertex OCR support has no Python unit test coverage yet"}, + {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_deepseek_request_uses_ocr_endpoint_shape", "reason": "vertex OCR support has no Python unit test coverage yet"}, + {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_deepseek_response_wraps_markdown_content", "reason": "vertex OCR support has no Python unit test coverage yet"}, + {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_rejects_non_object_document", "reason": "non-object document rejection has no dedicated Python unit test"}, + {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_response_normalizes_mistral_json", "reason": "Python's response tests target OCR4-specific fields only, none asserts the same base normalization this Rust test checks"}, + {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "complete_url_defaults_and_dedupes_v1", "reason": "URL-building/defaulting for Mistral has no Python unit test"}, + {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "resolve_api_key_prefers_param_then_env", "reason": "API key resolution precedence at the Rust provider-config layer has no Python unit test"}, + {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "resolve_api_key_errors_when_absent", "reason": "API key resolution error path at the Rust provider-config layer has no Python unit test"}, + {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "rust_custom_logger_reads_success_payload_for_ocr", "reason": "Rust-internal custom-logger dispatch for OCR payloads has no Python unit test at this layer"} + ] +} diff --git a/tests/rust-python-harness/strategies/unit_tests/mapping_validator.py b/tests/rust-python-harness/strategies/unit_tests/mapping_validator.py new file mode 100644 index 00000000000..b8355e04266 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests/mapping_validator.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from ...shared.parity.ledger import TestLedger, load_ledger +from .python_runner import enumerate_python_tests +from .rust_runner import enumerate_rust_tests + +REPO_ROOT = Path(__file__).resolve().parents[4] +LEDGER_ROOT = Path(__file__).parent / "ledgers" + + +def ledger_path_for(sdk_function: str) -> Path: + return LEDGER_ROOT / sdk_function / f"{sdk_function}_test_ledger.json" + + +@dataclass(frozen=True, slots=True) +class AuditReport: + missing_python_tests: tuple[str, ...] + stale_python_tests: tuple[str, ...] + missing_rust_tests: tuple[str, ...] + stale_rust_tests: tuple[str, ...] + + @property + def is_clean(self) -> bool: + return not ( + self.missing_python_tests + or self.stale_python_tests + or self.missing_rust_tests + or self.stale_rust_tests + ) + + +def _ledger_python_tests_by_file(ledger: TestLedger) -> dict[str, set[str]]: + grouping: dict[str, set[str]] = {path: set() for path in ledger.python_scope} + for entry in ledger.entries: + grouping.setdefault(entry.python_file, set()).add(entry.python_test) + return grouping + + +def _ledger_rust_tests_by_file(ledger: TestLedger) -> dict[str, set[str]]: + grouping: dict[str, set[str]] = {path: set() for path in ledger.rust_scope} + for entry in ledger.entries: + if entry.status == "mapped": + grouping.setdefault(entry.rust_file, set()).add(entry.rust_test) + for rust_only in ledger.rust_only_tests: + grouping.setdefault(rust_only.rust_file, set()).add(rust_only.rust_test) + return grouping + + +def audit_ledger(ledger: TestLedger, repo_root: Path = REPO_ROOT) -> AuditReport: + missing_python: list[str] = [] + stale_python: list[str] = [] + for python_file, ledger_tests in _ledger_python_tests_by_file(ledger).items(): + actual_tests = enumerate_python_tests(repo_root, python_file) + for missing in sorted(ledger_tests - actual_tests): + missing_python.append(f"{python_file}:{missing}") + for stale in sorted(actual_tests - ledger_tests): + stale_python.append(f"{python_file}:{stale}") + + missing_rust: list[str] = [] + stale_rust: list[str] = [] + for rust_file, ledger_tests in _ledger_rust_tests_by_file(ledger).items(): + actual_tests = enumerate_rust_tests(repo_root, rust_file) + for missing in sorted(ledger_tests - actual_tests): + missing_rust.append(f"{rust_file}:{missing}") + for stale in sorted(actual_tests - ledger_tests): + stale_rust.append(f"{rust_file}:{stale}") + + return AuditReport( + missing_python_tests=tuple(missing_python), + stale_python_tests=tuple(stale_python), + missing_rust_tests=tuple(missing_rust), + stale_rust_tests=tuple(stale_rust), + ) + + +@dataclass(frozen=True, slots=True) +class FunctionReport: + sdk_function: str + ledger: TestLedger | None + audit: AuditReport | None + + @property + def has_ledger(self) -> bool: + return self.ledger is not None + + @property + def is_clean(self) -> bool: + return self.audit is None or self.audit.is_clean + + +def build_function_report(sdk_function: str, repo_root: Path = REPO_ROOT) -> FunctionReport: + path = ledger_path_for(sdk_function) + if not path.exists(): + return FunctionReport(sdk_function=sdk_function, ledger=None, audit=None) + ledger = load_ledger(path) + return FunctionReport( + sdk_function=sdk_function, ledger=ledger, audit=audit_ledger(ledger, repo_root) + ) diff --git a/tests/rust-python-harness/strategies/unit_tests/python_runner.py b/tests/rust-python-harness/strategies/unit_tests/python_runner.py new file mode 100644 index 00000000000..a7528d27756 --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests/python_runner.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +import ast +from pathlib import Path + + +def enumerate_python_tests(repo_root: Path, relative_path: str) -> frozenset[str]: + source = (repo_root / relative_path).read_text(encoding="utf-8") + tree = ast.parse(source, filename=relative_path) + + module_level: list[str] = [] + for node in ast.iter_child_nodes(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name.startswith("test_"): + module_level.append(node.name) + elif isinstance(node, ast.ClassDef): + for child in ast.iter_child_nodes(node): + if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) and child.name.startswith( + "test_" + ): + module_level.append(f"{node.name}::{child.name}") + + return frozenset(module_level) diff --git a/tests/rust-python-harness/strategies/unit_tests/rust_runner.py b/tests/rust-python-harness/strategies/unit_tests/rust_runner.py new file mode 100644 index 00000000000..6b855adbc4b --- /dev/null +++ b/tests/rust-python-harness/strategies/unit_tests/rust_runner.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +import re +from pathlib import Path + +_RUST_TEST_PATTERN = re.compile( + r"#\[(?:test|tokio::test)\][^\n]*\n(?:[^\n]*\n)*?\s*(?:async\s+)?fn\s+(\w+)\s*\(" +) + + +def enumerate_rust_tests(repo_root: Path, relative_path: str) -> frozenset[str]: + source = (repo_root / relative_path).read_text(encoding="utf-8") + return frozenset(match.group(1) for match in _RUST_TEST_PATTERN.finditer(source)) diff --git a/tests/test_rust_python_harness.py b/tests/test_rust_python_harness.py index 90dc38663dd..6a8fa8d35cc 100644 --- a/tests/test_rust_python_harness.py +++ b/tests/test_rust_python_harness.py @@ -8,14 +8,24 @@ import pytest catalog = importlib.import_module("tests.rust-python-harness.catalog") cli = importlib.import_module("tests.rust-python-harness.cli") +ledger_module = importlib.import_module("tests.rust-python-harness.shared.parity.ledger") +mapping_validator = importlib.import_module( + "tests.rust-python-harness.strategies.unit_tests.mapping_validator" +) models = importlib.import_module("tests.rust-python-harness.models") runner = importlib.import_module("tests.rust-python-harness.runner") ui = importlib.import_module("tests.rust-python-harness.ui") load_catalog = catalog.load_catalog +load_ledger = ledger_module.load_ledger +ledger_path_for = mapping_validator.ledger_path_for +REPO_ROOT = mapping_validator.REPO_ROOT +audit_ledger = mapping_validator.audit_ledger +build_function_report = mapping_validator.build_function_report _pick_values = cli._pick_values _coverage_pytest_args = cli._coverage_pytest_args _select = cli._select +_validate_ledger = cli._validate_ledger CaseResult = models.CaseResult Coverage = models.Coverage HarnessCase = models.HarnessCase @@ -234,3 +244,48 @@ def test_should_report_confidence_for_each_sdk_section() -> None: assert scores["responses"].level.value == "MEDIUM" assert scores["count_tokens"].percentage == 0 assert scores["count_tokens"].level.value == "LOW" + + + +def test_should_report_no_ledger_for_a_function_without_one() -> None: + report = build_function_report("messages", repo_root=REPO_ROOT) + + assert report.has_ledger is False + assert report.is_clean is True + + +def test_should_report_ocr_ledger_stats_and_a_clean_audit() -> None: + ledger = load_ledger(ledger_path_for("ocr")) + + report = build_function_report("ocr", repo_root=REPO_ROOT) + + assert report.has_ledger is True + assert report.ledger.mapped_count == ledger.mapped_count + assert report.ledger.total_count == ledger.total_count + assert report.is_clean is True + + +def test_should_scope_validate_ledger_to_the_requested_function( + capsys: pytest.CaptureFixture[str], +) -> None: + exit_code = _validate_ledger({"messages"}) + + captured = capsys.readouterr() + assert exit_code == 0 + assert "messages" in captured.out + assert "no ledger yet" in captured.out + assert "ocr" not in captured.out + + +def test_should_have_every_python_and_rust_ocr_test_accounted_for_in_the_ledger() -> None: + ledger = load_ledger(ledger_path_for("ocr")) + + report = audit_ledger(ledger, repo_root=REPO_ROOT) + + assert report.is_clean, ( + "\nOCR test-parity ledger is out of sync with the live test files.\n" + f"Ledger references a Python test that no longer exists: {list(report.missing_python_tests)}\n" + f"Python test exists but is not tracked in the ledger: {list(report.stale_python_tests)}\n" + f"Ledger references a Rust test that no longer exists: {list(report.missing_rust_tests)}\n" + f"Rust test exists but is not tracked in the ledger: {list(report.stale_rust_tests)}\n" + ) From 64e45a069db6794577a94104a20d78e6b5803ffb Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 2 Sep 2026 19:36:09 -0700 Subject: [PATCH 19/23] feat(complexity_router): opt-in modality override of a kept session-affinity pin (#39454) A session pinned to a text-only model failed every image turn with a provider 400, because the modality gate exempts a kept pin by cause. Add modality_pin_override so that exemption is conditional: the image turn is re-placed on a capable model for that request only, reported as cause modality_pin_override, and the stored pin is left untouched so the next text turn replays it. The pin write on the replay path already happens upstream of the gate and stores the session's own model, so pin survival is structural rather than bookkeeping. The new cause joins the non-pinnable set. Default off at every layer. --- .../public_endpoints/autorouter_presets.json | 4 ++ .../complexity_router/README.md | 12 +++- .../complexity_router/complexity_router.py | 23 +++++-- .../complexity_router/config.py | 15 +++- litellm/types/utils.py | 4 ++ .../router_strategy/test_complexity_router.py | 68 ++++++++++++++++++- .../add_model/ComplexityRouterConfig.test.tsx | 38 +++++++++++ .../add_model/ComplexityRouterConfig.tsx | 1 + .../add_model/ModalityRoutingControls.tsx | 50 +++++++++----- .../add_model/add_auto_router_tab.test.tsx | 38 +++++++++++ .../add_model/add_auto_router_tab.tsx | 1 + .../build_complexity_router_config.test.ts | 9 +++ .../build_complexity_router_config.ts | 4 ++ ...d_updated_complexity_router_config.test.ts | 26 +++++++ .../edit_auto_router_modal.test.ts | 22 ++++++ .../edit_auto_router_modal.test.tsx | 44 ++++++++++++ .../edit_auto_router_modal.tsx | 5 ++ .../RoutingDecisionCard.test.tsx | 6 ++ .../LogDetailsDrawer/RoutingDecisionCard.tsx | 1 + .../src/lib/autorouter_presets.test.ts | 19 ++++++ .../src/lib/autorouter_presets.ts | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 ++- 22 files changed, 374 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/public_endpoints/autorouter_presets.json b/litellm/proxy/public_endpoints/autorouter_presets.json index 1e0d40567d7..c2b13b81542 100644 --- a/litellm/proxy/public_endpoints/autorouter_presets.json +++ b/litellm/proxy/public_endpoints/autorouter_presets.json @@ -17,6 +17,7 @@ "classification_mode": "every_request", "session_affinity": false, "modality_routing": false, + "modality_pin_override": false, "deployment_affinity": true } }, @@ -35,6 +36,7 @@ "classification_mode": "every_request", "session_affinity": false, "modality_routing": false, + "modality_pin_override": false, "deployment_affinity": true } }, @@ -63,6 +65,7 @@ "classification_mode": "every_request", "session_affinity": false, "modality_routing": false, + "modality_pin_override": false, "deployment_affinity": true } }, @@ -84,6 +87,7 @@ "classification_mode": "every_request", "session_affinity": false, "modality_routing": false, + "modality_pin_override": false, "deployment_affinity": true } } diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index ad8b67d5e8f..ee51add1ca1 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -187,6 +187,9 @@ model_list: # Replace a routed model that cannot take image input (default: false) modality_routing: true + + # Let that replacement also override a kept session pin, for image turns only (default: false) + modality_pin_override: true ``` ## Usage @@ -227,9 +230,16 @@ vision model sits below the decided tier gets the 400 and an actionable message A same-tier re-pick keeps the decision's cause and adds `modality:image` to `signals`; a tier change or default takeover records `cause: modality_escalation` with the displaced placement (`modality_escalated_from:` or `modality_displaced_default_model`). Escalations are never -pinned by session affinity, and a KEPT session pin bypasses the gate entirely: a session pinned +pinned by session affinity, and by default a KEPT session pin bypasses the gate: a session pinned to a text-only model keeps it even when an image arrives. +Add `modality_pin_override: true` to lift that last exemption. The image turn is then re-placed +the same way every other decision is, and records `cause: modality_pin_override` whether or not +the tier moved, since the model left the pin either way. The pin itself is untouched: the session +affinity write happens upstream of the gate and stores the session's own model, so the next text +turn replays the original pin and the override is never pinned in its place. It does nothing +unless `modality_routing` is also on. + ### Heuristic-first chaining `classifier_type: heuristic_first` runs the local scorer on every request and only calls the LLM diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 430efe339a2..a00ae6bee80 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -751,7 +751,8 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo A modality escalation is transient the same way: it describes what this one call carries (an image), not what the session's traffic looks like, and pinning it would hold every following - text turn on the vision-capable model the image forced. + text turn on the vision-capable model the image forced. A modality pin override is the same + fact on a session that already holds a pin, so it must not overwrite the pin it displaced. """ return decision is None or ( decision.get("cause") @@ -760,6 +761,7 @@ def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bo "plan_mode", "housekeeping", "modality_escalation", + "modality_pin_override", ) and not decision.get("context_escalated") ) @@ -2393,8 +2395,11 @@ class ComplexityRouter(CustomLogger): """Replace a routed model that cannot accept this request's image input. The single modality owner, applied to the decided response at the hook's exits so every - routing path is covered uniformly. A KEPT session pin is exempt by design (its cause); - replacement picks and every other path are just responses. The re-placement walks + routing path is covered uniformly. A KEPT session pin is exempt by design (its cause) + unless modality_pin_override is set, in which case the image turn is re-placed and reported + as modality_pin_override while the stored pin, written upstream from the session's own + model, is left for the next text turn; replacement picks and every other path are just + responses. The re-placement walks UPWARD-ONLY from the decision's tier (so a plan-mode floor can never be undercut), picks through `_pick_model_for_tier` so routing plugins still apply, then falls to default_model (never on plugin routers, and never on a plan-floored decision, since @@ -2407,7 +2412,11 @@ class ComplexityRouter(CustomLogger): not self.config.modality_routing or not resolved_messages or response.model is None - or (decision is not None and decision.get("cause") == "session_affinity_pin") + or ( + decision is not None + and decision.get("cause") == "session_affinity_pin" + and not self.config.modality_pin_override + ) or not request_contains_image_content(resolved_messages) or self._model_accepts_image_input(response.model) ): @@ -2449,6 +2458,10 @@ class ComplexityRouter(CustomLogger): self._restamp_adaptive_choice(request_kwargs, response.model, new_model) same_tier: Final = capable is not None and decided == capable base_cause: Final = (decision.get("cause") if decision is not None else None) or "default_fallback" + # Reaching here on a kept pin means modality_pin_override is on, since the guard above + # returns otherwise. The model moved off the pin even on a same-tier repick, so reporting + # the pin's own cause would claim the session's model served a request it did not. + displaced_pin: Final = base_cause == "session_affinity_pin" displaced_default: Final = decided is None and response.model == self.config.default_model markers: Final = ( "modality:image", @@ -2458,7 +2471,7 @@ class ComplexityRouter(CustomLogger): old_signals: Final = tuple(decision.get("signals") or ()) if decision is not None else () new_decision: Final = self._build_routing_decision( routed_model=new_model, - cause=base_cause if same_tier else "modality_escalation", + cause="modality_pin_override" if displaced_pin else (base_cause if same_tier else "modality_escalation"), tier=new_tier, score=decision.get("score") if decision is not None else None, signals=(*old_signals, *markers), diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 0ae0db63fad..9f2054dda01 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -883,7 +883,20 @@ class ComplexityRouterConfig(BaseModel): "a routed model explicitly declared supports_vision false (deployment model_info " "or the model cost map; unmapped names stay routable) is replaced by the nearest " "HIGHER tier holding a capable model, then default_model, else a clear 400. A kept " - "session-affinity pin still wins even when an image arrives." + "session-affinity pin still wins even when an image arrives, unless " + "modality_pin_override is also enabled." + ), + ) + modality_pin_override: bool = Field( + default=False, + description=( + "Let modality_routing replace a kept session-affinity pin on the turns that carry an " + "image. Without this, a session pinned to a text-only model fails every image turn with " + "a provider 400, since the pin is exempt from the modality gate. When enabled, such a " + "turn routes to a capable model for that request only and the stored pin is left " + "untouched, so the next text turn replays the session's own model; the override is " + "reported as cause modality_pin_override and is never itself pinned. Inert unless " + "modality_routing is also enabled." ), ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ee6f09e05dc..569fce4f7b8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2879,6 +2879,10 @@ RoutingDecisionCause = Literal[ # routed model does not accept image input, so the nearest higher capable tier or # default_model served instead. The displaced placement rides in signals. "modality_escalation", + # modality_pin_override replaced a KEPT session-affinity pin for this request only: the turn + # carries an image the pinned model cannot accept. The stored pin is untouched, so the next + # text turn replays it. Distinct from "modality_escalation", which never displaces a pin. + "modality_pin_override", "session_affinity_pin", "session_affinity_escalation", # classification_mode 'user_turn': the request is an agent loop's continuation turn (no new diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index d7d02544efb..3ecccb673f3 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -10797,6 +10797,9 @@ class TestModalityRouting: ("custom_tiers_walk", "premium-model", "modality_escalation"), ("pin_kept_bypasses", "text-cheap", "session_affinity_pin"), ("pin_replacement_gated", "vision-big", "modality_escalation"), + ("pin_override_escalates", "vision-mid", "modality_pin_override"), + ("pin_override_same_tier", "vision-cheap", "modality_pin_override"), + ("pin_override_inert_without_modality_routing", "text-cheap", "session_affinity_pin"), ("adaptive_pick_rewritten", "vision-mid", "modality_escalation"), ], ) @@ -10847,7 +10850,7 @@ class TestModalityRouting: messages = [ {"role": "user", "content": [{"type": "text", "text": "quick lookup: what is this?"}, IMG_PART]} ] - elif path in ("pin_kept_bypasses", "pin_replacement_gated"): + elif path.startswith(("pin_kept", "pin_replacement", "pin_override")): cache = AsyncMock() cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) mock_router_instance.cache = cache @@ -10859,6 +10862,13 @@ class TestModalityRouting: messages = [ {"role": "user", "content": [{"type": "text", "text": "LITELLM ESCALATE describe this"}, IMG_PART]} ] + elif path == "pin_override_same_tier": + config["modality_pin_override"] = True + config["tiers"]["SIMPLE"] = ["text-cheap", "vision-cheap"] + vision["vision-cheap"] = True + elif path == "pin_override_inert_without_modality_routing": + config["modality_routing"] = False + config["modality_pin_override"] = path.startswith("pin_override") elif path == "adaptive_pick_rewritten": config["adaptive"] = True mock_router_instance.model_list = [] @@ -11025,4 +11035,60 @@ class TestModalityRouting: from litellm.router_strategy.complexity_router.complexity_router import _decision_is_pinnable assert _decision_is_pinnable({"cause": "modality_escalation"}) is False + assert _decision_is_pinnable({"cause": "modality_pin_override"}) is False assert _decision_is_pinnable({"cause": "heuristic_scorer"}) is True + + @pytest.mark.asyncio + async def test_pin_override_serves_the_image_turn_without_repinning(self, mock_router_instance): + """The override is for one request: the session keeps the model it was pinned to.""" + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) + mock_router_instance.cache = cache + router = self._router( + mock_router_instance, + { + "tiers": dict(self.BASE_TIERS), + "modality_routing": True, + "modality_pin_override": True, + "session_affinity": True, + }, + dict(self.BASE_VISION), + ) + request_kwargs = {"metadata": {"session_id": "s1"}} + + image_turn = await router.async_pre_routing_hook( + model="m", request_kwargs=request_kwargs, messages=self.IMAGE_MESSAGE + ) + assert image_turn.model == "vision-mid" + assert image_turn.routing_decision["cause"] == "modality_pin_override" + assert "modality_escalated_from:SIMPLE" in image_turn.routing_decision["signals"] + + assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"} + + text_turn = await router.async_pre_routing_hook( + model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=[{"role": "user", "content": "hi"}] + ) + assert text_turn.model == "text-cheap" + assert text_turn.routing_decision["cause"] == "session_affinity_pin" + + @pytest.mark.asyncio + async def test_pin_override_with_no_capable_model_rejects_and_keeps_the_pin(self, mock_router_instance): + """The clear 400 replaces the provider's, and a rejected turn must not cost the session its pin.""" + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) + mock_router_instance.cache = cache + router = self._router( + mock_router_instance, + { + "tiers": {"SIMPLE": "text-cheap", "COMPLEX": "text-big"}, + "modality_routing": True, + "modality_pin_override": True, + "session_affinity": True, + }, + {"text-cheap": False, "text-big": False}, + ) + with pytest.raises(litellm.BadRequestError, match="no model"): + await router.async_pre_routing_hook( + model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=self.IMAGE_MESSAGE + ) + assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"} diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 349d03356f8..9b647badf69 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -880,6 +880,44 @@ describe("ComplexityRouterConfig modality panel", () => { expect(screen.getByRole("switch", { name: "Route image requests to vision-capable models" })).toBeChecked(); }); + + // The backend ignores modality_pin_override unless modality_routing is on, so offering it while + // image routing is off would let an operator save a flag that does nothing. + it("disables the pin-override switch while image routing is off", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Modality Routing")); + + const override = screen.getByRole("switch", { name: "Override session pin for image requests" }); + expect(override).toHaveAttribute("aria-disabled", "true"); + fireEvent.click(override); + expect(onChange).not.toHaveBeenCalled(); + }); + + it("writes modality_pin_override through onChange once image routing is on", () => { + const onChange = vi.fn(); + const value = { ...defaultValue, modality_routing: true }; + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Modality Routing")); + + const override = screen.getByRole("switch", { name: "Override session pin for image requests" }); + expect(override).not.toBeChecked(); + fireEvent.click(override); + + expect(onChange).toHaveBeenCalledWith({ ...value, modality_pin_override: true }); + }); + + it("renders a stored modality_pin_override=true as on", () => { + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Modality Routing")); + + expect(screen.getByRole("switch", { name: "Override session pin for image requests" })).toBeChecked(); + }); }); describe("ComplexityRouterConfig affinity panel", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 73ab4f0abef..29dd3d9ea83 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -410,6 +410,7 @@ export interface ComplexityRouterConfigValue { classification_mode?: ClassificationMode; session_affinity?: boolean; modality_routing?: boolean; + modality_pin_override?: boolean; deployment_affinity?: boolean; /** Plan-mode floor as a tier ROW ID, unset meaning off. The wire carries the row's name. */ plan_mode_min_tier?: string; diff --git a/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx b/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx index dd697b35239..8d6c4ed5d5e 100644 --- a/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ModalityRoutingControls.tsx @@ -7,20 +7,36 @@ import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; export const ModalityRoutingControls: React.FC<{ value: ComplexityRouterConfigValue; onChange: (value: ComplexityRouterConfigValue) => void; -}> = ({ value, onChange }) => ( - <> -
- onChange({ ...value, modality_routing: modalityRouting })} - aria-label="Route image requests to vision-capable models" - /> - Route image requests to vision-capable models -
- - Replaces a routed model that cannot take image input with the nearest higher tier that can, then the default - model, instead of failing with a provider 400. Only models explicitly declared supports_vision false are replaced, - and a kept session pin still wins. - - -); +}> = ({ value, onChange }) => { + const modalityRouting = value.modality_routing ?? false; + return ( + <> +
+ onChange({ ...value, modality_routing: nextModalityRouting })} + aria-label="Route image requests to vision-capable models" + /> + Route image requests to vision-capable models +
+ + Replaces a routed model that cannot take image input with the nearest higher tier that can, then the default + model, instead of failing with a provider 400. Only models explicitly declared supports_vision false are + replaced, and a kept session pin still wins unless you turn on the override below. + +
+ onChange({ ...value, modality_pin_override: modalityPinOverride })} + disabled={!modalityRouting} + aria-label="Override session pin for image requests" + /> + Override session pin for image requests +
+ + Route an image turn to a capable model even when the session is pinned to one that cannot take images. The pin + is kept, so the next text turn goes back to it. Needs image routing turned on. + + + ); +}; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 7769042a0c7..27a169524b9 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -590,6 +590,44 @@ describe("AddAutoRouterTab", () => { }); }); + it("writes both modality flags as false into the create payload when the panel stays untouched", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + fireEvent.change(screen.getByPlaceholderText(/smart_router/i), { target: { value: "modality-router" } }); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ + modality_routing: false, + modality_pin_override: false, + }); + }); + + it("carries the pin override through to the create payload once image routing unlocks it", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + fireEvent.change(screen.getByPlaceholderText(/smart_router/i), { target: { value: "modality-router" } }); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Modality Routing")); + await user.click(await screen.findByRole("switch", { name: "Route image requests to vision-capable models" })); + await user.click(await screen.findByRole("switch", { name: "Override session pin for image requests" })); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ + modality_routing: true, + modality_pin_override: true, + }); + }); + // Custom is the escape hatch, not the headline choice, so it's listed after every bundled preset // rather than first. it("lists Custom Configuration after the bundled presets", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index ab5ea4cbfc9..fc3fdaf39cd 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -358,6 +358,7 @@ const AddAutoRouterTab: React.FC = ({ classifierFallback: complexityRouterConfig.classifier_fallback, sessionAffinity: complexityRouterConfig.session_affinity ?? DEFAULT_SESSION_AFFINITY, modalityRouting: complexityRouterConfig.modality_routing ?? false, + modalityPinOverride: complexityRouterConfig.modality_pin_override ?? false, deploymentAffinity: complexityRouterConfig.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, customTechnicalKeywords, keywordTierRules, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index c6370541a45..3ef55b9295d 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -56,6 +56,7 @@ describe("buildComplexityRouterConfig", () => { session_affinity: false, deployment_affinity: true, modality_routing: false, + modality_pin_override: false, escalation_keywords: ["LITELLM ESCALATE"], }; expect(config).toEqual(expected); @@ -270,6 +271,14 @@ describe("buildComplexityRouterConfig", () => { expect(buildComplexityRouterConfig({ ...baseParams, modalityRouting: false }).modality_routing).toBe(false); }); + it("writes modality_pin_override explicitly both ways, so the stored config never relies on the backend default", () => { + expect(buildComplexityRouterConfig({ ...baseParams, modalityPinOverride: true }).modality_pin_override).toBe(true); + expect(buildComplexityRouterConfig(baseParams).modality_pin_override).toBe(false); + expect(buildComplexityRouterConfig({ ...baseParams, modalityPinOverride: false }).modality_pin_override).toBe( + false, + ); + }); + it("writes session_affinity=true so turning the toggle on overrides the backend's off-by-default", () => { const config = buildComplexityRouterConfig({ ...baseParams, sessionAffinity: true }); expect(config.session_affinity).toBe(true); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 94c5badf6d1..91e4fdb23b5 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -111,6 +111,7 @@ export interface BuildComplexityRouterConfigParams { classificationMode: ClassificationMode | undefined; sessionAffinity: boolean; modalityRouting?: boolean; + modalityPinOverride?: boolean; deploymentAffinity: boolean; customTechnicalKeywords: string[]; keywordTierRules: KeywordTierRule[]; @@ -169,6 +170,7 @@ export interface ComplexityRouterConfigPayload { session_affinity: boolean; deployment_affinity: boolean; modality_routing: boolean; + modality_pin_override: boolean; custom_technical_keywords?: string[]; keyword_tier_rules?: { keywords: string[]; tier: KeywordTierRule["tier"] }[]; semantic_keyword_matching?: boolean; @@ -409,6 +411,7 @@ export const buildComplexityRouterConfig = ({ classificationMode, sessionAffinity, modalityRouting, + modalityPinOverride, deploymentAffinity, customTechnicalKeywords, keywordTierRules, @@ -472,6 +475,7 @@ export const buildComplexityRouterConfig = ({ session_affinity: sessionAffinity, deployment_affinity: deploymentAffinity, modality_routing: modalityRouting ?? false, + modality_pin_override: modalityPinOverride ?? false, ...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }), ...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }), escalation_keywords: cleanedEscalationKeywords, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index fae6744d3ac..01198bcc118 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -257,6 +257,30 @@ describe("buildUpdatedComplexityRouterConfig session affinity", () => { }); }); +describe("buildUpdatedComplexityRouterConfig modality pin override", () => { + it("writes modality_pin_override explicitly both ways", () => { + expect( + buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, modality_pin_override: true }).modality_pin_override, + ).toBe(true); + expect( + buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, modality_pin_override: false }).modality_pin_override, + ).toBe(false); + }); + + it("re-asserts the backend's off-by-default when the form value is absent, rather than dropping the key", () => { + const result = buildUpdatedComplexityRouterConfig({ ...STORED, modality_pin_override: true }, FORM_VALUE); + expect(result.modality_pin_override).toBe(false); + }); + + it("round-trips a stored modality_pin_override=true through hydrate then save", () => { + const stored = { ...STORED, modality_routing: true, modality_pin_override: true }; + const hydrated = hydrateComplexityRouterConfig(stored, undefined); + + expect(hydrated.modality_pin_override).toBe(true); + expect(buildUpdatedComplexityRouterConfig(stored, hydrated).modality_pin_override).toBe(true); + }); +}); + describe("buildUpdatedComplexityRouterConfig classification mode", () => { it("round-trips a stored user_turn through hydrate then save", () => { const stored = { ...STORED, classification_mode: "user_turn" }; @@ -505,6 +529,8 @@ describe("managed keys survive an untouched open-and-save", () => { classifier_fallback: "default_model", classification_mode: "user_turn", session_affinity: true, + modality_routing: true, + modality_pin_override: true, deployment_affinity: false, adaptive: true, adaptive_weights: { quality: 0.4, cost: 0.6 }, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts index 481ba2b6b00..651d5880db4 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts @@ -51,6 +51,7 @@ const expectedClassifiedTierConfig = { session_affinity: false, deployment_affinity: true, modality_routing: false, + modality_pin_override: false, adaptive: true, adaptive_weights: { quality: 0.4, cost: 0.6 }, adaptive_eligible: "classified_tier", @@ -74,6 +75,7 @@ const expectedAdaptiveDisabledConfig = { session_affinity: false, deployment_affinity: true, modality_routing: false, + modality_pin_override: false, }; describe("buildUpdatedComplexityRouterConfig", () => { @@ -109,6 +111,26 @@ describe("buildUpdatedComplexityRouterConfig", () => { expect(disabled.modality_routing).toBe(false); }); + it("hydrates a stored modality_pin_override into form state and defaults absent to off", () => { + expect( + hydrateComplexityRouterConfig({ ...storedConfig, modality_pin_override: true }, null).modality_pin_override, + ).toBe(true); + expect(hydrateComplexityRouterConfig(storedConfig, null).modality_pin_override).toBe(false); + }); + + it("round-trips modality_pin_override explicitly in both directions", () => { + const enabled = buildUpdatedComplexityRouterConfig(storedConfig, { + ...classifiedTierValue, + modality_pin_override: true, + }); + expect(enabled.modality_pin_override).toBe(true); + const disabled = buildUpdatedComplexityRouterConfig( + { ...storedConfig, modality_pin_override: true }, + { ...classifiedTierValue, modality_pin_override: false }, + ); + expect(disabled.modality_pin_override).toBe(false); + }); + it("includes return_raw_model_name only when enabled", () => { const updatedConfig = buildUpdatedComplexityRouterConfig(storedConfig, { ...classifiedTierValue, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index 96e8549eac8..6067e72e547 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -549,6 +549,50 @@ describe("EditAutoRouterModal deployment affinity", () => { await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); expect(savedConfig().deployment_affinity).toBe(false); }); + + // modality_pin_override is a managed key, so the modal rewrites it from form state on save. A + // hydration gap would silently turn a stored override off on the next untouched save. + it("shows a stored modality_pin_override=true as on and preserves it through an untouched save", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, modality_routing: true, modality_pin_override: true }); + + await user.click(await screen.findByText("Advanced: Modality Routing")); + expect(await screen.findByRole("switch", { name: "Override session pin for image requests" })).toBeChecked(); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().modality_pin_override).toBe(true); + }); + + it("persists turning the modality pin override on", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, modality_routing: true }); + + await user.click(await screen.findByText("Advanced: Modality Routing")); + await user.click(await screen.findByRole("switch", { name: "Override session pin for image requests" })); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().modality_pin_override).toBe(true); + }); + + it("writes modality_pin_override=false for a stored config that never carried the key", async () => { + const user = userEvent.setup(); + renderWithStoredConfig(STORED_CONFIG); + + await user.click(await screen.findByText("Advanced: Modality Routing")); + expect(await screen.findByRole("switch", { name: "Override session pin for image requests" })).toHaveAttribute( + "aria-disabled", + "true", + ); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().modality_pin_override).toBe(false); + }); }); describe("EditAutoRouterModal custom classifier prompt and fallback", () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 7aa5f4beb1f..442e1ec6039 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -104,6 +104,7 @@ export interface StoredComplexityRouterConfig { reasoning_override_min_score?: unknown; session_affinity?: unknown; modality_routing?: unknown; + modality_pin_override?: unknown; deployment_affinity?: unknown; adaptive?: boolean; adaptive_weights?: AdaptiveRouterWeights; @@ -181,6 +182,8 @@ export const hydrateComplexityRouterConfig = ( session_affinity: typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY, modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false, + modality_pin_override: + typeof parsedConfig.modality_pin_override === "boolean" ? parsedConfig.modality_pin_override : false, deployment_affinity: typeof parsedConfig.deployment_affinity === "boolean" ? parsedConfig.deployment_affinity @@ -221,6 +224,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "classification_mode", "session_affinity", "modality_routing", + "modality_pin_override", "deployment_affinity", "adaptive", "adaptive_weights", @@ -318,6 +322,7 @@ export const buildUpdatedComplexityRouterConfig = ( classifierFallback: value.classifier_fallback, sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY, modalityRouting: value.modality_routing ?? false, + modalityPinOverride: value.modality_pin_override ?? false, deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, customTechnicalKeywords: customTechnicalKeywords ?? [], keywordTierRules: keywordMatching?.keywordTierRules ?? [], diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx index bf6a20a4af8..fd1777f802c 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.test.tsx @@ -186,6 +186,12 @@ describe("RoutingDecisionCard", () => { expect(screen.queryByText("housekeeping")).not.toBeInTheDocument(); }); + it("labels a modality pin override instead of showing the raw cause token", () => { + render(); + expect(screen.getByText("Overrode session pin for image input")).toBeInTheDocument(); + expect(screen.queryByText("modality_pin_override")).not.toBeInTheDocument(); + }); + it("labels a modality escalation instead of showing the raw cause token", () => { render(); expect(screen.getByText("Escalated for image input")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx index ffef01010b7..cf2c71e64c6 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/RoutingDecisionCard.tsx @@ -93,6 +93,7 @@ const CONSTANT_CAUSE_LABELS: Record = { session_affinity_escalation: "Escalated from session pin", user_turn_continuation: "Continuation turn, classifier skipped", modality_escalation: "Escalated for image input", + modality_pin_override: "Overrode session pin for image input", quality_tier: "Quality tier mapping", bandit: "Adaptive bandit", default_fallback: "Default model, no route matched", diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index e3132067f4c..44ac6973a1c 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -148,6 +148,25 @@ describe("autorouter_presets", () => { expect(withoutFlag.complexityRouterConfig.modality_routing).toBe(false); }); + it("carries a preset's modality_pin_override into the prefilled form state", () => { + const preset = getPresetByKey("anthropic_family")!; + const withFlag = { ...preset.complexity_router_config, modality_routing: true, modality_pin_override: true }; + const prefill = buildPresetPrefill(withFlag, groupsOnly(getRequiredModelsInPreset(preset))); + expect(prefill.complexityRouterConfig.modality_pin_override).toBe(true); + const withoutFlag = buildPresetPrefill( + preset.complexity_router_config, + groupsOnly(getRequiredModelsInPreset(preset)), + ); + expect(withoutFlag.complexityRouterConfig.modality_pin_override).toBe(false); + }); + + it("ships every bundled preset with both modality flags written out, since the payload type requires them", () => { + for (const preset of getAllPresets()) { + expect(preset.complexity_router_config.modality_routing, preset.key).toBe(false); + expect(preset.complexity_router_config.modality_pin_override, preset.key).toBe(false); + } + }); + it("prefills the anthropic preset's effort through to tier_model_params", () => { const preset = getPresetByKey("anthropic_family")!; const prefill = buildPresetPrefill(preset.complexity_router_config, groupsOnly(getRequiredModelsInPreset(preset))); diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index 30f8ed99e74..ff482bd23b8 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -286,6 +286,7 @@ export const buildPresetPrefill = ( session_affinity: config.session_affinity ?? DEFAULT_SESSION_AFFINITY, deployment_affinity: config.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, modality_routing: config.modality_routing ?? false, + modality_pin_override: config.modality_pin_override ?? false, adaptive: config.adaptive, adaptive_weights: config.adaptive_weights, tier_distance_penalty: config.tier_distance_penalty, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 491c4fb6a44..b0f8e618645 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34740,9 +34740,15 @@ export interface components { * @default 0.5 */ match_threshold: number; + /** + * Modality Pin Override + * @description Let modality_routing replace a kept session-affinity pin on the turns that carry an image. Without this, a session pinned to a text-only model fails every image turn with a provider 400, since the pin is exempt from the modality gate. When enabled, such a turn routes to a capable model for that request only and the stored pin is left untouched, so the next text turn replays the session's own model; the override is reported as cause modality_pin_override and is never itself pinned. Inert unless modality_routing is also enabled. + * @default false + */ + modality_pin_override: boolean; /** * Modality Routing - * @description Route image-bearing requests only to models that can accept image input. The classifier reads text alone, so an image request whose text classifies cheap otherwise lands on a text-only model and fails with a provider 400. When enabled, a routed model explicitly declared supports_vision false (deployment model_info or the model cost map; unmapped names stay routable) is replaced by the nearest HIGHER tier holding a capable model, then default_model, else a clear 400. A kept session-affinity pin still wins even when an image arrives. + * @description Route image-bearing requests only to models that can accept image input. The classifier reads text alone, so an image request whose text classifies cheap otherwise lands on a text-only model and fails with a provider 400. When enabled, a routed model explicitly declared supports_vision false (deployment model_info or the model cost map; unmapped names stay routable) is replaced by the nearest HIGHER tier holding a capable model, then default_model, else a clear 400. A kept session-affinity pin still wins even when an image arrives, unless modality_pin_override is also enabled. * @default false */ modality_routing: boolean; @@ -35945,7 +35951,7 @@ export interface components { * Cause * @enum {string} */ - cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; + cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "modality_pin_override" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; /** Classifier Cost */ classifier_cost?: number; /** Classifier Model */ From 291e84e565ca8a2f91981f32885970e7585d5ae6 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 2 Sep 2026 19:46:09 -0700 Subject: [PATCH 20/23] feat(datadog_llm_obs): cost tag dimensions, router decision fields, reasoning token metric, redaction gating (#39402) * feat(datadog_llm_obs): cost tag dimensions, router decision fields, reasoning token metric, redaction gating * test(datadog_llm_obs): satisfy test quality gate * fix: forward integer parent_id as its string form * fix(datadog): sanitize redacted message roles * fix(datadog): keep the A2A agent role on redacted spans * fix(datadog): merge current staging budget * style(datadog): format redaction tests * fix(datadog): handle malformed redacted roles * test(datadog): put the test quality suppression on the reported line Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 6 +- .../integrations/datadog/datadog_llm_obs.py | 217 ++++++++-- litellm/types/integrations/datadog_llm_obs.py | 1 + .../datadog/test_datadog_llm_obs.py | 390 +++++++++++++++++- type-discipline-budget.json | 4 +- 5 files changed, 577 insertions(+), 41 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 3f96531cf6f..2967fc2a505 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15290 + "limit": 15288 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,7 +105,7 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38332 + "limit": 38324 }, "reportUnknownParameterType": { "limit": 19625 @@ -123,7 +123,7 @@ "limit": 4 }, "reportUnnecessaryIsInstance": { - "limit": 823 + "limit": 819 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 5e116b7301a..ec86c0ae1d9 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -19,11 +19,13 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.constants import REDACTED_BY_LITELLM from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.integrations.datadog.datadog_handler import ( get_datadog_base_url_from_env, get_datadog_service, get_datadog_tags, + normalize_datadog_tag_value, ) from litellm.integrations.datadog.datadog_mock_client import ( create_mock_datadog_client, @@ -34,6 +36,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, handle_any_messages_to_chat_completion_str_messages_conversion, ) +from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import ( @@ -43,6 +46,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens from litellm.types.integrations.datadog_llm_obs import * from litellm.types.utils import ( + PROMPT_QUOTING_ROUTING_DECISION_FIELDS, CallTypes, StandardLoggingGuardrailInformation, StandardLoggingPayload, @@ -52,6 +56,120 @@ from litellm.types.utils import ( _EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({}) _EMPTY_MESSAGE: Final[Message] = {"role": "", "content": ""} _MAX_PARSED_TOOL_ARGUMENT_CHARS: Final = 256 * 1024 +_SAFE_REDACTED_MESSAGE_ROLES: Final = frozenset( + {"agent", "assistant", "developer", "function", "model", "system", "tool", "user"} +) + +_PROMPT_CARRYING_METADATA_FIELDS: Final = frozenset( + { + "routing_decision", + "requester_metadata", + "prompt_management_metadata", + "mcp_tool_call_metadata", + "vector_store_request_metadata", + } +) + +_ROUTER_SPAN_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + { + "tier": "router_tier", + "cause": "router_cause", + "score": "router_score", + "escalated": "router_escalated", + "signals": "router_signals", + "routed_model": "routed_model", + } +) +_ROUTER_DIMENSIONS: Final[tuple[str, ...]] = ("router_tier", "router_cause", "router_escalated", "routed_model") +_COST_DIMENSIONS: Final[tuple[str, ...]] = ("team", "user", "key_alias", "model_group", *_ROUTER_DIMENSIONS) + + +def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, Any]: + metadata: Final = standard_logging_payload.get("metadata") + return metadata or _EMPTY_MAPPING + + +def _router_span_fields( + standard_logging_payload: StandardLoggingPayload, redact_prompt_text: bool +) -> Mapping[str, object]: + """Flatten the auto-router decision, omitting prompt-quoting fields when redaction is enabled.""" + routing_decision: Final = _mapping_field(_metadata_of(standard_logging_payload), "routing_decision") + if not routing_decision: + return _EMPTY_MAPPING + escalated: Final = bool(routing_decision.get("escalated") or routing_decision.get("context_escalated")) + return MappingProxyType( + { + _ROUTER_SPAN_FIELDS[record_field]: value + for record_field, value in (*routing_decision.items(), ("escalated", escalated)) + if record_field in _ROUTER_SPAN_FIELDS + and value is not None + and not (redact_prompt_text and record_field in PROMPT_QUOTING_ROUTING_DECISION_FIELDS) + } + ) + + +def _metadata_without_prompt_carriers(standard_logging_metadata: Mapping[str, Any]) -> Mapping[str, Any]: + """The metadata minus the records that quote prompts, tool arguments, tool results, or retrieved text.""" + return MappingProxyType( + { + field: value + for field, value in standard_logging_metadata.items() + if field not in _PROMPT_CARRYING_METADATA_FIELDS + } + ) + + +def _redact_messages(messages: Sequence[Message]) -> tuple[Message, ...]: + """Each message's shape with its content replaced and tool payloads dropped; no message is invented.""" + return tuple( + { + "role": role if isinstance(role, str) and role in _SAFE_REDACTED_MESSAGE_ROLES else "", + "content": REDACTED_BY_LITELLM, + } + for message in messages + for role in (message.get("role", ""),) + ) + + +def _cost_dimension_tags( + standard_logging_payload: StandardLoggingPayload, router_fields: Mapping[str, object] +) -> tuple[str, ...]: + """The dimensions LLM Obs breaks token and cost metrics down by, as span tags.""" + metadata: Final = _metadata_of(standard_logging_payload) + dimensions: Final = ( + ("user", metadata.get("user_api_key_user_id")), + ("key_alias", metadata.get("user_api_key_alias")), + ("model_group", standard_logging_payload.get("model_group")), + *((dimension, router_fields.get(dimension)) for dimension in _ROUTER_DIMENSIONS), + ) + return tuple( + f"{key}:{normalized}" + for key, value in dimensions + if value is not None and (normalized := normalize_datadog_tag_value(value)) != "" + ) + + +def _declared_cost_tags(span_tags: Sequence[str]) -> tuple[str, ...]: + """Declare only cost dimensions carrying a value on this span.""" + present: Final = frozenset(key for tag in span_tags if (key := tag.partition(":")[0]) and tag.partition(":")[2]) + return tuple(dimension for dimension in _COST_DIMENSIONS if dimension in present) + + +def _reasoning_output_tokens(usage_object: Mapping[str, Any] | None) -> float: + """The provider's reasoning-token count, from either the chat or the responses spelling.""" + if usage_object is None: + return 0.0 + return next( + ( + float(reasoning_tokens) + for details_field in ("completion_tokens_details", "output_tokens_details") + if isinstance( + reasoning_tokens := _mapping_field(usage_object, details_field).get("reasoning_tokens"), (int, float) + ) + and not isinstance(reasoning_tokens, bool) + ), + 0.0, + ) def _mapping_field(source: Mapping[str, Any], key: str) -> Mapping[str, Any]: @@ -316,12 +434,12 @@ class DataDogLLMObsLogger(CustomBatchLogger): dict_datadog_llm_obs_params: dict = {} if litellm.datadog_llm_observability_params is not None: if isinstance(litellm.datadog_llm_observability_params, DatadogLLMObsInitParams): - dict_datadog_llm_obs_params = litellm.datadog_llm_observability_params.model_dump() + dict_datadog_llm_obs_params = litellm.datadog_llm_observability_params.model_dump(exclude_unset=True) elif isinstance(litellm.datadog_llm_observability_params, dict): # only allow params that are of DatadogLLMObsInitParams dict_datadog_llm_obs_params = DatadogLLMObsInitParams( **litellm.datadog_llm_observability_params - ).model_dump() + ).model_dump(exclude_unset=True) return dict_datadog_llm_obs_params async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -410,25 +528,40 @@ class DataDogLLMObsLogger(CustomBatchLogger): if standard_logging_payload is None: raise Exception("DataDogLLMObs: standard_logging_object is not set") - metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {}) + raw_metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {}) + metadata: Final = raw_metadata if isinstance(raw_metadata, dict) else {} + redact_payload: Final = self._payload_logging_is_off(kwargs) - input_meta: Final = InputMeta(messages=_to_dd_messages(standard_logging_payload["messages"])) + input_messages: Final = _to_dd_messages(standard_logging_payload.get("messages")) + output_messages: Final = self._get_response_messages( + standard_logging_payload=standard_logging_payload, + call_type=standard_logging_payload.get("call_type"), + ) + input_meta: Final = InputMeta(messages=_redact_messages(input_messages) if redact_payload else input_messages) output_meta: Final = OutputMeta( - messages=self._get_response_messages( - standard_logging_payload=standard_logging_payload, - call_type=standard_logging_payload.get("call_type"), - ) + messages=_redact_messages(output_messages) if redact_payload else output_messages ) error_info: Final = self._assemble_error_info(standard_logging_payload) - metadata_parent_id: str | None = None - if isinstance(metadata, dict): - metadata_parent_id = metadata.get("parent_id") + raw_parent_id: Final = metadata.get("parent_id") + metadata_parent_id: Final[str | None] = str(raw_parent_id) if raw_parent_id else None - tool_definitions: Final = _to_dd_tool_definitions(standard_logging_payload.get("model_parameters")) + tool_definitions: Final = ( + () if redact_payload else _to_dd_tool_definitions(standard_logging_payload.get("model_parameters")) + ) span_kind: Final = self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id) - payload_metadata: Final = self._get_dd_llm_obs_payload_metadata(standard_logging_payload) + router_fields: Final = _router_span_fields(standard_logging_payload, redact_prompt_text=redact_payload) + span_tags: Final = [ + *get_datadog_tags(standard_logging_object=standard_logging_payload), + *_cost_dimension_tags(standard_logging_payload, router_fields), + ] + payload_metadata: Final = self._get_dd_llm_obs_payload_metadata( + standard_logging_payload, + router_fields=router_fields, + cost_tags=_declared_cost_tags(span_tags), + redact_prompt_text=redact_payload, + ) meta: Final[Meta] = { "kind": span_kind, @@ -451,7 +584,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): duration=int((end_time - start_time).total_seconds() * 1e9), metrics=metrics, status="error" if error_info else "ok", - tags=get_datadog_tags(standard_logging_object=standard_logging_payload), + tags=span_tags, ) apm_trace_id: Final = self._get_apm_trace_id() @@ -497,6 +630,13 @@ class DataDogLLMObsLogger(CustomBatchLogger): ) return error_info + def _payload_logging_is_off(self, kwargs: Mapping[str, Any]) -> bool: + return ( + bool(self.turn_off_message_logging) + or self.message_logging is not True + or should_redact_message_logging(dict(kwargs)) + ) + def _assemble_metrics(self, standard_logging_payload: StandardLoggingPayload) -> LLMMetrics: """ Build the span metrics, including the prompt-cache counts LLM Obs charts cache savings from. @@ -513,10 +653,11 @@ class DataDogLLMObsLogger(CustomBatchLogger): total_cost: Final = float(standard_logging_payload.get("response_cost", 0)) time_to_first_token: Final = self._get_time_to_first_token_seconds(standard_logging_payload) - raw_usage: Final = (standard_logging_payload.get("metadata") or {}).get("usage_object") + raw_usage: Final = _metadata_of(standard_logging_payload).get("usage_object") usage_object: Final = raw_usage if isinstance(raw_usage, dict) else None cache_read: Final = float(extract_cache_read_tokens(usage_object)) cache_write: Final = float(extract_cache_creation_tokens(usage_object)) + reasoning_output_tokens: Final = _reasoning_output_tokens(usage_object) metrics: Final[LLMMetrics] = { "input_tokens": prompt_tokens, @@ -533,6 +674,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): if cache_read or cache_write else {} ), + **({"reasoning_output_tokens": reasoning_output_tokens} if reasoning_output_tokens else {}), } return metrics @@ -707,11 +849,21 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Default fallback for unknown or passthrough operations return "llm" - def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]: + def _get_dd_llm_obs_payload_metadata( + self, + standard_logging_payload: StandardLoggingPayload, + router_fields: Mapping[str, object] | None = None, + cost_tags: Sequence[str] = (), + redact_prompt_text: bool = False, + ) -> dict[str, object]: """ Fields to track in DD LLM Observability metadata from litellm standard logging payload """ - _metadata: Final[dict[str, object]] = { + raw_metadata: Final = _metadata_of(standard_logging_payload) + standard_logging_metadata: Final = ( + _metadata_without_prompt_carriers(raw_metadata) if redact_prompt_text else raw_metadata + ) + return { "model_name": standard_logging_payload.get("model", "unknown"), "model_provider": standard_logging_payload.get("custom_llm_provider", "unknown"), "id": standard_logging_payload.get("id", "unknown"), @@ -719,26 +871,21 @@ class DataDogLLMObsLogger(CustomBatchLogger): "cache_hit": standard_logging_payload.get("cache_hit", "unknown"), "cache_key": standard_logging_payload.get("cache_key", "unknown"), "saved_cache_cost": standard_logging_payload.get("saved_cache_cost", 0), - "guardrail_information": standard_logging_payload.get("guardrail_information", None), + "guardrail_information": ( + None if redact_prompt_text else standard_logging_payload.get("guardrail_information", None) + ), "is_streamed_request": self._get_stream_value_from_payload(standard_logging_payload), + "latency_metrics": dict(self._get_latency_metrics(standard_logging_payload)), + "spend_metrics": dict(self._get_spend_metrics(standard_logging_payload)), + **standard_logging_metadata, + **(router_fields or _EMPTY_MAPPING), + **( + {"_dd": {**_mapping_field(standard_logging_metadata, "_dd"), "cost_tags": list(cost_tags)}} + if cost_tags + else _EMPTY_MAPPING + ), } - ######################################################### - # Add latency metrics to metadata - ######################################################### - latency_metrics: Final = self._get_latency_metrics(standard_logging_payload) - _metadata.update({"latency_metrics": dict(latency_metrics)}) - - ######################################################### - # Add spend metrics to metadata - ######################################################### - spend_metrics: Final = self._get_spend_metrics(standard_logging_payload) - _metadata.update({"spend_metrics": dict(spend_metrics)}) - - _standard_logging_metadata: Final[dict] = dict(standard_logging_payload.get("metadata", {})) or {} - _metadata.update(_standard_logging_metadata) - return _metadata - def _get_latency_metrics(self, standard_logging_payload: StandardLoggingPayload) -> DDLLMObsLatencyMetrics: """ Get the latency metrics from the standard logging payload @@ -808,7 +955,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): spend_metrics["response_cost"] = standard_logging_payload.get("response_cost", 0.0) # Get budget information from metadata - metadata: Final = standard_logging_payload.get("metadata", {}) + metadata: Final = _metadata_of(standard_logging_payload) # API key max budget user_api_key_max_budget: Final = metadata.get("user_api_key_max_budget") diff --git a/litellm/types/integrations/datadog_llm_obs.py b/litellm/types/integrations/datadog_llm_obs.py index bae876dfdd9..17cf5831c96 100644 --- a/litellm/types/integrations/datadog_llm_obs.py +++ b/litellm/types/integrations/datadog_llm_obs.py @@ -86,6 +86,7 @@ class LLMMetrics(TypedDict, total=False): cache_read_input_tokens: ReadOnly[float] cache_write_input_tokens: ReadOnly[float] non_cached_input_tokens: ReadOnly[float] + reasoning_output_tokens: ReadOnly[float] class LLMObsPayload(TypedDict, total=False): diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py index 2d0605e3b7f..34c62864c4e 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py @@ -17,6 +17,7 @@ from unittest.mock import patch import pytest +import litellm from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -55,15 +56,22 @@ def build_payload( response_message: dict[str, Any] | None = None, usage_object: dict[str, Any] | None = None, model_parameters: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, + model_group: str | None = None, prompt_tokens: int = 4447, ) -> dict[str, Any]: + standard_logging_metadata: dict[str, Any] = { + **(metadata or {}), + **({"usage_object": usage_object} if usage_object is not None else {}), + } return { "standard_logging_object": { "call_type": "acompletion", "messages": [{"role": "user", "content": "hi"}] if messages is NOT_GIVEN else messages, "response": {"choices": [{"message": response_message or {"role": "assistant", "content": "hello"}}]}, "model_parameters": model_parameters or {}, - "metadata": {"usage_object": usage_object} if usage_object is not None else {}, + "metadata": standard_logging_metadata, + "model_group": model_group, "prompt_tokens": prompt_tokens, "completion_tokens": 507, "total_tokens": prompt_tokens + 507, @@ -244,6 +252,43 @@ def test_no_cache_keys_when_the_provider_reports_no_caching(logger: DataDogLLMOb assert "non_cached_input_tokens" not in payload["metrics"] +def test_reasoning_tokens_are_reported_as_span_metrics(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, usage_object={"completion_tokens_details": {"reasoning_tokens": 128}}) + + assert payload["metrics"]["reasoning_output_tokens"] == 128.0 + + +def test_responses_reasoning_tokens_are_reported_as_span_metrics(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, usage_object={"output_tokens_details": {"reasoning_tokens": 64}}) + + assert payload["metrics"]["reasoning_output_tokens"] == 64.0 + + +def test_zero_reasoning_tokens_are_not_reported(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, usage_object={"completion_tokens_details": {"reasoning_tokens": 0}}) + + assert "reasoning_output_tokens" not in payload["metrics"] + + +def test_reasoning_tokens_come_from_the_spelling_that_reports_them(logger: DataDogLLMObsLogger) -> None: + """A chat-details mapping without the count must not shadow the responses spelling that has it.""" + payload = build( + logger, + usage_object={ + "completion_tokens_details": {"accepted_prediction_tokens": 5}, + "output_tokens_details": {"reasoning_tokens": 64}, + }, + ) + + assert payload["metrics"]["reasoning_output_tokens"] == 64.0 + + +def test_boolean_reasoning_tokens_are_not_a_count(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, usage_object={"completion_tokens_details": {"reasoning_tokens": True}}) + + assert "reasoning_output_tokens" not in payload["metrics"] + + def test_tool_definitions_are_sent_on_meta(logger: DataDogLLMObsLogger) -> None: payload = build(logger, model_parameters={"tools": [TOOL_DEFINITION]}) @@ -256,6 +301,340 @@ def test_tool_definitions_are_sent_on_meta(logger: DataDogLLMObsLogger) -> None: ] +def test_cost_tags_include_present_categories_and_dimensions(logger: DataDogLLMObsLogger) -> None: + payload = build( + logger, + metadata={ + "user_api_key_user_id": "User 42", + "user_api_key_alias": "Primary Key", + "team_alias": "Platform", + "routing_decision": { + "tier": "premium", + "cause": "high_complexity", + "score": 0.91, + "escalated": True, + "signals": ["long prompt"], + "routed_model": "openai/gpt-5", + }, + }, + model_group="premium-models", + ) + + assert payload["tags"][-8:] == [ + "team:platform", + "user:user_42", + "key_alias:primary_key", + "model_group:premium-models", + "router_tier:premium", + "router_cause:high_complexity", + "router_escalated:true", + "routed_model:openai/gpt-5", + ] + assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == [ + "team", + "user", + "key_alias", + "model_group", + "router_tier", + "router_cause", + "router_escalated", + "routed_model", + ] + + +def test_missing_cost_tag_values_are_not_declared(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, metadata={"team_alias": "Platform"}) + + assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == ["team"] + assert not any(tag.startswith(("user:", "key_alias:", "model_group:")) for tag in payload["tags"]) + + +def test_values_that_normalize_to_empty_are_not_tagged_or_declared(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, metadata={"user_api_key_user_id": "___", "user_api_key_alias": "!!!"}, model_group="tier-1") + + assert not any(tag in ("user:", "key_alias:") for tag in payload["tags"]) + assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == ["model_group"] + + +def test_a_valueless_tag_from_the_shared_builder_is_not_declared(logger: DataDogLLMObsLogger) -> None: + """The team tag comes from the shared builder, which emits it bare when the alias normalizes away.""" + payload = build(logger, metadata={"team_alias": "!!!"}, model_group="tier-1") + + assert "team:" in payload["tags"] + assert payload["meta"]["metadata"]["_dd"]["cost_tags"] == ["model_group"] + + +def test_router_fields_are_flattened(logger: DataDogLLMObsLogger) -> None: + payload = build( + logger, + metadata={ + "routing_decision": { + "tier": "premium", + "cause": "high_complexity", + "score": 0.91, + "escalated": True, + "signals": ["secret prompt text"], + "routed_model": "openai/gpt-5", + } + }, + model_group="premium-models", + ) + + assert payload["meta"]["metadata"]["router_tier"] == "premium" + assert payload["meta"]["metadata"]["router_cause"] == "high_complexity" + assert payload["meta"]["metadata"]["router_score"] == 0.91 + assert payload["meta"]["metadata"]["router_escalated"] is True + assert payload["meta"]["metadata"]["router_signals"] == ["secret prompt text"] + assert payload["meta"]["metadata"]["routed_model"] == "openai/gpt-5" + + +def test_a_context_escalated_route_reports_as_escalated(logger: DataDogLLMObsLogger) -> None: + """The router records a size-driven escalation under its own key, and it is still an escalation.""" + payload = build(logger, metadata={"routing_decision": {"tier": "premium", "context_escalated": True}}) + + assert payload["meta"]["metadata"]["router_escalated"] is True + assert "router_escalated:true" in payload["tags"] + + +def test_a_routed_request_that_did_not_escalate_reports_false(logger: DataDogLLMObsLogger) -> None: + """Without this the escalation dimension is absent on ordinary traffic, so nothing can group by it.""" + payload = build(logger, metadata={"routing_decision": {"tier": "simple", "cause": "heuristic_scorer"}}) + + assert payload["meta"]["metadata"]["router_escalated"] is False + assert "router_escalated:false" in payload["tags"] + assert "router_escalated" in payload["meta"]["metadata"]["_dd"]["cost_tags"] + + +def test_a_request_that_never_reached_a_router_has_no_router_fields(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, model_group="premium-models") + + assert "router_escalated" not in payload["meta"]["metadata"] + assert not any(tag.startswith("router_") for tag in payload["tags"]) + + +def test_redacted_payload_keeps_metrics_and_removes_sensitive_fields(logger: DataDogLLMObsLogger) -> None: + payload = build_payload( + messages=[{"role": "user", "content": "secret prompt"}], + response_message={"role": "assistant", "content": "secret response"}, + usage_object={"prompt_tokens_details": {"cached_tokens": 128}}, + metadata={"routing_decision": {"tier": "premium", "signals": ["secret prompt text"]}}, + model_parameters={"tools": [TOOL_DEFINITION]}, + ) + with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True): + with patch("asyncio.create_task"): + redacted_logger = DataDogLLMObsLogger(turn_off_message_logging=True) + redacted_payload = redacted_logger.redact_standard_logging_payload_from_model_call_details(payload) + result = json.loads( + safe_dumps( + redacted_logger.create_llm_obs_payload( + redacted_payload, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2) + ) + ) + ) + + assert result["meta"]["input"]["messages"][0]["content"] == "redacted-by-litellm" + assert result["meta"]["output"]["messages"][0]["content"] == "redacted-by-litellm" + assert result["meta"]["metadata"]["router_tier"] == "premium" + assert "router_signals" not in result["meta"]["metadata"] + assert "routing_decision" not in result["meta"]["metadata"] + assert "tool_definitions" not in result["meta"] + assert result["metrics"]["cache_read_input_tokens"] == 128.0 + assert result["metrics"]["total_cost"] == 0.02 + + +def test_redaction_drops_the_routing_record_carried_in_metadata(logger: DataDogLLMObsLogger) -> None: + """The whole routing record rides along in metadata, so dropping the flat copy alone leaks the prompt.""" + with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True): + with patch("asyncio.create_task"): + redacted_logger = DataDogLLMObsLogger(turn_off_message_logging=True) + result = json.loads( + safe_dumps( + redacted_logger.create_llm_obs_payload( + build_payload( + metadata={ + "routing_decision": { + "tier": "premium", + "cause": "keyword_rule", + "signals": ["secret prompt text"], + "matched_keyword": "secret keyword", + "escalation_keyword": "secret escalation", + } + } + ), + datetime(2026, 9, 1, 12, 0, 0), + datetime(2026, 9, 1, 12, 0, 2), + ) + ) + ) + + assert "routing_decision" not in result["meta"]["metadata"] + assert result["meta"]["metadata"]["router_tier"] == "premium" + assert result["meta"]["metadata"]["router_cause"] == "keyword_rule" + assert "secret" not in safe_dumps(result["meta"]["metadata"]) + + +def test_a_failure_span_redacts_its_messages(logger: DataDogLLMObsLogger) -> None: + """The redaction hook only runs on success, so the failure span has to redact for itself.""" + failed = build_payload(messages=[{"role": "user", "content": "secret prompt"}]) + failed["standard_logging_object"]["status"] = "failure" + failed["standard_logging_object"]["response"] = None + failed["standard_logging_object"]["error_information"] = {"error_message": "boom", "error_class": "BadRequestError"} + with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True): + with patch("asyncio.create_task"): + redacted_logger = DataDogLLMObsLogger(turn_off_message_logging=True) + result = json.loads( + safe_dumps( + redacted_logger.create_llm_obs_payload( + failed, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2) + ) + ) + ) + + assert result["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}] + assert result["meta"]["output"]["messages"] == [] + assert result["status"] == "error" + + +def test_excluding_messages_from_the_logging_payload_still_ships_the_span(logger: DataDogLLMObsLogger) -> None: + """`standard_logging_payload_excluded_fields` deletes the key, and a span with no prompt is still a span.""" + payload = build_payload() + del payload["standard_logging_object"]["messages"] + + span = json.loads( + safe_dumps( + logger.create_llm_obs_payload(payload, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2)) + ) + ) + + assert span["meta"]["input"]["messages"] == [] + assert span["metrics"]["total_cost"] == 0.02 + + +def test_an_explicit_redaction_setting_survives_the_global_params(logger: DataDogLLMObsLogger) -> None: + """Global params carry defaults for keys the operator never set, and those must not win.""" + with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True): + with patch("asyncio.create_task"): + with patch.object( # test-quality-ok: the ctor reads this module global with no injection seam + litellm, "datadog_llm_observability_params", {} + ): + configured_logger = DataDogLLMObsLogger( + turn_off_message_logging=True + ) # test-quality-ok: verifies ctor setting + + assert configured_logger.turn_off_message_logging is True + + +def _redacting_logger( + **kwargs: Any, +) -> DataDogLLMObsLogger: # test-quality-ok: shared test factory accepts init variants + with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True): + with patch("asyncio.create_task"): + return DataDogLLMObsLogger(**kwargs) + + +def _span_json(logger_under_test: DataDogLLMObsLogger, payload: dict[str, Any]) -> dict[str, Any]: + span = logger_under_test.create_llm_obs_payload( + payload, datetime(2026, 9, 1, 12, 0, 0), datetime(2026, 9, 1, 12, 0, 2) + ) + return json.loads(safe_dumps(span)) + + +def test_redaction_keeps_the_conversation_shape_without_its_content() -> None: + """Roles and message count survive so the trace stays legible; contents and tool payloads do not.""" + result = _span_json( + _redacting_logger(turn_off_message_logging=True), + build_payload( + messages=[ + {"role": "user", "content": "secret prompt"}, + {"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]}, + ], + response_message={"role": "assistant", "content": "secret response"}, + ), + ) + + assert result["meta"]["input"]["messages"] == [ + {"role": "user", "content": "redacted-by-litellm"}, + {"role": "assistant", "content": "redacted-by-litellm"}, + ] + assert result["meta"]["output"]["messages"] == [{"role": "assistant", "content": "redacted-by-litellm"}] + + +def test_redaction_drops_unrecognized_and_malformed_message_roles() -> None: + """Caller-controlled role values must not bypass redaction or crash span creation.""" + result = _span_json( + _redacting_logger(turn_off_message_logging=True), + build_payload( + messages=[ + {"role": "SECRET-39402", "content": "hello"}, + {"role": ["SECRET-39402"], "content": "hello"}, + {"role": {"secret": "SECRET-39402"}, "content": "hello"}, + {"role": "agent", "content": "hello"}, + ] + ), + ) + + assert result["meta"]["input"]["messages"] == [ + {"role": "", "content": "redacted-by-litellm"}, + {"role": "", "content": "redacted-by-litellm"}, + {"role": "", "content": "redacted-by-litellm"}, + {"role": "agent", "content": "redacted-by-litellm"}, + ] + assert "SECRET-39402" not in safe_dumps(result) + + +def test_the_deprecated_message_logging_flag_engages_the_same_redaction() -> None: + """The platform redacts for `message_logging is not True`, so this callback's own gate must agree.""" + result = _span_json( + _redacting_logger(message_logging=False), + build_payload( + messages=[{"role": "user", "content": "secret prompt"}], + model_parameters={"tools": [TOOL_DEFINITION]}, + metadata={"routing_decision": {"tier": "premium", "signals": ["secret prompt text"]}}, + ), + ) + + assert result["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}] + assert "tool_definitions" not in result["meta"] + assert "routing_decision" not in result["meta"]["metadata"] + + +def test_a_truthy_redaction_setting_redacts_like_the_shared_hook() -> None: + """The shared hook redacts on truthiness, so a config-provided string must not half-redact the span.""" + result = _span_json( + _redacting_logger(turn_off_message_logging="yes"), + build_payload(messages=[{"role": "user", "content": "secret prompt"}]), + ) + + assert result["meta"]["input"]["messages"] == [{"role": "user", "content": "redacted-by-litellm"}] + + +def test_redaction_drops_every_prompt_carrying_metadata_record(logger: DataDogLLMObsLogger) -> None: + """Tool arguments, retrieved text, and the guardrail's copy of the request ride in metadata records too.""" + sensitive_metadata: dict[str, Any] = { + "requester_metadata": {"note": "secret prompt text"}, + "prompt_management_metadata": {"prompt_id": "p1", "prompt_variables": {"topic": "secret"}}, + "mcp_tool_call_metadata": {"name": "search", "arguments": {"query": "secret"}}, + "vector_store_request_metadata": [{"query": "secret"}], + } + + def sensitive_payload() -> dict[str, Any]: + payload = build_payload(metadata=sensitive_metadata) + payload["standard_logging_object"]["guardrail_information"] = [ + {"guardrail_name": "g", "guardrail_request": {"messages": [{"content": "secret prompt"}]}} + ] + return payload + + redacted = _span_json(_redacting_logger(turn_off_message_logging=True), sensitive_payload()) + unredacted = _span_json(logger, sensitive_payload()) + + assert "secret" not in safe_dumps(redacted["meta"]["metadata"]) + for record in sensitive_metadata: + assert record not in redacted["meta"]["metadata"] + assert record in unredacted["meta"]["metadata"] + assert redacted["meta"]["metadata"]["guardrail_information"] is None + assert unredacted["meta"]["metadata"]["guardrail_information"] is not None + + def test_tool_definitions_accept_the_bare_anthropic_shape(logger: DataDogLLMObsLogger) -> None: """The Anthropic surface declares tools unwrapped, with input_schema instead of parameters.""" payload = build( @@ -272,6 +651,15 @@ def test_meta_omits_tool_definitions_when_no_tools_were_offered(logger: DataDogL assert "tool_definitions" not in build(logger)["meta"] +def test_a_ddtrace_integer_parent_id_is_forwarded_as_its_string(logger: DataDogLLMObsLogger) -> None: + """ddtrace hands span ids as ints; dropping them detaches the span from its APM trace.""" + kwargs = build_payload() + kwargs["litellm_params"]["metadata"]["parent_id"] = 8675309 + start = datetime(2026, 9, 1, 12, 0, 0) + span = json.loads(safe_dumps(logger.create_llm_obs_payload(kwargs, start, start + timedelta(seconds=2)))) + assert span["parent_id"] == "8675309" + + def test_unparseable_tool_arguments_are_preserved_rather_than_dropped(logger: DataDogLLMObsLogger) -> None: """A truncated argument string is still the only record of what the model tried to call.""" payload = build( diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f3c4c7760c6..5c25b8722ef 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22334 + "limit": 22330 }, "LIT002": { "limit": 26763 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16480 + "limit": 16478 }, "LIT011": { "limit": 5520 From bcd3e2d94d3faf06cdb70c4648d872e8ddd86f5b Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:50:08 -0700 Subject: [PATCH 21/23] feat(rust-python-harness): wire existing e2e SDK tests into the matrix (#39463) Adds chat_completions and transcription as SDK function columns, backed by the existing rust_bridge test files. Adds a fourth strategy folder, existing_e2e_test_sdk, that points at already-existing live-API SDK tests (tests/ocr_tests/ as a whole folder, plus chat completion and Whisper transcription tests) instead of writing new parity tests. Extends selector_matches_node with trailing-slash folder selectors so a whole test folder can back one matrix cell. --- tests/rust-python-harness/README.md | 9 ++++++--- tests/rust-python-harness/cli.py | 9 ++++----- .../e2e_fuzz_tests/strategy.json | 4 +++- .../existing_e2e_test_sdk/README.md | 3 +++ .../existing_e2e_test_sdk/strategy.json | 14 ++++++++++++++ tests/rust-python-harness/models.py | 2 +- tests/rust-python-harness/runner.py | 2 ++ tests/rust-python-harness/ui.py | 2 +- .../unit_tests_rust/strategy.json | 4 +++- .../validate_sub_methods/strategy.json | 4 +++- tests/test_rust_python_harness.py | 16 +++++++++++++--- 11 files changed, 53 insertions(+), 16 deletions(-) create mode 100644 tests/rust-python-harness/existing_e2e_test_sdk/README.md create mode 100644 tests/rust-python-harness/existing_e2e_test_sdk/strategy.json diff --git a/tests/rust-python-harness/README.md b/tests/rust-python-harness/README.md index e94ac87c3b3..a34e1a1ab73 100644 --- a/tests/rust-python-harness/README.md +++ b/tests/rust-python-harness/README.md @@ -8,14 +8,17 @@ The matrix always has these SDK columns: - `messages / amessages` - `responses / aresponses` - `count_tokens` +- `chat_completions / acompletion` +- `transcription / atranscription` -The harness has three deliberately broad test-strategy folders: +The harness has four deliberately broad test-strategy folders: | Strategy | Folder | | --- | --- | | Public SDK parity over generated and recorded inputs | [`e2e_fuzz_tests/`](e2e_fuzz_tests/) | | Focused tests of Rust-owned behavior | [`unit_tests_rust/`](unit_tests_rust/) | | Isolated transform and Python-to-Rust helper coverage | [`validate_sub_methods/`](validate_sub_methods/) | +| Already-existing live-API SDK tests | [`existing_e2e_test_sdk/`](existing_e2e_test_sdk/) | ## Run it @@ -112,7 +115,7 @@ The initial end-to-end entries deliberately show `◐`: the repository has Rust ## Attach parity tests -Each of the three folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list: +Each of the four folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list: ```json { @@ -123,7 +126,7 @@ Each of the three folders contains a concise `README.md` and a `strategy.json`. } ``` -Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice. +Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family; a selector ending in `/` aggregates every test in that folder, recursively. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice. Use these coverage values: diff --git a/tests/rust-python-harness/cli.py b/tests/rust-python-harness/cli.py index f9e97d7ad43..c996b68846c 100644 --- a/tests/rust-python-harness/cli.py +++ b/tests/rust-python-harness/cli.py @@ -6,14 +6,13 @@ from collections.abc import Sequence from pathlib import Path from .catalog import load_catalog -from .models import HarnessCase, Strategy +from .models import SDK_FUNCTIONS, HarnessCase, Strategy from .runner import run_pytest from .ui import make_dashboard from .strategies.unit_tests.mapping_validator import FunctionReport, build_function_report REPO_ROOT = Path(__file__).resolve().parents[2] COVERAGE_ROOT = REPO_ROOT / "target" / "rust-python-harness" -SDK_FUNCTION_CHOICES = ("ocr", "messages", "responses", "count_tokens") def _parser() -> argparse.ArgumentParser: @@ -42,7 +41,7 @@ def _parser() -> argparse.ArgumentParser: action="append", default=[], dest="sdk_functions", - choices=SDK_FUNCTION_CHOICES, + choices=SDK_FUNCTIONS, help="run only this SDK function", ) parser.add_argument( @@ -110,7 +109,7 @@ def _interactive_filters(strategies: Sequence[Strategy]) -> tuple[set[str], set[ ) sdk_functions = _pick_values( "SDK functions", - [(name, name) for name in SDK_FUNCTION_CHOICES], + [(name, name) for name in SDK_FUNCTIONS], ) return strategy_ids, sdk_functions @@ -166,7 +165,7 @@ def _print_function_report(report: FunctionReport) -> None: def _validate_ledger(sdk_functions: set[str]) -> int: - functions = sdk_functions or set(SDK_FUNCTION_CHOICES) + functions = sdk_functions or set(SDK_FUNCTIONS) reports = tuple(build_function_report(function) for function in sorted(functions)) for report in reports: _print_function_report(report) diff --git a/tests/rust-python-harness/e2e_fuzz_tests/strategy.json b/tests/rust-python-harness/e2e_fuzz_tests/strategy.json index abeea01d9b5..d838486772d 100644 --- a/tests/rust-python-harness/e2e_fuzz_tests/strategy.json +++ b/tests/rust-python-harness/e2e_fuzz_tests/strategy.json @@ -7,6 +7,8 @@ "ocr": {"coverage": "partial", "selectors": ["tests/test_litellm/ocr/test_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."}, "messages": {"coverage": "partial", "selectors": ["tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."}, "responses": {"coverage": "partial", "selectors": ["tests/test_litellm/responses/test_rust_bridge_websocket.py"], "note": "Covers the websocket bridge; full responses parity is still being added."}, - "count_tokens": {"coverage": "planned", "selectors": [], "note": "No Rust count_tokens parity test is present yet."} + "count_tokens": {"coverage": "planned", "selectors": [], "note": "No Rust count_tokens parity test is present yet."}, + "chat_completions": {"coverage": "partial", "selectors": ["tests/test_litellm/rust_bridge/test_chat_completions.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."}, + "transcription": {"coverage": "partial", "selectors": ["tests/test_litellm/test_audio_transcription_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."} } } diff --git a/tests/rust-python-harness/existing_e2e_test_sdk/README.md b/tests/rust-python-harness/existing_e2e_test_sdk/README.md new file mode 100644 index 00000000000..fb84f170703 --- /dev/null +++ b/tests/rust-python-harness/existing_e2e_test_sdk/README.md @@ -0,0 +1,3 @@ +# Existing e2e SDK tests + +Wires already-existing live-API SDK tests into the matrix instead of writing new parity tests. Selectors point at real test files and folders, such as `tests/ocr_tests/`, rather than individual node IDs, so future tests added to those folders are picked up automatically. diff --git a/tests/rust-python-harness/existing_e2e_test_sdk/strategy.json b/tests/rust-python-harness/existing_e2e_test_sdk/strategy.json new file mode 100644 index 00000000000..eefceea1a75 --- /dev/null +++ b/tests/rust-python-harness/existing_e2e_test_sdk/strategy.json @@ -0,0 +1,14 @@ +{ + "order": 40, + "id": "existing_e2e_test_sdk", + "label": "Existing e2e SDK tests", + "description": "Wire already-existing live-API SDK tests into the matrix instead of writing new parity tests.", + "functions": { + "ocr": {"coverage": "partial", "selectors": ["tests/ocr_tests/"], "note": "Existing live OCR provider tests; not yet a frozen Rust/Python oracle comparison."}, + "messages": {"coverage": "planned", "selectors": []}, + "responses": {"coverage": "planned", "selectors": []}, + "count_tokens": {"coverage": "planned", "selectors": []}, + "chat_completions": {"coverage": "partial", "selectors": ["tests/llm_translation/test_anthropic_completion.py", "tests/llm_translation/test_bedrock_completion.py"], "note": "Existing live chat completion tests for providers with confirmed Rust bridge regressions."}, + "transcription": {"coverage": "partial", "selectors": ["tests/audio_tests/test_whisper.py"], "note": "Existing live Whisper transcription test."} + } +} diff --git a/tests/rust-python-harness/models.py b/tests/rust-python-harness/models.py index 21097e0f7d0..a02684dc63d 100644 --- a/tests/rust-python-harness/models.py +++ b/tests/rust-python-harness/models.py @@ -33,7 +33,7 @@ class ConfidenceLevel(str, Enum): LOW = "LOW" -SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens") +SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens", "chat_completions", "transcription") @dataclass(frozen=True) diff --git a/tests/rust-python-harness/runner.py b/tests/rust-python-harness/runner.py index 82393ef234e..286f2c4116d 100644 --- a/tests/rust-python-harness/runner.py +++ b/tests/rust-python-harness/runner.py @@ -15,6 +15,8 @@ UpdateCallback = Callable[[HarnessRun], None] def selector_matches_node(selector: str, nodeid: str) -> bool: normalized_selector = selector.replace("\\", "/") normalized_nodeid = nodeid.replace("\\", "/") + if normalized_selector.endswith("/"): + return normalized_nodeid.startswith(normalized_selector) if "::" in normalized_selector: return normalized_nodeid == normalized_selector or normalized_nodeid.startswith( f"{normalized_selector}[" diff --git a/tests/rust-python-harness/ui.py b/tests/rust-python-harness/ui.py index 57fedd17fa6..bf143a4e100 100644 --- a/tests/rust-python-harness/ui.py +++ b/tests/rust-python-harness/ui.py @@ -119,7 +119,7 @@ class RichDashboard(AbstractContextManager["RichDashboard"]): table = Table(box=box.ROUNDED, expand=True, title="Strategy × SDK function") table.add_column("Strategy", ratio=3) - for label in ("ocr/aocr", "messages", "responses", "count_tokens"): + for label in SDK_FUNCTIONS: table.add_column(label, justify="center", ratio=1) for strategy in self.strategies: cells = [] diff --git a/tests/rust-python-harness/unit_tests_rust/strategy.json b/tests/rust-python-harness/unit_tests_rust/strategy.json index 89e897c872d..bfb2bc0dad0 100644 --- a/tests/rust-python-harness/unit_tests_rust/strategy.json +++ b/tests/rust-python-harness/unit_tests_rust/strategy.json @@ -7,6 +7,8 @@ "ocr": {"coverage": "planned", "selectors": []}, "messages": {"coverage": "planned", "selectors": []}, "responses": {"coverage": "planned", "selectors": []}, - "count_tokens": {"coverage": "planned", "selectors": []} + "count_tokens": {"coverage": "planned", "selectors": []}, + "chat_completions": {"coverage": "planned", "selectors": []}, + "transcription": {"coverage": "planned", "selectors": []} } } diff --git a/tests/rust-python-harness/validate_sub_methods/strategy.json b/tests/rust-python-harness/validate_sub_methods/strategy.json index 6e6381678e0..a26bc2d70ed 100644 --- a/tests/rust-python-harness/validate_sub_methods/strategy.json +++ b/tests/rust-python-harness/validate_sub_methods/strategy.json @@ -7,6 +7,8 @@ "ocr": {"coverage": "planned", "selectors": []}, "messages": {"coverage": "planned", "selectors": []}, "responses": {"coverage": "planned", "selectors": []}, - "count_tokens": {"coverage": "planned", "selectors": []} + "count_tokens": {"coverage": "planned", "selectors": []}, + "chat_completions": {"coverage": "planned", "selectors": []}, + "transcription": {"coverage": "planned", "selectors": []} } } diff --git a/tests/test_rust_python_harness.py b/tests/test_rust_python_harness.py index 6a8fa8d35cc..514446577fd 100644 --- a/tests/test_rust_python_harness.py +++ b/tests/test_rust_python_harness.py @@ -66,13 +66,14 @@ def _manifest() -> dict[str, object]: } -def test_should_load_the_three_harness_strategies_in_order() -> None: +def test_should_load_the_four_harness_strategies_in_order() -> None: strategies = load_catalog() assert [strategy.id for strategy in strategies] == [ "e2e_fuzz_tests", "unit_tests_rust", "validate_sub_methods", + "existing_e2e_test_sdk", ] assert all( tuple(case.sdk_function for case in strategy.cases) == SDK_FUNCTIONS @@ -104,6 +105,8 @@ def test_should_reject_a_manifest_missing_an_sdk_function(tmp_path: Path) -> Non True, ), ("tests/test_parity.py::test_one", "tests/test_parity.py::test_two", False), + ("tests/ocr_tests/", "tests/ocr_tests/test_ocr_mistral.py::test_one", True), + ("tests/ocr_tests/", "tests/other_tests/test_ocr_mistral.py::test_one", False), ], ) def test_should_match_pytest_file_and_node_selectors( @@ -123,6 +126,13 @@ def test_should_only_return_selectors_whose_files_exist(tmp_path: Path) -> None: assert runnable_selectors((case,), tmp_path) == ("tests/test_parity.py",) +def test_should_treat_an_existing_folder_selector_as_runnable(tmp_path: Path) -> None: + (tmp_path / "tests" / "ocr_tests").mkdir(parents=True) + case = _case(selectors=("tests/ocr_tests/",)) + + assert runnable_selectors((case,), tmp_path) == ("tests/ocr_tests/",) + + def test_should_mark_planned_and_not_applicable_cases_without_running() -> None: planned = CaseResult(case=_case(coverage=Coverage.PLANNED)) not_applicable = CaseResult(case=_case(coverage=Coverage.NOT_APPLICABLE)) @@ -239,8 +249,8 @@ def test_should_report_confidence_for_each_sdk_section() -> None: } assert scores["responses"].verified_strategies == 1 - assert scores["responses"].required_strategies == 3 - assert scores["responses"].percentage == 33 + assert scores["responses"].required_strategies == 4 + assert scores["responses"].percentage == 25 assert scores["responses"].level.value == "MEDIUM" assert scores["count_tokens"].percentage == 0 assert scores["count_tokens"].level.value == "LOW" From 291d02f8aa90e0ed2cc6a08b38c4976559f9cafe Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 20:32:43 -0700 Subject: [PATCH 22/23] fix(mcp): never exchange the LiteLLM virtual key as the upstream subject token (#39446) --- .../mcp_server/mcp_server_manager.py | 97 ++++-- .../proxy/_experimental/mcp_server/server.py | 9 + .../mcp_server/test_mcp_server.py | 6 +- .../mcp_server/test_mcp_server_manager.py | 311 +++++++++++++++++- 4 files changed, 396 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1f552ff3e13..1434fa5bfea 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -820,20 +820,46 @@ def _should_strip_caller_authorization( if not (mcp_server.is_oauth_passthrough or mcp_server.is_oauth_delegate): return False - normalized_raw_headers: Final = {str(k).lower(): v for k, v in (raw_headers or {}).items() if isinstance(k, str)} - has_explicit_litellm_admission_header: Final = normalized_raw_headers.get("x-litellm-api-key") is not None + has_explicit_litellm_admission_header: Final = _has_explicit_litellm_admission_header(raw_headers) if mcp_server.is_oauth_delegate: return not has_explicit_litellm_admission_header - admission_consumed_authorization_as_litellm_key: Final = ( - user_api_key_auth is not None - and bool(getattr(user_api_key_auth, "api_key", None)) - and not has_explicit_litellm_admission_header - ) - return admission_consumed_authorization_as_litellm_key or ( + return _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth) or ( user_api_key_auth is None and not has_explicit_litellm_admission_header ) +LITELLM_VIRTUAL_KEY_PREFIX: Final = "sk-" + + +def _raw_header_value(raw_headers: Mapping[str, str] | None, name: str) -> str | None: + return next((v for k, v in (raw_headers or {}).items() if isinstance(k, str) and k.lower() == name), None) + + +def _has_explicit_litellm_admission_header(raw_headers: Mapping[str, str] | None) -> bool: + """Admission only consumes a non-empty ``x-litellm-api-key``; an empty one falls back to ``Authorization``.""" + return bool(_raw_header_value(raw_headers, "x-litellm-api-key")) + + +def _authorization_is_litellm_admission_credential( + raw_headers: Mapping[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None, +) -> bool: + """True when ``Authorization`` carries the LiteLLM key admission validated. + + That is the case when no usable ``x-litellm-api-key`` was sent, or when the client repeated the + same key in both headers. + """ + if user_api_key_auth is None or not user_api_key_auth.api_key: + return False + admission_header: Final = _raw_header_value(raw_headers, "x-litellm-api-key") + if not admission_header: + return True + authorization: Final = _raw_header_value(raw_headers, "authorization") + return authorization is not None and strip_auth_scheme(authorization, "Bearer") == strip_auth_scheme( + admission_header, "Bearer" + ) + + def _format_byok_openapi_auth_header(mcp_server: MCPServer, mcp_auth_header: str) -> str: """Format a raw BYOK credential for OpenAPI tool ``Authorization`` injection. @@ -3277,8 +3303,8 @@ class MCPServerManager: ######################################################### @staticmethod def _extract_bearer_token( - oauth2_headers: dict[str, str] | None, - raw_headers: dict[str, str] | None, + oauth2_headers: Mapping[str, str] | None, + raw_headers: Mapping[str, str] | None, ) -> str | None: """Extract the bare Bearer token from oauth2_headers or raw_headers. @@ -3298,10 +3324,29 @@ class MCPServerManager: return auth_value return None + @staticmethod + def _extract_subject_token( + oauth2_headers: Mapping[str, str] | None, + raw_headers: Mapping[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None, + ) -> str | None: + """The caller's upstream identity token, or ``None`` when the bearer is a LiteLLM key. + + Rejects the key admission validated and, because virtual keys always carry the ``sk-`` prefix, + any other LiteLLM key a client puts in ``Authorization`` next to ``x-litellm-api-key``. + """ + if _authorization_is_litellm_admission_credential(raw_headers, user_api_key_auth): + return None + bearer: Final = MCPServerManager._extract_bearer_token(oauth2_headers, raw_headers) + if bearer is not None and bearer.startswith(LITELLM_VIRTUAL_KEY_PREFIX): + return None + return bearer + def _obo_subject_token( self, server: MCPServer, - raw_headers: dict[str, str] | None, + raw_headers: Mapping[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None, ) -> str | None: """The caller's bearer as the token_exchange (OBO) subject token, for that mode only. @@ -3311,7 +3356,7 @@ class MCPServerManager: """ if server.auth_type != MCPAuth.oauth2_token_exchange: return None - return self._extract_bearer_token(None, raw_headers) + return self._extract_subject_token(None, raw_headers, user_api_key_auth) def _build_stdio_env( self, @@ -3566,6 +3611,7 @@ class MCPServerManager: server: MCPServer, oauth2_headers: dict[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None, + raw_headers: Mapping[str, str] | None = None, ) -> None: """Run the OBO exchange for a caller-supplied subject at the transport edge. @@ -3577,13 +3623,15 @@ class MCPServerManager: """ if server.auth_type != MCPAuth.oauth2_token_exchange: return - subject_token: Final = self._extract_bearer_token(oauth2_headers, None) - if not subject_token: + if not self._extract_bearer_token(oauth2_headers, None): return resolved_server: Final = await self.ensure_oauth_metadata_discovered(server) spec: Final = to_server_spec(resolved_server) if spec is None or not isinstance(spec.config, TokenExchangeConfig): return + subject_token: Final = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth) + if subject_token is None: + raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path()) match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec): case Ok(_): return @@ -3851,7 +3899,7 @@ class MCPServerManager: # token (mirrors the call path), not v1's deleted client_credentials fallback. Other modes # never read the inbound bearer, so leave subject_token None to avoid forwarding it. subject_token: Final = ( - self._extract_bearer_token(oauth2_headers, raw_headers) + self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth) if server.auth_type == MCPAuth.oauth2_token_exchange else None ) @@ -3931,6 +3979,7 @@ class MCPServerManager: async def get_prompts_from_server( self, server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, add_prefix: bool = True, @@ -3959,7 +4008,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env: Final = self._build_stdio_env(server, raw_headers) - subject_token: Final = self._obo_subject_token(server, raw_headers) + subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) client = await self._create_mcp_client( server=server, @@ -3982,6 +4031,7 @@ class MCPServerManager: async def get_resources_from_server( self, server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, add_prefix: bool = True, @@ -4001,7 +4051,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env: Final = self._build_stdio_env(server, raw_headers) - subject_token: Final = self._obo_subject_token(server, raw_headers) + subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) client = await self._create_mcp_client( server=server, @@ -4024,6 +4074,7 @@ class MCPServerManager: async def get_resource_templates_from_server( self, server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, mcp_auth_header: str | dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, add_prefix: bool = True, @@ -4043,7 +4094,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env: Final = self._build_stdio_env(server, raw_headers) - subject_token: Final = self._obo_subject_token(server, raw_headers) + subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) client = await self._create_mcp_client( server=server, @@ -4068,6 +4119,7 @@ class MCPServerManager: async def read_resource_from_server( self, server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, url: AnyUrl, mcp_auth_header: str | dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -4084,7 +4136,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env: Final = self._build_stdio_env(server, raw_headers) - subject_token: Final = self._obo_subject_token(server, raw_headers) + subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) client: Final = await self._create_mcp_client( server=server, @@ -4099,6 +4151,7 @@ class MCPServerManager: async def get_prompt_from_server( self, server: MCPServer, + user_api_key_auth: UserAPIKeyAuth | None, prompt_name: str, arguments: dict[str, str] | None = None, mcp_auth_header: str | dict[str, str] | None = None, @@ -4116,7 +4169,7 @@ class MCPServerManager: extra_headers.update(server.static_headers) stdio_env: Final = self._build_stdio_env(server, raw_headers) - subject_token: Final = self._obo_subject_token(server, raw_headers) + subject_token: Final = self._obo_subject_token(server, raw_headers, user_api_key_auth) client: Final = await self._create_mcp_client( server=server, @@ -5290,7 +5343,7 @@ class MCPServerManager: MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag, ): - subject_token = self._extract_bearer_token(oauth2_headers, raw_headers) + subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth) elif mcp_server.auth_type == MCPAuth.oauth2: if mcp_server.has_client_credentials: # For M2M OAuth servers, Authorization must come from token fetch. @@ -5638,7 +5691,7 @@ class MCPServerManager: subject_token: str | None = None if isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)): - subject_token = self._extract_bearer_token(oauth2_headers, raw_headers) + subject_token = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth) elif isinstance(spec.config, PassthroughConfig): inbound_token, forwarded_headers = _take_forwarded_authorization(forwarded_headers) per_server_token: Final = _passthrough_token_from_mcp_auth_header(mcp_auth_header) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index af3ff6714af..bb075220530 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2189,6 +2189,7 @@ if MCP_AVAILABLE: try: prompts = await global_mcp_server_manager.get_prompts_from_server( server=server, + user_api_key_auth=user_api_key_auth, mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=True, # Always add server prefix @@ -2242,6 +2243,7 @@ if MCP_AVAILABLE: try: resources = await global_mcp_server_manager.get_resources_from_server( server=server, + user_api_key_auth=user_api_key_auth, mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=True, # Always add server prefix @@ -2293,6 +2295,7 @@ if MCP_AVAILABLE: try: resource_templates = await global_mcp_server_manager.get_resource_templates_from_server( server=server, + user_api_key_auth=user_api_key_auth, mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=True, # Always add server prefix @@ -3211,6 +3214,7 @@ if MCP_AVAILABLE: return await global_mcp_server_manager.get_prompt_from_server( server=server, + user_api_key_auth=user_api_key_auth, prompt_name=original_prompt_name, arguments=arguments, mcp_auth_header=server_auth_header, @@ -3261,6 +3265,7 @@ if MCP_AVAILABLE: return await global_mcp_server_manager.read_resource_from_server( server=server, + user_api_key_auth=user_api_key_auth, url=url, mcp_auth_header=server_auth_header, extra_headers=extra_headers, @@ -3723,6 +3728,7 @@ if MCP_AVAILABLE: user_api_key_auth: UserAPIKeyAuth | None, client_ip: str | None, allowed_server_ids: set[str] | None = None, + raw_headers: Mapping[str, str] | None = None, ) -> None: """Fail fast with HTTP 401 for MCP servers that need user auth but didn't receive it on this request. Covers both gateway-managed OAuth2 @@ -3867,6 +3873,7 @@ if MCP_AVAILABLE: server=server, oauth2_headers=oauth2_headers, user_api_key_auth=user_api_key_auth, + raw_headers=raw_headers, ) # Pass-through OAuth: when the admin has opted a server into @@ -4195,6 +4202,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, client_ip=_client_ip, allowed_server_ids=toolset_allowed_server_ids, + raw_headers=raw_headers, ) # Pre-flight auth check for pass-through servers. Must run after @@ -4518,6 +4526,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, client_ip=_sse_client_ip, allowed_server_ids=toolset_allowed_server_ids, + raw_headers=raw_headers, ) # Pre-flight auth check for pass-through servers: surface upstream diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 3f6d8f8837c..ff0007f47aa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -914,6 +914,7 @@ async def test_mcp_get_prompt_success(): ) mock_manager.get_prompt_from_server.assert_awaited_once_with( server=server, + user_api_key_auth=user_api_key_auth, prompt_name="hello", arguments={"foo": "bar"}, mcp_auth_header={"Authorization": "token"}, @@ -976,6 +977,7 @@ async def test_mcp_read_resource_success(): ) mock_manager.read_resource_from_server.assert_awaited_once_with( server=server, + user_api_key_auth=user_api_key_auth, url="https://example.com/resource", mcp_auth_header={"Authorization": "token"}, extra_headers={"X-Test": "1"}, @@ -8268,7 +8270,9 @@ class TestOboPreflightScopedToAllowedServers: _, preflight = await self._run(requested, allowed=[requested], user_api_key_auth=key) - preflight.assert_awaited_once_with(server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key) + preflight.assert_awaited_once_with( + server=requested, oauth2_headers=self.SUBJECT_HEADERS, user_api_key_auth=key, raw_headers=None + ) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 5508259273d..5cde2e83f62 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -30,6 +30,7 @@ from mcp.types import ( TextResourceContents, ) from mcp.types import Tool as MCPTool +from pydantic import AnyUrl from litellm.constants import MCP_METADATA_TIMEOUT from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -2270,7 +2271,9 @@ class TestMCPServerManager: """prompts/list on an OBO server must exchange the caller's bearer, not connect with none.""" server = self._token_exchange_server("te-prompts") st = await self._capture_subject_token( - lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"}) + lambda m: m.get_prompts_from_server( + server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"} + ) ) assert st == "subj-jwt" @@ -2279,7 +2282,9 @@ class TestMCPServerManager: """resources/list on an OBO server must exchange the caller's bearer.""" server = self._token_exchange_server("te-resources") st = await self._capture_subject_token( - lambda m: m.get_resources_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"}) + lambda m: m.get_resources_from_server( + server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"} + ) ) assert st == "subj-jwt" @@ -2290,6 +2295,7 @@ class TestMCPServerManager: st = await self._capture_subject_token( lambda m: m.read_resource_from_server( server=server, + user_api_key_auth=None, url="https://up.example.com/r", raw_headers={"authorization": "Bearer subj-jwt"}, ) @@ -2307,7 +2313,9 @@ class TestMCPServerManager: auth_type=MCPAuth.none, ) st = await self._capture_subject_token( - lambda m: m.get_prompts_from_server(server=server, raw_headers={"authorization": "Bearer subj-jwt"}) + lambda m: m.get_prompts_from_server( + server=server, user_api_key_auth=None, raw_headers={"authorization": "Bearer subj-jwt"} + ) ) assert st is None @@ -3254,7 +3262,7 @@ class TestMCPServerManager: new_callable=AsyncMock, return_value=mock_client, ): - prompts = await manager.get_prompts_from_server(server, add_prefix=True) + prompts = await manager.get_prompts_from_server(server, user_api_key_auth=None, add_prefix=True) mock_client.list_prompts.assert_awaited_once() assert len(prompts) == 1 @@ -3289,6 +3297,7 @@ class TestMCPServerManager: ): result = await manager.get_prompt_from_server( server=server, + user_api_key_auth=None, prompt_name="hello", arguments={"tone": "casual"}, ) @@ -3334,6 +3343,7 @@ class TestMCPServerManager: ): result = await manager.get_resources_from_server( server=server, + user_api_key_auth=None, mcp_auth_header="auth", extra_headers={"X-Test": "1"}, add_prefix=True, @@ -3391,6 +3401,7 @@ class TestMCPServerManager: ): result = await manager.get_resource_templates_from_server( server=server, + user_api_key_auth=None, mcp_auth_header="auth", extra_headers=None, add_prefix=False, @@ -3441,6 +3452,7 @@ class TestMCPServerManager: ) as mock_create_client: result = await manager.read_resource_from_server( server=server, + user_api_key_auth=None, url="https://example.com/resource", mcp_auth_header="auth", extra_headers={"X-Test": "1"}, @@ -11006,3 +11018,294 @@ class TestOpenApiHandlerRelaysUpstreamAuth: assert result.isError is True assert "upstream returned HTTP 503" in result.content[0].text + + +class TestLitellmAdmissionKeyIsNeverTheSubjectToken: + """The bearer that admitted the request as a LiteLLM key must not be sent to the IdP as the + RFC 8693 subject_token (or ID-JAG assertion). Only ``x-litellm-api-key`` disambiguates: with it + present, ``Authorization`` is the caller's own identity token and is exchanged as before.""" + + _ADMISSION_KEY: Final = "sk-litellm-virtual-key" + _USER_TOKEN: Final = "user-idp-jwt" + + @staticmethod + def _token_exchange_server(server_id: str) -> MCPServer: + return MCPServer( + server_id=server_id, + name=f"{server_id}-server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + + @staticmethod + def _id_jag_server(server_id: str) -> MCPServer: + return MCPServer( + server_id=server_id, + name=f"{server_id}-server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_id_jag, + client_id="cid", + client_secret="csec", + token_exchange_endpoint="https://idp.example.com/token", + id_jag_resource_token_endpoint="https://resource-as.example.com/token", + ) + + @staticmethod + def _recording_provider() -> MagicMock: + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + provider: Final = MagicMock() + provider.resolve_credentials = AsyncMock( + return_value=Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) + ) + return provider + + @staticmethod + def _subjects_seen_by(provider: MagicMock) -> list[str | None]: + return [ + call.args[0].inbound_token.get_secret_value() if call.args[0].inbound_token else None + for call in provider.resolve_credentials.call_args_list + ] + + @staticmethod + def _manager_with_recording_client() -> MCPServerManager: + manager: Final = MCPServerManager() + client: Final = AsyncMock() + client.call_tool = AsyncMock(return_value=CallToolResult(content=[], isError=False)) + client.list_prompts = AsyncMock(return_value=[]) + client.read_resource = AsyncMock(return_value=ReadResourceResult(contents=[])) + manager._create_mcp_client = AsyncMock(return_value=client) + return manager + + @staticmethod + def _subject_token_given_to_client(manager: MCPServerManager) -> str | None: + return manager._create_mcp_client.call_args.kwargs["subject_token"] + + async def _call_tool_subject(self, server: MCPServer, oauth2_headers, raw_headers, user_api_key_auth): + manager: Final = self._manager_with_recording_client() + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="tool", + arguments={}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + proxy_logging_obj=None, + user_api_key_auth=user_api_key_auth, + ) + return self._subject_token_given_to_client(manager) + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag]) + async def test_tools_call_with_only_the_litellm_key_has_no_subject(self, auth_type): + server = ( + self._token_exchange_server("te-call") + if auth_type == MCPAuth.oauth2_token_exchange + else self._id_jag_server("jag-call") + ) + subject_token = await self._call_tool_subject( + server, + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token is None + + @pytest.mark.asyncio + async def test_rest_tools_call_with_only_the_litellm_key_has_no_subject(self): + """The REST facade passes no oauth2_headers; the bearer is reached through raw_headers only.""" + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-rest"), + oauth2_headers=None, + raw_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token is None + + @pytest.mark.asyncio + async def test_tools_call_exchanges_the_user_token_when_x_litellm_api_key_admits(self): + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-split"), + oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"}, + raw_headers={ + "X-LiteLLM-API-Key": f"Bearer {self._ADMISSION_KEY}", + "authorization": f"Bearer {self._USER_TOKEN}", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token == self._USER_TOKEN + + @pytest.mark.asyncio + async def test_tools_call_with_an_empty_x_litellm_api_key_has_no_subject(self): + """Admission ignores an empty ``x-litellm-api-key`` and validates ``Authorization`` instead.""" + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-empty-header"), + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + raw_headers={"x-litellm-api-key": "", "authorization": f"Bearer {self._ADMISSION_KEY}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token is None + + @pytest.mark.asyncio + async def test_tools_call_with_the_same_litellm_key_in_both_headers_has_no_subject(self): + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-same-key"), + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + raw_headers={ + "x-litellm-api-key": self._ADMISSION_KEY, + "authorization": f"Bearer {self._ADMISSION_KEY}", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token is None + + @pytest.mark.asyncio + async def test_tools_call_with_a_different_litellm_key_in_authorization_has_no_subject(self): + """A second ``sk-`` virtual key next to ``x-litellm-api-key`` is still a gateway credential.""" + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-second-key"), + oauth2_headers={"Authorization": "Bearer sk-another-virtual-key"}, + raw_headers={ + "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}", + "authorization": "Bearer sk-another-virtual-key", + }, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert subject_token is None + + @pytest.mark.asyncio + async def test_tools_call_exchanges_the_bearer_when_jwt_admission_left_api_key_unset(self): + subject_token = await self._call_tool_subject( + self._token_exchange_server("te-jwt"), + oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"}, + raw_headers={"authorization": f"Bearer {self._USER_TOKEN}"}, + user_api_key_auth=UserAPIKeyAuth(api_key=None, user_id="alice"), + ) + assert subject_token == self._USER_TOKEN + + @pytest.mark.asyncio + async def test_tools_list_with_only_the_litellm_key_has_no_subject(self): + manager: Final = self._manager_with_recording_client() + manager._fetch_tools_with_timeout = AsyncMock(return_value=[]) + await manager._get_tools_from_server( + server=self._token_exchange_server("te-list-key"), + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + ) + assert self._subject_token_given_to_client(manager) is None + + @pytest.mark.asyncio + async def test_prompts_list_with_only_the_litellm_key_has_no_subject(self): + manager: Final = self._manager_with_recording_client() + await manager.get_prompts_from_server( + server=self._token_exchange_server("te-prompts-key"), + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + ) + assert self._subject_token_given_to_client(manager) is None + + @pytest.mark.asyncio + async def test_resource_read_with_only_the_litellm_key_has_no_subject(self): + manager: Final = self._manager_with_recording_client() + await manager.read_resource_from_server( + server=self._token_exchange_server("te-read-key"), + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + url=AnyUrl("file:///notes.txt"), + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + ) + assert self._subject_token_given_to_client(manager) is None + + @pytest.mark.asyncio + async def test_resource_read_exchanges_the_user_token_when_x_litellm_api_key_admits(self): + manager: Final = self._manager_with_recording_client() + await manager.read_resource_from_server( + server=self._token_exchange_server("te-read-split"), + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + url=AnyUrl("file:///notes.txt"), + raw_headers={ + "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}", + "authorization": f"Bearer {self._USER_TOKEN}", + }, + ) + assert self._subject_token_given_to_client(manager) == self._USER_TOKEN + + @pytest.mark.asyncio + async def test_openapi_call_never_hands_the_litellm_key_to_the_exchanger(self): + provider: Final = self._recording_provider() + manager = MCPServerManager(cred_provider=provider) + server = MCPServer( + server_id="te-openapi", + name="te_openapi", + server_name="te_openapi", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2_token_exchange, + token_exchange_endpoint="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + spec_path="https://api.example.com/openapi.json", + ) + user_auth = UserAPIKeyAuth(api_key="hashed-key", user_id="alice") + + await manager.resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + mcp_auth_header=None, + user_api_key_auth=user_auth, + forwarded_headers=None, + ) + await manager.resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"}, + raw_headers={ + "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}", + "authorization": f"Bearer {self._USER_TOKEN}", + }, + mcp_auth_header=None, + user_api_key_auth=user_auth, + forwarded_headers=None, + ) + assert self._subjects_seen_by(provider) == [None, self._USER_TOKEN] + + @pytest.mark.asyncio + async def test_preflight_challenges_instead_of_exchanging_the_litellm_key(self): + provider: Final = self._recording_provider() + manager = MCPServerManager(cred_provider=provider) + + with pytest.raises(HTTPException) as exc_info: + await manager.preflight_token_exchange( + server=self._token_exchange_server("te-preflight-key"), + oauth2_headers={"Authorization": f"Bearer {self._ADMISSION_KEY}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + raw_headers={"authorization": f"Bearer {self._ADMISSION_KEY}"}, + ) + assert exc_info.value.status_code == 401 + headers = exc_info.value.headers or {} + assert "resource_metadata" in (headers.get("WWW-Authenticate") or headers.get("www-authenticate") or "") + assert self._subjects_seen_by(provider) == [] + + @pytest.mark.asyncio + async def test_preflight_exchanges_the_user_token_when_x_litellm_api_key_admits(self): + provider: Final = self._recording_provider() + manager = MCPServerManager(cred_provider=provider) + + await manager.preflight_token_exchange( + server=self._token_exchange_server("te-preflight-split"), + oauth2_headers={"Authorization": f"Bearer {self._USER_TOKEN}"}, + user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="alice"), + raw_headers={ + "x-litellm-api-key": f"Bearer {self._ADMISSION_KEY}", + "authorization": f"Bearer {self._USER_TOKEN}", + }, + ) + assert self._subjects_seen_by(provider) == [self._USER_TOKEN] From e058aa68c4331ce7ffce2802528eacc2f2333642 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 20:40:00 -0700 Subject: [PATCH 23/23] test: add mistral ocr transformation parity coverage (#39482) * test: cover mistral ocr transformation parity Co-Authored-By: Claude Code * test: map mistral ocr parity contracts Co-Authored-By: Claude Code --------- Co-authored-by: Claude Code --- .../providers/mistral/ocr/transformation.rs | 188 ++- .../ledgers/ocr/ocr_test_ledger.json | 1240 ++++++++++++++--- 2 files changed, 1208 insertions(+), 220 deletions(-) diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index 9648321d7ff..0125886aac1 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -186,68 +186,180 @@ mod tests { use serde_json::json; #[test] - fn supported_params_match_python_mistral_ocr_config() { + fn extract_header_is_a_supported_ocr_param() { + assert!(supported_ocr_params().contains(&"extract_header")); + } + + #[test] + fn extract_footer_is_a_supported_ocr_param() { + assert!(supported_ocr_params().contains(&"extract_footer")); + } + + #[test] + fn existing_ocr_params_remain_supported() { + for param in [ + "pages", + "include_image_base64", + "image_limit", + "image_min_size", + "bbox_annotation_format", + "document_annotation_format", + ] { + assert!(supported_ocr_params().contains(¶m)); + } + } + + #[test] + fn map_ocr_params_forwards_extract_header() { + let params = json!({"extract_header": true}); assert_eq!( - supported_ocr_params(), - &[ - "pages", - "include_image_base64", - "image_limit", - "image_min_size", - "bbox_annotation_format", - "document_annotation_format", - "document_annotation_prompt", - "extract_header", - "extract_footer", - "table_format", - "confidence_scores_granularity", - "include_blocks", - "id", - ] + map_ocr_params(params.as_object().unwrap()), + params.as_object().unwrap().clone() + ); + } + + #[test] + fn map_ocr_params_forwards_extract_footer() { + let params = json!({"extract_footer": true}); + assert_eq!( + map_ocr_params(params.as_object().unwrap()), + params.as_object().unwrap().clone() + ); + } + + #[test] + fn map_ocr_params_forwards_extract_header_and_footer() { + let params = json!({"extract_header": true, "extract_footer": false}); + assert_eq!( + map_ocr_params(params.as_object().unwrap()), + params.as_object().unwrap().clone() ); } #[test] fn map_ocr_params_drops_unknown_params() { - let params = json!({ - "extract_header": true, - "unsupported_param": "value", - "pages": [0, 1] - }); + let params = json!({"extract_header": true, "unsupported_param": "value"}); let mapped = map_ocr_params(params.as_object().unwrap()); - assert_eq!(mapped.get("extract_header"), Some(&json!(true))); - assert_eq!(mapped.get("pages"), Some(&json!([0, 1]))); assert!(!mapped.contains_key("unsupported_param")); } #[test] - fn transform_ocr_request_builds_mistral_body() { + fn new_ocr_params_are_supported() { + for param in [ + "table_format", + "confidence_scores_granularity", + "document_annotation_prompt", + "include_blocks", + "id", + ] { + assert!(supported_ocr_params().contains(¶m)); + } + } + + #[test] + fn map_ocr_params_forwards_new_ocr_params() { + for (param, value) in [ + ("table_format", json!("html")), + ("confidence_scores_granularity", json!("word")), + ( + "document_annotation_prompt", + json!("Extract all invoice line items"), + ), + ("include_blocks", json!(true)), + ("id", json!("req-123")), + ] { + let params = json!({param: value}); + assert_eq!( + map_ocr_params(params.as_object().unwrap()), + params.as_object().unwrap().clone() + ); + } + } + + #[test] + fn transform_ocr_request_includes_each_optional_param() { + let document = json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }); + for (param, value) in [ + ("table_format", json!("html")), + ("confidence_scores_granularity", json!("word")), + ( + "document_annotation_prompt", + json!("Extract all invoice line items"), + ), + ("id", json!("req-123")), + ("extract_header", json!(true)), + ("include_blocks", json!(true)), + ("pages", json!([0, 1])), + ] { + let result = transform_ocr_request( + "mistral-ocr-latest", + document.clone(), + json!({param: value}).as_object().unwrap().clone(), + ) + .expect("request should transform"); + assert_eq!(result.data.get(param), Some(&value)); + assert_eq!(result.data.get("model"), Some(&json!("mistral-ocr-latest"))); + assert_eq!(result.data.get("document"), Some(&document)); + assert_eq!(result.files, None); + } + } + + #[test] + fn transform_ocr_request_includes_multiple_new_params() { let document = json!({ "type": "document_url", "document_url": "https://example.com/doc.pdf" }); let optional_params = json!({ - "include_image_base64": true, - "table_format": "html" + "table_format": "html", + "confidence_scores_granularity": "page", + "extract_header": true }) .as_object() .unwrap() .clone(); - - let result = transform_ocr_request("mistral-ocr-latest", document.clone(), optional_params) + let result = transform_ocr_request("mistral-ocr-latest", document, optional_params) .expect("request should transform"); - + assert_eq!(result.data.get("table_format"), Some(&json!("html"))); assert_eq!( - result.data, - json!({ - "model": "mistral-ocr-latest", - "document": document, - "include_image_base64": true, - "table_format": "html" - }) + result.data.get("confidence_scores_granularity"), + Some(&json!("page")) ); - assert_eq!(result.files, None); + assert_eq!(result.data.get("extract_header"), Some(&json!(true))); + } + + #[test] + fn transform_ocr_response_preserves_blocks_and_confidence_scores() { + let blocks = json!([{"type": "title", "content": "Invoice"}]); + let confidence_scores = json!({"page": 0.98}); + let response = json!({ + "pages": [{"index": 0, "markdown": "# Invoice", "blocks": blocks, "confidence_scores": confidence_scores}], + "model": "mistral-ocr-4-0", + "usage_info": {"pages_processed": 1} + }); + let result = + transform_ocr_response("mistral-ocr-4-0", response).expect("response should transform"); + assert_eq!(result.pages[0].get("blocks"), Some(&blocks)); + assert_eq!( + result.pages[0].get("confidence_scores"), + Some(&confidence_scores) + ); + } + + #[test] + fn transform_ocr_response_preserves_ocr4_page_fields() { + let response = json!({ + "pages": [{"index": 0, "markdown": "table page", "tables": [{"rows": 2, "cols": 3}], "hyperlinks": ["https://example.com"], "header": "Acme Corp", "footer": "Page 1"}], + "model": "mistral-ocr-4-0", + "usage_info": {"pages_processed": 1} + }); + let result = transform_ocr_response("mistral-ocr-4-0", response.clone()) + .expect("response should transform"); + assert_eq!(result.pages[0], response["pages"][0]); } #[test] diff --git a/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json b/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json index 799a1320463..e617ddf8f94 100644 --- a/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json +++ b/tests/rust-python-harness/strategies/unit_tests/ledgers/ocr/ocr_test_ledger.json @@ -20,189 +20,1065 @@ "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs" ], "entries": [ - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_should_encode_azure_document_intelligence_model_id", "status": "unmapped", "reason": "model-id URL percent-encoding has no Rust test; Rust only tests pages/features query building"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_should_reject_dot_segment_azure_document_intelligence_model_id", "status": "unmapped", "reason": "model-id dot-segment validation has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_preserves_azure_native_fields", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_response_normalizes_pages", "justification": "both assert page markdown, dimension (inch-to-pixel) normalization, and usage_info.pages_processed from the same Azure succeeded response shape"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_async_transform_ocr_response_preserves_azure_native_fields", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_response_normalizes_pages", "justification": "async twin of the sync case above, same underlying transform is exercised on the Rust side"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_tolerates_missing_native_fields", "status": "unmapped", "reason": "tables/keyValuePairs absence tolerance is not asserted by the Rust response test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_non_succeeded_status_raises", "status": "unmapped", "reason": "no Rust test asserts on a non-succeeded Azure DI status"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_supported_ocr_params_includes_features", "status": "unmapped", "reason": "supported-params list content has no Rust equivalent for Azure"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_native_format_carries_raw_operation", "status": "unmapped", "reason": "native req_format raw-operation passthrough is not tested in Rust"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_async_transform_ocr_response_native_format_carries_raw_operation", "status": "unmapped", "reason": "native req_format raw-operation passthrough is not tested in Rust"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_transform_ocr_response_default_format_omits_raw_operation", "status": "unmapped", "reason": "req_format gating of raw-operation output has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_passes_through_req_format", "status": "unmapped", "reason": "req_format passthrough in map_ocr_params has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_rejects_unknown_req_format_as_bad_request", "status": "unmapped", "reason": "req_format validation error path has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_omits_req_format_query_param", "status": "unmapped", "reason": "no Rust test asserts req_format is excluded from the built URL"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_features", "status": "unmapped", "reason": "features-string normalization in map_ocr_params has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_empty_features_list_omitted", "status": "unmapped", "reason": "empty-features omission has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_map_ocr_params_invalid_features_raises", "status": "unmapped", "reason": "features validation error path has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_appends_features_query", "status": "unmapped", "reason": "features query-param construction has no Rust test"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_get_complete_url_combines_pages_and_features", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_url_normalizes_zero_based_pages", "justification": "both assert 0-based, duplicate page indices are deduped, sorted, and rewritten 1-based into the request URL"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_validate_environment_uses_subscription_key", "status": "unmapped", "reason": "Python-side header derivation from litellm_params; Rust's poll test only checks the header is present, not how it was resolved"}, - {"python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", "python_test": "test_validate_environment_falls_back_to_entra_token", "status": "unmapped", "reason": "Entra bearer-token fallback logic has no Rust test"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_doc_intelligence_route", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_documentintelligence_and_is_case_insensitive", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestIsAzureDocumentIntelligenceModel::test_does_not_match_mistral_route", "status": "unmapped", "reason": "model-route string matching is Python-only dispatch logic"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_does_not_hijack_doc_intelligence", "status": "unmapped", "reason": "api_base resolution from the secret manager runs before the Rust bridge is called, no Rust test exists for it"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_explicit_api_base_is_honoured_for_doc_intelligence", "status": "unmapped", "reason": "api_base precedence resolution is Python-only"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_still_applies_to_mistral_ocr", "status": "unmapped", "reason": "api_base precedence resolution is Python-only"}, - - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_use_litellm_rust_toggles_flag", "status": "unmapped", "reason": "bridge-plumbing: Python-side feature-flag toggle, no Rust equivalent"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_env_var_enables_rust_ocr", "status": "unmapped", "reason": "bridge-plumbing: Python-side env-var flag gating"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_returns_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: dependency-injection test hook, not provider behavior"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_loader_returns_none_when_extension_absent", "status": "unmapped", "reason": "bridge-plumbing: native-extension import/loader fallback"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_loader_caches_absent_extension", "status": "unmapped", "reason": "bridge-plumbing: loader caching behavior"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_native_bridge_available_reflects_loader", "status": "unmapped", "reason": "bridge-plumbing: loader availability check"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_aocr_returns_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: dependency-injection test hook"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_toggle_without_ocr_arg_preserves_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: injected-impl state retention regression"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_explicit_ocr_none_clears_injected_impl", "status": "unmapped", "reason": "bridge-plumbing: injected-impl clearing behavior"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_none_when_extension_absent", "status": "unmapped", "reason": "bridge-plumbing: degrade path when the native extension is missing"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_load_rust_ocr_uses_compiled_extension", "status": "unmapped", "reason": "bridge-plumbing: native module resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_timeout_to_seconds_handles_float_timeout_and_none", "status": "unmapped", "reason": "bridge-plumbing: Python-side timeout normalization helper"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_bridge_wrapper_forwards_prepared_args_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: wrapper argument forwarding, asserted against a fake bridge not the real Rust code"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: async wrapper argument forwarding"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_prepares_request_and_wraps_response", "status": "unmapped", "reason": "bridge-plumbing: request preparation and response wrapping in Python"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_resolves_key_via_secret_manager_when_missing", "status": "unmapped", "reason": "secret-manager: API key resolution happens in Python before the bridge is invoked"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_prefers_explicit_key_over_resolver", "status": "unmapped", "reason": "secret-manager: key precedence resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_uses_provider_api_key_env_var", "status": "unmapped", "reason": "secret-manager: provider-specific env var name resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_forwards_vertex_routing_metadata", "status": "unmapped", "reason": "secret-manager: vertex routing metadata merge happens in Python"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager", "status": "unmapped", "reason": "secret-manager: vertex project/location resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager", "status": "unmapped", "reason": "secret-manager: azure_ai api_base resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint", "status": "unmapped", "reason": "secret-manager: doc-intelligence endpoint resolution"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_run_rust_ocr_runs_pre_call_logging", "status": "unmapped", "reason": "bridge-plumbing: Python logging-object pre_call invocation"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_routes_to_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: routing to a fake bridge, not the real Rust transform"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_routes_azure_ai_to_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: provider-prefix stripping before routing"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_rust_path_converts_file_document_before_bridge", "status": "unmapped", "reason": "file-normalization: raw-bytes-to-data-URI conversion happens in Python before the bridge call"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_exception_type_uses_resolved_provider_context", "status": "unmapped", "reason": "bridge-plumbing: Python exception-type mapping on bridge failure"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_aocr_routes_to_async_rust_when_enabled", "status": "unmapped", "reason": "bridge-plumbing: async routing to a fake bridge"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_aocr_exception_type_uses_resolved_provider_context", "status": "unmapped", "reason": "bridge-plumbing: async exception-type mapping on bridge failure"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_does_not_route_to_rust_when_disabled", "status": "unmapped", "reason": "bridge-plumbing: Python control flow for the toggle-disabled branch, no Rust-owned behavior runs"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_falls_back_to_python_when_bridge_unavailable", "status": "unmapped", "reason": "bridge-plumbing: Python-only fallback when the compiled Rust extension is absent, Rust cannot test its own absence"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_forwards_timeout_to_rust", "status": "unmapped", "reason": "bridge-plumbing: asserts the Python call site forwards a timeout kwarg, Rust receives an already-constructed request"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_passes_default_request_timeout_to_rust", "status": "unmapped", "reason": "bridge-plumbing: asserts the Python call site supplies a default timeout kwarg, no Rust equivalent"}, - {"python_file": "tests/test_litellm/ocr/test_rust_bridge.py", "python_test": "test_ocr_provider_configs_expose_api_key_env_vars", "status": "unmapped", "reason": "asserts per-provider get_api_key_env_var() strings; the closest Rust test (ocr_dispatch_supports_migrated_providers) asserts provider dispatch/param resolution instead, not API key env var names"}, - - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_pdf_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection is Python-only preprocessing before the bridge call"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_png_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_jpg_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_jpeg_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_gif_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_webp_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_tiff_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_tif_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_detect_bmp_mime_type", "status": "unmapped", "reason": "file-normalization: MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_be_case_insensitive", "status": "unmapped", "reason": "file-normalization: MIME detection case handling"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestGetMimeType::test_should_fallback_for_unknown_extension", "status": "unmapped", "reason": "file-normalization: MIME detection fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pdf_pathlib_path_to_document_url", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion happens in Python"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_image_pathlib_path_to_image_url", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_reject_bare_str_path", "status": "unmapped", "reason": "file-normalization: arbitrary-file-read guard on bare str paths"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pathlib_path", "status": "unmapped", "reason": "file-normalization: local-path-to-data-URI conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes", "status": "unmapped", "reason": "file-normalization: raw-bytes-to-data-URI conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_explicit_mime_type", "status": "unmapped", "reason": "file-normalization: explicit MIME override on raw bytes"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_image_mime_type", "status": "unmapped", "reason": "file-normalization: explicit MIME override on raw bytes"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object", "status": "unmapped", "reason": "file-normalization: file-like-object conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object_with_name", "status": "unmapped", "reason": "file-normalization: file-like-object name-based MIME detection"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_missing_file_field", "status": "unmapped", "reason": "file-normalization: missing-field validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_nonexistent_pathlib_path", "status": "unmapped", "reason": "file-normalization: missing-file validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_empty_file", "status": "unmapped", "reason": "file-normalization: empty-file validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_unsupported_type", "status": "unmapped", "reason": "file-normalization: unsupported input type validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_invalid_mime_type", "status": "unmapped", "reason": "file-normalization: MIME-type injection validation"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestConvertFileDocumentToUrlDocument::test_should_override_mime_type_for_pathlib_path", "status": "unmapped", "reason": "file-normalization: explicit MIME override precedence"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_document_url_for_pdf", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_png", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_jpeg", "status": "unmapped", "reason": "file-normalization: multipart upload conversion"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_octet_stream", "status": "unmapped", "reason": "file-normalization: filename-based MIME fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_none", "status": "unmapped", "reason": "file-normalization: filename-based MIME fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_fallback_to_octet_stream_for_unknown", "status": "unmapped", "reason": "file-normalization: default MIME fallback"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_preserve_base64_content_correctly", "status": "unmapped", "reason": "file-normalization: binary round-trip through base64"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_from_content_type", "status": "unmapped", "reason": "file-normalization: content-type parameter stripping"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_with_multiple_params", "status": "unmapped", "reason": "file-normalization: content-type parameter stripping"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_reject_file_type_document_in_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body file-type guard, a different mechanism than Rust's URL-fetch SSRF guard"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_accept_document_url_type_in_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body parsing"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_raise_on_invalid_json_body", "status": "unmapped", "reason": "proxy-layer JSON-body parsing error path"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", "python_test": "TestProxySecurityGuard::test_should_ignore_document_form_field_injection", "status": "unmapped", "reason": "proxy-layer multipart form-field injection guard, a different mechanism than Rust's URL-fetch SSRF guard"}, - - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_extract_header_in_supported_params", "status": "unmapped", "reason": "Rust's fixed-list test checks the full list as one assertion, not this individual param"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_extract_footer_in_supported_params", "status": "unmapped", "reason": "Rust's fixed-list test checks the full list as one assertion, not this individual param"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestGetSupportedOcrParams::test_existing_params_still_present", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "supported_params_match_python_mistral_ocr_config", "justification": "both assert the full supported_ocr_params list matches the same fixed set of param names"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_header_passed_through", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "map_ocr_params_drops_unknown_params", "justification": "both assert extract_header survives map_ocr_params filtering unchanged"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_footer_passed_through", "status": "unmapped", "reason": "Rust's map_ocr_params test does not assert on extract_footer specifically"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_extract_header_and_footer_together", "status": "unmapped", "reason": "combined extract_header+extract_footer passthrough is not asserted together in Rust"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestMapOcrParams::test_unknown_param_is_dropped", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "map_ocr_params_drops_unknown_params", "justification": "both assert an unrecognized param key is dropped while a known one is kept"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestNewSupportedParams::test_new_param_in_supported_list", "status": "unmapped", "reason": "OCR4-specific new params (table_format etc) are not individually verified against the Rust fixed-list test"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestNewParamsMapOcr::test_new_param_passed_through", "status": "unmapped", "reason": "OCR4-specific new params are not individually asserted in the Rust map_ocr_params test"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrRequest::test_param_included_in_request_body", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_builds_mistral_body", "justification": "both assert an optional param value ends up in the built request body alongside model/document"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrRequest::test_multiple_new_params_together", "status": "mapped", "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_builds_mistral_body", "justification": "both assert multiple optional params (table_format/include_image_base64) land correctly in the same request body"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrResponseOcr4Fields::test_blocks_and_confidence_scores_preserved", "status": "unmapped", "reason": "OCR4 blocks/confidence_scores fields are not asserted by the Rust response test"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", "python_test": "TestTransformOcrResponseOcr4Fields::test_ocr4_fields_survive_model_dump", "status": "unmapped", "reason": "OCR4 tables/hyperlinks/header/footer fields are not asserted by the Rust response test"}, - - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_model_info_ocr4_price", "status": "unmapped", "reason": "cost-calc: pricing/model-info lookup is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr4_cost_scales_with_pages", "status": "unmapped", "reason": "cost-calc: per-page pricing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_pricing_entry", "status": "unmapped", "reason": "cost-calc: cost-map JSON entry validation is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_model_info_price", "status": "unmapped", "reason": "cost-calc: pricing/model-info lookup is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_cost_scales_with_pages", "status": "unmapped", "reason": "cost-calc: per-page pricing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_ocr_and_annotation_pages_at_their_own_rates", "status": "unmapped", "reason": "cost-calc: mixed-rate billing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_annotation_only_response", "status": "unmapped", "reason": "cost-calc: annotation-only billing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_ocr3_bills_annotation_pages_when_pages_processed_missing", "status": "unmapped", "reason": "cost-calc: fallback billing math is Python-only"}, - {"python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", "python_test": "test_azure_doc_ai_annotation_pages_fall_back_to_ocr_rate", "status": "unmapped", "reason": "cost-calc: fallback billing math is Python-only"}, - - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_rust_ocr_serves_default_format", "status": "unmapped", "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_rust_ocr_skipped_for_native_format", "status": "unmapped", "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_native_format_rejected_for_provider_without_support_as_bad_request", "status": "unmapped", "reason": "provider-support validation for req_format happens in Python"}, - {"python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", "python_test": "test_unknown_format_rejected_for_provider_without_support_as_bad_request", "status": "unmapped", "reason": "req_format validation error path is Python-only"}, - - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestHandlerDiscovery::test_handler_discovered_for_ocr", "status": "unmapped", "reason": "guardrail-translation handler discovery is a Python proxy-layer concern"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestHandlerDiscovery::test_handler_discovered_for_aocr", "status": "unmapped", "reason": "guardrail-translation handler discovery is a Python proxy-layer concern"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_document_url", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_image_url", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_no_document", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_process_invalid_document", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestInputProcessing::test_input_blocking_guardrail", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_single_page", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_multiple_pages", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_empty_pages", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_page_with_empty_markdown", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_process_preserves_page_metadata", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestOutputProcessing::test_output_blocking_guardrail", "status": "unmapped", "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle"}, - {"python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", "python_test": "TestPIIMaskingScenario::test_pii_masking_in_ocr_pages", "status": "unmapped", "reason": "PII redaction in the translation handler has no Rust equivalent"}, - - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_read_req_format_from_header", "status": "unmapped", "reason": "proxy-layer header parsing has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_prefer_body_req_format_over_header", "status": "unmapped", "reason": "proxy-layer body-vs-header precedence has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_omit_req_format_when_header_absent", "status": "unmapped", "reason": "proxy-layer parsing has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_reject_unknown_req_format", "status": "unmapped", "reason": "proxy-layer validation has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_return_native_payload_with_litellm_response_headers", "status": "unmapped", "reason": "proxy-layer response construction has no Rust equivalent"}, - {"python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", "python_test": "test_should_return_normalized_response_when_no_native_payload", "status": "unmapped", "reason": "proxy-layer response construction has no Rust equivalent"} + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_should_encode_azure_document_intelligence_model_id", + "status": "unmapped", + "reason": "model-id URL percent-encoding has no Rust test; Rust only tests pages/features query building" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_should_reject_dot_segment_azure_document_intelligence_model_id", + "status": "unmapped", + "reason": "model-id dot-segment validation has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_preserves_azure_native_fields", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_response_normalizes_pages", + "justification": "both assert page markdown, dimension (inch-to-pixel) normalization, and usage_info.pages_processed from the same Azure succeeded response shape" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_async_transform_ocr_response_preserves_azure_native_fields", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_response_normalizes_pages", + "justification": "async twin of the sync case above, same underlying transform is exercised on the Rust side" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_tolerates_missing_native_fields", + "status": "unmapped", + "reason": "tables/keyValuePairs absence tolerance is not asserted by the Rust response test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_non_succeeded_status_raises", + "status": "unmapped", + "reason": "no Rust test asserts on a non-succeeded Azure DI status" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_supported_ocr_params_includes_features", + "status": "unmapped", + "reason": "supported-params list content has no Rust equivalent for Azure" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_native_format_carries_raw_operation", + "status": "unmapped", + "reason": "native req_format raw-operation passthrough is not tested in Rust" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_async_transform_ocr_response_native_format_carries_raw_operation", + "status": "unmapped", + "reason": "native req_format raw-operation passthrough is not tested in Rust" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_transform_ocr_response_default_format_omits_raw_operation", + "status": "unmapped", + "reason": "req_format gating of raw-operation output has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_passes_through_req_format", + "status": "unmapped", + "reason": "req_format passthrough in map_ocr_params has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_rejects_unknown_req_format_as_bad_request", + "status": "unmapped", + "reason": "req_format validation error path has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_complete_url_omits_req_format_query_param", + "status": "unmapped", + "reason": "no Rust test asserts req_format is excluded from the built URL" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_features", + "status": "unmapped", + "reason": "features-string normalization in map_ocr_params has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_empty_features_list_omitted", + "status": "unmapped", + "reason": "empty-features omission has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_map_ocr_params_invalid_features_raises", + "status": "unmapped", + "reason": "features validation error path has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_complete_url_appends_features_query", + "status": "unmapped", + "reason": "features query-param construction has no Rust test" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_get_complete_url_combines_pages_and_features", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_url_normalizes_zero_based_pages", + "justification": "both assert 0-based, duplicate page indices are deduped, sorted, and rewritten 1-based into the request URL" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_validate_environment_uses_subscription_key", + "status": "unmapped", + "reason": "Python-side header derivation from litellm_params; Rust's poll test only checks the header is present, not how it was resolved" + }, + { + "python_file": "tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py", + "python_test": "test_validate_environment_falls_back_to_entra_token", + "status": "unmapped", + "reason": "Entra bearer-token fallback logic has no Rust test" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_doc_intelligence_route", + "status": "unmapped", + "reason": "model-route string matching is Python-only dispatch logic" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestIsAzureDocumentIntelligenceModel::test_matches_documentintelligence_and_is_case_insensitive", + "status": "unmapped", + "reason": "model-route string matching is Python-only dispatch logic" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestIsAzureDocumentIntelligenceModel::test_does_not_match_mistral_route", + "status": "unmapped", + "reason": "model-route string matching is Python-only dispatch logic" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_does_not_hijack_doc_intelligence", + "status": "unmapped", + "reason": "api_base resolution from the secret manager runs before the Rust bridge is called, no Rust test exists for it" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestDocIntelligenceApiBaseResolution::test_explicit_api_base_is_honoured_for_doc_intelligence", + "status": "unmapped", + "reason": "api_base precedence resolution is Python-only" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py", + "python_test": "TestDocIntelligenceApiBaseResolution::test_generic_azure_ai_base_still_applies_to_mistral_ocr", + "status": "unmapped", + "reason": "api_base precedence resolution is Python-only" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_use_litellm_rust_toggles_flag", + "status": "unmapped", + "reason": "bridge-plumbing: Python-side feature-flag toggle, no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_env_var_enables_rust_ocr", + "status": "unmapped", + "reason": "bridge-plumbing: Python-side env-var flag gating" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_ocr_returns_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: dependency-injection test hook, not provider behavior" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_native_bridge_loader_returns_none_when_extension_absent", + "status": "unmapped", + "reason": "bridge-plumbing: native-extension import/loader fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_native_bridge_loader_caches_absent_extension", + "status": "unmapped", + "reason": "bridge-plumbing: loader caching behavior" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_native_bridge_available_reflects_loader", + "status": "unmapped", + "reason": "bridge-plumbing: loader availability check" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_aocr_returns_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: dependency-injection test hook" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_toggle_without_ocr_arg_preserves_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: injected-impl state retention regression" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_explicit_ocr_none_clears_injected_impl", + "status": "unmapped", + "reason": "bridge-plumbing: injected-impl clearing behavior" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_ocr_none_when_extension_absent", + "status": "unmapped", + "reason": "bridge-plumbing: degrade path when the native extension is missing" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_load_rust_ocr_uses_compiled_extension", + "status": "unmapped", + "reason": "bridge-plumbing: native module resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_timeout_to_seconds_handles_float_timeout_and_none", + "status": "unmapped", + "reason": "bridge-plumbing: Python-side timeout normalization helper" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_bridge_wrapper_forwards_prepared_args_and_wraps_response", + "status": "unmapped", + "reason": "bridge-plumbing: wrapper argument forwarding, asserted against a fake bridge not the real Rust code" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response", + "status": "unmapped", + "reason": "bridge-plumbing: async wrapper argument forwarding" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_prepares_request_and_wraps_response", + "status": "unmapped", + "reason": "bridge-plumbing: request preparation and response wrapping in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_resolves_key_via_secret_manager_when_missing", + "status": "unmapped", + "reason": "secret-manager: API key resolution happens in Python before the bridge is invoked" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_prefers_explicit_key_over_resolver", + "status": "unmapped", + "reason": "secret-manager: key precedence resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_uses_provider_api_key_env_var", + "status": "unmapped", + "reason": "secret-manager: provider-specific env var name resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_forwards_vertex_routing_metadata", + "status": "unmapped", + "reason": "secret-manager: vertex routing metadata merge happens in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager", + "status": "unmapped", + "reason": "secret-manager: vertex project/location resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager", + "status": "unmapped", + "reason": "secret-manager: azure_ai api_base resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint", + "status": "unmapped", + "reason": "secret-manager: doc-intelligence endpoint resolution" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_run_rust_ocr_runs_pre_call_logging", + "status": "unmapped", + "reason": "bridge-plumbing: Python logging-object pre_call invocation" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_routes_to_rust_when_enabled", + "status": "unmapped", + "reason": "bridge-plumbing: routing to a fake bridge, not the real Rust transform" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_routes_azure_ai_to_rust_when_enabled", + "status": "unmapped", + "reason": "bridge-plumbing: provider-prefix stripping before routing" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_rust_path_converts_file_document_before_bridge", + "status": "unmapped", + "reason": "file-normalization: raw-bytes-to-data-URI conversion happens in Python before the bridge call" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_exception_type_uses_resolved_provider_context", + "status": "unmapped", + "reason": "bridge-plumbing: Python exception-type mapping on bridge failure" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_aocr_routes_to_async_rust_when_enabled", + "status": "unmapped", + "reason": "bridge-plumbing: async routing to a fake bridge" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_aocr_exception_type_uses_resolved_provider_context", + "status": "unmapped", + "reason": "bridge-plumbing: async exception-type mapping on bridge failure" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_does_not_route_to_rust_when_disabled", + "status": "unmapped", + "reason": "bridge-plumbing: Python control flow for the toggle-disabled branch, no Rust-owned behavior runs" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_falls_back_to_python_when_bridge_unavailable", + "status": "unmapped", + "reason": "bridge-plumbing: Python-only fallback when the compiled Rust extension is absent, Rust cannot test its own absence" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_forwards_timeout_to_rust", + "status": "unmapped", + "reason": "bridge-plumbing: asserts the Python call site forwards a timeout kwarg, Rust receives an already-constructed request" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_passes_default_request_timeout_to_rust", + "status": "unmapped", + "reason": "bridge-plumbing: asserts the Python call site supplies a default timeout kwarg, no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_ocr_provider_configs_expose_api_key_env_vars", + "status": "unmapped", + "reason": "asserts per-provider get_api_key_env_var() strings; the closest Rust test (ocr_dispatch_supports_migrated_providers) asserts provider dispatch/param resolution instead, not API key env var names" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_pdf_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection is Python-only preprocessing before the bridge call" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_png_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_jpg_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_jpeg_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_gif_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_webp_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_tiff_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_tif_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_detect_bmp_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_be_case_insensitive", + "status": "unmapped", + "reason": "file-normalization: MIME detection case handling" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestGetMimeType::test_should_fallback_for_unknown_extension", + "status": "unmapped", + "reason": "file-normalization: MIME detection fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pdf_pathlib_path_to_document_url", + "status": "unmapped", + "reason": "file-normalization: local-path-to-data-URI conversion happens in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_image_pathlib_path_to_image_url", + "status": "unmapped", + "reason": "file-normalization: local-path-to-data-URI conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_reject_bare_str_path", + "status": "unmapped", + "reason": "file-normalization: arbitrary-file-read guard on bare str paths" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_pathlib_path", + "status": "unmapped", + "reason": "file-normalization: local-path-to-data-URI conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes", + "status": "unmapped", + "reason": "file-normalization: raw-bytes-to-data-URI conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_explicit_mime_type", + "status": "unmapped", + "reason": "file-normalization: explicit MIME override on raw bytes" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_raw_bytes_with_image_mime_type", + "status": "unmapped", + "reason": "file-normalization: explicit MIME override on raw bytes" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object", + "status": "unmapped", + "reason": "file-normalization: file-like-object conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_convert_file_like_object_with_name", + "status": "unmapped", + "reason": "file-normalization: file-like-object name-based MIME detection" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_missing_file_field", + "status": "unmapped", + "reason": "file-normalization: missing-field validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_nonexistent_pathlib_path", + "status": "unmapped", + "reason": "file-normalization: missing-file validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_empty_file", + "status": "unmapped", + "reason": "file-normalization: empty-file validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_unsupported_type", + "status": "unmapped", + "reason": "file-normalization: unsupported input type validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_raise_error_for_invalid_mime_type", + "status": "unmapped", + "reason": "file-normalization: MIME-type injection validation" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestConvertFileDocumentToUrlDocument::test_should_override_mime_type_for_pathlib_path", + "status": "unmapped", + "reason": "file-normalization: explicit MIME override precedence" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_build_document_url_for_pdf", + "status": "unmapped", + "reason": "file-normalization: multipart upload conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_png", + "status": "unmapped", + "reason": "file-normalization: multipart upload conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_build_image_url_for_jpeg", + "status": "unmapped", + "reason": "file-normalization: multipart upload conversion" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_octet_stream", + "status": "unmapped", + "reason": "file-normalization: filename-based MIME fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_detect_mime_from_filename_when_content_type_is_none", + "status": "unmapped", + "reason": "file-normalization: filename-based MIME fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_fallback_to_octet_stream_for_unknown", + "status": "unmapped", + "reason": "file-normalization: default MIME fallback" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_preserve_base64_content_correctly", + "status": "unmapped", + "reason": "file-normalization: binary round-trip through base64" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_from_content_type", + "status": "unmapped", + "reason": "file-normalization: content-type parameter stripping" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestBuildDocumentFromUpload::test_should_strip_mime_parameters_with_multiple_params", + "status": "unmapped", + "reason": "file-normalization: content-type parameter stripping" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_reject_file_type_document_in_json_body", + "status": "unmapped", + "reason": "proxy-layer JSON-body file-type guard, a different mechanism than Rust's URL-fetch SSRF guard" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_accept_document_url_type_in_json_body", + "status": "unmapped", + "reason": "proxy-layer JSON-body parsing" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_raise_on_invalid_json_body", + "status": "unmapped", + "reason": "proxy-layer JSON-body parsing error path" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_file_input.py", + "python_test": "TestProxySecurityGuard::test_should_ignore_document_form_field_injection", + "status": "unmapped", + "reason": "proxy-layer multipart form-field injection guard, a different mechanism than Rust's URL-fetch SSRF guard" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestGetSupportedOcrParams::test_extract_header_in_supported_params", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "extract_header_is_a_supported_ocr_param", + "justification": "both assert extract_header appears in the Mistral supported OCR params list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestGetSupportedOcrParams::test_extract_footer_in_supported_params", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "extract_footer_is_a_supported_ocr_param", + "justification": "both assert extract_footer appears in the Mistral supported OCR params list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestGetSupportedOcrParams::test_existing_params_still_present", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "existing_ocr_params_remain_supported", + "justification": "both assert the previously supported params are still present in the supported list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_extract_header_passed_through", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_extract_header", + "justification": "both assert extract_header alone survives map_ocr_params unchanged" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_extract_footer_passed_through", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_extract_footer", + "justification": "both assert extract_footer alone survives map_ocr_params unchanged" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_extract_header_and_footer_together", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_extract_header_and_footer", + "justification": "both assert header and footer passed together are both forwarded with their given values" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestMapOcrParams::test_unknown_param_is_dropped", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_drops_unknown_params", + "justification": "both assert an unrecognized param key is dropped while a known one is kept" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestNewSupportedParams::test_new_param_in_supported_list", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "new_ocr_params_are_supported", + "justification": "both assert each OCR4 param is in the supported list" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestNewParamsMapOcr::test_new_param_passed_through", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "map_ocr_params_forwards_new_ocr_params", + "justification": "both assert each OCR4 param/value pair survives map_ocr_params unchanged" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrRequest::test_param_included_in_request_body", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_request_includes_each_optional_param", + "justification": "both assert each optional param value lands in the built request body alongside model/document with no files" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrRequest::test_multiple_new_params_together", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_request_includes_multiple_new_params", + "justification": "both assert multiple OCR4 params passed together all land in the same request body" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrResponseOcr4Fields::test_blocks_and_confidence_scores_preserved", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_response_preserves_blocks_and_confidence_scores", + "justification": "both assert blocks and confidence_scores survive the OCR response transform on the returned page" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_transformation.py", + "python_test": "TestTransformOcrResponseOcr4Fields::test_ocr4_fields_survive_model_dump", + "status": "mapped", + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_response_preserves_ocr4_page_fields", + "justification": "both assert tables, hyperlinks, header and footer survive the OCR response transform on the returned page" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_model_info_ocr4_price", + "status": "unmapped", + "reason": "cost-calc: pricing/model-info lookup is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr4_cost_scales_with_pages", + "status": "unmapped", + "reason": "cost-calc: per-page pricing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_pricing_entry", + "status": "unmapped", + "reason": "cost-calc: cost-map JSON entry validation is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_model_info_price", + "status": "unmapped", + "reason": "cost-calc: pricing/model-info lookup is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_cost_scales_with_pages", + "status": "unmapped", + "reason": "cost-calc: per-page pricing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_bills_ocr_and_annotation_pages_at_their_own_rates", + "status": "unmapped", + "reason": "cost-calc: mixed-rate billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_bills_annotation_only_response", + "status": "unmapped", + "reason": "cost-calc: annotation-only billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_ocr3_bills_annotation_pages_when_pages_processed_missing", + "status": "unmapped", + "reason": "cost-calc: fallback billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py", + "python_test": "test_azure_doc_ai_annotation_pages_fall_back_to_ocr_rate", + "status": "unmapped", + "reason": "cost-calc: fallback billing math is Python-only" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_rust_ocr_serves_default_format", + "status": "unmapped", + "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_rust_ocr_skipped_for_native_format", + "status": "unmapped", + "reason": "request-format gating decision is made in Python before the Rust bridge is ever invoked" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_native_format_rejected_for_provider_without_support_as_bad_request", + "status": "unmapped", + "reason": "provider-support validation for req_format happens in Python" + }, + { + "python_file": "tests/test_litellm/ocr/test_ocr_native_format.py", + "python_test": "test_unknown_format_rejected_for_provider_without_support_as_bad_request", + "status": "unmapped", + "reason": "req_format validation error path is Python-only" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestHandlerDiscovery::test_handler_discovered_for_ocr", + "status": "unmapped", + "reason": "guardrail-translation handler discovery is a Python proxy-layer concern" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestHandlerDiscovery::test_handler_discovered_for_aocr", + "status": "unmapped", + "reason": "guardrail-translation handler discovery is a Python proxy-layer concern" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_document_url", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_image_url", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_no_document", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_process_invalid_document", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestInputProcessing::test_input_blocking_guardrail", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_single_page", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_multiple_pages", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_empty_pages", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_page_with_empty_markdown", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_process_preserves_page_metadata", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestOutputProcessing::test_output_blocking_guardrail", + "status": "unmapped", + "reason": "scoped to the Python translation handler, not the Rust gateway's guardrail hook lifecycle" + }, + { + "python_file": "tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py", + "python_test": "TestPIIMaskingScenario::test_pii_masking_in_ocr_pages", + "status": "unmapped", + "reason": "PII redaction in the translation handler has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_read_req_format_from_header", + "status": "unmapped", + "reason": "proxy-layer header parsing has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_prefer_body_req_format_over_header", + "status": "unmapped", + "reason": "proxy-layer body-vs-header precedence has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_omit_req_format_when_header_absent", + "status": "unmapped", + "reason": "proxy-layer parsing has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_reject_unknown_req_format", + "status": "unmapped", + "reason": "proxy-layer validation has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_return_native_payload_with_litellm_response_headers", + "status": "unmapped", + "reason": "proxy-layer response construction has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/proxy/ocr_endpoints/test_endpoints.py", + "python_test": "test_should_return_normalized_response_when_no_native_payload", + "status": "unmapped", + "reason": "proxy-layer response construction has no Rust equivalent" + }, + { + "python_file": "tests/test_litellm/ocr/test_rust_bridge.py", + "python_test": "test_explicit_false_overrides_process_enable", + "status": "unmapped", + "reason": "bridge-plumbing: asserts the explicit-per-request rust:False override wins over the Python process flag, no Rust-owned behavior runs" + } ], "rust_only_tests": [ - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "rust_custom_logger_reads_failure_payload_for_non_ocr_call_type", "reason": "exercises the non-OCR (acompletion) call-type branch of the logger; the OCR branch is covered separately by rust_custom_logger_reads_success_payload_for_ocr"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "no_callback_fast_path_dispatches_nothing", "reason": "Rust-only fast-path optimization test for when zero callbacks are registered; Python has no equivalent no-op dispatch path"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "with_standard_logging_payload_keeps_top_level_fields_in_sync", "reason": "Rust-internal builder-method invariant, Python has no equivalent internal builder"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "blocks_private_and_metadata_ips", "reason": "SSRF IP-blocking helper has no Python unit test; Python relies on the proxy-layer JSON/form guards instead"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "convert_document_url_rejects_loopback_fetch", "reason": "URL-fetch SSRF protection is Rust-gateway-only"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", "rust_test": "convert_document_url_leaves_data_uri_untouched", "reason": "URL-fetch SSRF protection is Rust-gateway-only"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_passes_short_strings_through", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_caps_long_payloads", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "truncate_error_body_does_not_split_multibyte_chars", "reason": "Rust-gateway error-body truncation helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_dispatch_supports_migrated_providers", "reason": "Rust-internal provider-config dispatch table has no equivalent Python unit test"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "string_headers_accepts_string_values", "reason": "Rust-gateway header-coercion helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "auth_header_detection_is_case_insensitive", "reason": "Rust-gateway header-detection helper has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_runs_pre_during_and_success_hooks", "reason": "full gateway-level guardrail-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_runs_failure_hook_on_provider_error", "reason": "full gateway-level failure-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_lifecycle_pre_call_block_skips_provider_socket", "reason": "full gateway-level pre-call-block-plus-socket-skip test with no Python equivalent at this integration scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "ocr_does_not_duplicate_authorization_header_when_header_is_supplied", "reason": "outgoing HTTP header dedup at the Rust gateway has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "document_intelligence_poll_uses_resolved_subscription_key", "reason": "full Azure DI poll-loop integration test with no Python equivalent at this scope"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", "rust_test": "string_headers_rejects_non_string_values", "reason": "Rust-gateway header-coercion error path has no Python counterpart"}, - {"rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "azure_ai_reuses_mistral_body_transform", "reason": "Rust-internal delegation-to-Mistral-transform implementation detail, no Python test asserts this delegation"}, - {"rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", "rust_test": "document_intelligence_request_uses_base64_source_for_data_uri", "reason": "no Python test asserts on the base64Source request body shape"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_mistral_url_uses_project_location_and_model", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_mistral_reuses_mistral_body_transform", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_deepseek_request_uses_ocr_endpoint_shape", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", "rust_test": "vertex_deepseek_response_wraps_markdown_content", "reason": "vertex OCR support has no Python unit test coverage yet"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_request_rejects_non_object_document", "reason": "non-object document rejection has no dedicated Python unit test"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "transform_ocr_response_normalizes_mistral_json", "reason": "Python's response tests target OCR4-specific fields only, none asserts the same base normalization this Rust test checks"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "complete_url_defaults_and_dedupes_v1", "reason": "URL-building/defaulting for Mistral has no Python unit test"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "resolve_api_key_prefers_param_then_env", "reason": "API key resolution precedence at the Rust provider-config layer has no Python unit test"}, - {"rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", "rust_test": "resolve_api_key_errors_when_absent", "reason": "API key resolution error path at the Rust provider-config layer has no Python unit test"}, - {"rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", "rust_test": "rust_custom_logger_reads_success_payload_for_ocr", "reason": "Rust-internal custom-logger dispatch for OCR payloads has no Python unit test at this layer"} + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "rust_custom_logger_reads_failure_payload_for_non_ocr_call_type", + "reason": "exercises the non-OCR (acompletion) call-type branch of the logger; the OCR branch is covered separately by rust_custom_logger_reads_success_payload_for_ocr" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "no_callback_fast_path_dispatches_nothing", + "reason": "Rust-only fast-path optimization test for when zero callbacks are registered; Python has no equivalent no-op dispatch path" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "with_standard_logging_payload_keeps_top_level_fields_in_sync", + "reason": "Rust-internal builder-method invariant, Python has no equivalent internal builder" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", + "rust_test": "blocks_private_and_metadata_ips", + "reason": "SSRF IP-blocking helper has no Python unit test; Python relies on the proxy-layer JSON/form guards instead" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", + "rust_test": "convert_document_url_rejects_loopback_fetch", + "reason": "URL-fetch SSRF protection is Rust-gateway-only" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs", + "rust_test": "convert_document_url_leaves_data_uri_untouched", + "reason": "URL-fetch SSRF protection is Rust-gateway-only" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "truncate_error_body_passes_short_strings_through", + "reason": "Rust-gateway error-body truncation helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "truncate_error_body_caps_long_payloads", + "reason": "Rust-gateway error-body truncation helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "truncate_error_body_does_not_split_multibyte_chars", + "reason": "Rust-gateway error-body truncation helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_dispatch_supports_migrated_providers", + "reason": "Rust-internal provider-config dispatch table has no equivalent Python unit test" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "string_headers_accepts_string_values", + "reason": "Rust-gateway header-coercion helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "auth_header_detection_is_case_insensitive", + "reason": "Rust-gateway header-detection helper has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_lifecycle_runs_pre_during_and_success_hooks", + "reason": "full gateway-level guardrail-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_lifecycle_runs_failure_hook_on_provider_error", + "reason": "full gateway-level failure-hook-plus-HTTP-lifecycle test with no Python equivalent at this integration scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_lifecycle_pre_call_block_skips_provider_socket", + "reason": "full gateway-level pre-call-block-plus-socket-skip test with no Python equivalent at this integration scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "ocr_does_not_duplicate_authorization_header_when_header_is_supplied", + "reason": "outgoing HTTP header dedup at the Rust gateway has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "document_intelligence_poll_uses_resolved_subscription_key", + "reason": "full Azure DI poll-loop integration test with no Python equivalent at this scope" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/ocr/tests.rs", + "rust_test": "string_headers_rejects_non_string_values", + "reason": "Rust-gateway header-coercion error path has no Python counterpart" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "azure_ai_reuses_mistral_body_transform", + "reason": "Rust-internal delegation-to-Mistral-transform implementation detail, no Python test asserts this delegation" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs", + "rust_test": "document_intelligence_request_uses_base64_source_for_data_uri", + "reason": "no Python test asserts on the base64Source request body shape" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_mistral_url_uses_project_location_and_model", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_mistral_reuses_mistral_body_transform", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_deepseek_request_uses_ocr_endpoint_shape", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs", + "rust_test": "vertex_deepseek_response_wraps_markdown_content", + "reason": "vertex OCR support has no Python unit test coverage yet" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_request_rejects_non_object_document", + "reason": "non-object document rejection has no dedicated Python unit test" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "transform_ocr_response_normalizes_mistral_json", + "reason": "Python's response tests target OCR4-specific fields only, none asserts the same base normalization this Rust test checks" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "complete_url_defaults_and_dedupes_v1", + "reason": "URL-building/defaulting for Mistral has no Python unit test" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "resolve_api_key_prefers_param_then_env", + "reason": "API key resolution precedence at the Rust provider-config layer has no Python unit test" + }, + { + "rust_file": "litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs", + "rust_test": "resolve_api_key_errors_when_absent", + "reason": "API key resolution error path at the Rust provider-config layer has no Python unit test" + }, + { + "rust_file": "litellm-rust/crates/ai-gateway/src/integrations/custom_logger/mod.rs", + "rust_test": "rust_custom_logger_reads_success_payload_for_ocr", + "reason": "Rust-internal custom-logger dispatch for OCR payloads has no Python unit test at this layer" + } ] }