From 12919628501340c8b7b596d33492bd9f5ef6eff0 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 16 Jul 2026 13:23:41 -0700 Subject: [PATCH 01/18] feat(ui): configure Anthropic automatic prompt caching from the Admin UI Register enable_anthropic_prompt_caching and anthropic_prompt_caching_ttl on the General Settings table so caching can be turned on without hand-writing config. The registry could not express either field: validation was hardcoded to a float in (0, 1], reset set every field to None (not a bool for a boolean flag), and the listing reported any non-None value as 'In Config', which a False default would always trip. Validation now dispatches on the declared type and reset restores each field's own default. ConfigList carries field_options so the table can render a Select for enums instead of no editor at all. --- litellm/proxy/_types.py | 3 +- litellm/proxy/proxy_server.py | 96 +++++++-- tests/test_litellm/proxy/test_proxy_server.py | 204 ++++++++++++++++++ .../_components/general_settings.tsx | 15 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 5 files changed, 300 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d102c1d1e37..e07c7b9ae78 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1011,10 +1011,10 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): mcp_tool_search_enabled: Optional[bool] = None +from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402 from litellm.types.object_permission import ( # noqa: E402 ObjectPermissionDict as ObjectPermissionDict, ) -from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402 class GenerateRequestBase(LiteLLMPydanticObjectBase): @@ -2122,6 +2122,7 @@ class ConfigList(LiteLLMPydanticObjectBase): field_default_value: Any premium_field: bool = False nested_fields: Optional[List[FieldDetail]] = None # For nested dictionary or Pydantic fields + field_options: Optional[list[str]] = None # Allowed values, for field_type == "Select" class UserHeaderMapping(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dbdfdd5fdd3..ae91eb70427 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -28,6 +28,7 @@ from typing import ( Optional, Set, Tuple, + TypedDict, Union, cast, get_args, @@ -39,6 +40,7 @@ import anyio import websockets import websockets.exceptions from pydantic import BaseModel, Json, JsonValue +from typing_extensions import NotRequired, assert_never from litellm._uuid import uuid from litellm.constants import ( @@ -363,15 +365,15 @@ from litellm.proxy.management_endpoints.cache_settings_endpoints import ( from litellm.proxy.management_endpoints.callback_management_endpoints import ( router as callback_management_endpoints_router, ) -from litellm.proxy.management_endpoints.coordination_redis_endpoints import ( - get_persisted_coordination_redis_settings, - router as coordination_redis_settings_router, -) from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_privileges, _user_has_admin_view, admin_can_invite_user, ) +from litellm.proxy.management_endpoints.coordination_redis_endpoints import ( + get_persisted_coordination_redis_settings, + router as coordination_redis_settings_router, +) from litellm.proxy.management_endpoints.cost_tracking_settings import ( router as cost_tracking_settings_router, ) @@ -14800,7 +14802,16 @@ async def get_config_general_settings( ) -_GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, dict[str, str]] = { +GeneralSettingsUILiteLLMValue = Union[float, bool, str, None] + + +class GeneralSettingsUILiteLLMFieldSpec(TypedDict): + type: Literal["Float", "Boolean", "Select"] + description: str + options: NotRequired[tuple[str, ...]] + + +_GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, GeneralSettingsUILiteLLMFieldSpec] = { "budget_exceeded_throttle_percentage": { "type": "Float", "description": ( @@ -14809,18 +14820,64 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, dict[str, str]] = { "over-budget keys." ), }, + "enable_anthropic_prompt_caching": { + "type": "Boolean", + "description": ( + "Automatically add Anthropic cache_control breakpoints to the system prompt and the " + "trailing turn, for Anthropic and Bedrock Claude models that support prompt caching. " + "Lets clients that never set cache_control themselves still get cached prompts. " + "Requests that already carry their own cache_control are left untouched." + ), + }, + "anthropic_prompt_caching_ttl": { + "type": "Select", + "options": ("5m", "1h"), + "description": ( + "Cache lifetime for the breakpoints added by 'enable_anthropic_prompt_caching'. " + "Leave empty for Anthropic's 5m default. 1h suits long agentic sessions but doubles " + "the cache write premium." + ), + }, } -def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> Optional[float]: +def _general_settings_ui_litellm_default( + field_type: Literal["Float", "Boolean", "Select"], +) -> GeneralSettingsUILiteLLMValue: + """The value a field falls back to when it is cleared or reset.""" + return False if field_type == "Boolean" else None + + +def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> GeneralSettingsUILiteLLMValue: + spec = _GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name] + field_type = spec["type"] if value is None or value == "": - return None - if isinstance(value, bool) or not isinstance(value, (int, float)) or not (0 < float(value) <= 1): - raise HTTPException( - status_code=400, - detail={"error": f"{field_name} must be a number in (0, 1] or empty"}, - ) - return float(value) + return _general_settings_ui_litellm_default(field_type) + match field_type: + case "Boolean": + if not isinstance(value, bool): + raise HTTPException( + status_code=400, + detail={"error": f"{field_name} must be true or false"}, + ) + return value + case "Select": + options = spec.get("options", ()) + if value not in options: + raise HTTPException( + status_code=400, + detail={"error": f"{field_name} must be one of: {', '.join(options)}, or empty"}, + ) + return cast(str, value) # cast-ok: membership in options proves it is one of the option strings + case "Float": + if isinstance(value, bool) or not isinstance(value, (int, float)) or not (0 < float(value) <= 1): + raise HTTPException( + status_code=400, + detail={"error": f"{field_name} must be a number in (0, 1] or empty"}, + ) + return float(value) + case _: + assert_never(field_type) async def _persist_general_settings_ui_litellm_field( @@ -14841,11 +14898,12 @@ async def _persist_general_settings_ui_litellm_field( async def _reset_general_settings_ui_litellm_field(field_name: str, user_api_key_dict: UserAPIKeyAuth) -> dict: config = await proxy_config.get_config() before_value = config.get("litellm_settings", {}).get(field_name) - setattr(litellm, field_name, None) + default_value = _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name]["type"]) + setattr(litellm, field_name, default_value) if "litellm_settings" in config: config["litellm_settings"].pop(field_name, None) await proxy_config.save_config(new_config=config) - asyncio.create_task(create_config_audit_log(field_name, "deleted", before_value, None, user_api_key_dict)) + asyncio.create_task(create_config_audit_log(field_name, "deleted", before_value, default_value, user_api_key_dict)) return {"message": f"Field {field_name} reset", "status": "success"} @@ -15013,11 +15071,12 @@ async def get_config_list( else {} ) for litellm_field_name, spec in _GENERAL_SETTINGS_UI_LITELLM_FIELDS.items(): - current_value: Optional[float] = getattr(litellm, litellm_field_name, None) + current_value: GeneralSettingsUILiteLLMValue = getattr(litellm, litellm_field_name, None) + default_value = _general_settings_ui_litellm_default(spec["type"]) stored_in_db_litellm: Optional[bool] if litellm_field_name in db_litellm_settings: stored_in_db_litellm = True - elif current_value is not None: + elif current_value != default_value: stored_in_db_litellm = False else: stored_in_db_litellm = None @@ -15028,7 +15087,8 @@ async def get_config_list( field_description=spec["description"], field_value=current_value, stored_in_db=stored_in_db_litellm, - field_default_value=None, + field_default_value=default_value, + field_options=list(spec.get("options", ())) or None, nested_fields=None, ) ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 54db0c0fd4f..86e807447c1 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8984,6 +8984,210 @@ async def test_update_config_field_throttle_persists_to_litellm_settings(monkeyp assert saved["litellm_settings"]["budget_exceeded_throttle_percentage"] == 0.1 +def test_get_config_list_includes_anthropic_prompt_caching_fields(monkeypatch): + """The auto prompt caching flag and its ttl are litellm_settings globals surfaced on the + General Settings table, so an admin can turn caching on without hand-writing config. The + ttl is a Select and must ship its allowed values, or the table renders no editor for it.""" + import types + from unittest.mock import AsyncMock, MagicMock + + from fastapi.testclient import TestClient + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.proxy_server import app + + mock_prisma = MagicMock() + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + monkeypatch.setattr(litellm, "anthropic_prompt_caching_ttl", "1h") + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + client = TestClient(app) + resp = client.get("/config/list", params={"config_type": "general_settings"}) + assert resp.status_code == 200, resp.text + fields = {item["field_name"]: item for item in resp.json()} + + assert fields["enable_anthropic_prompt_caching"]["field_type"] == "Boolean" + assert fields["enable_anthropic_prompt_caching"]["field_value"] is True + + assert fields["anthropic_prompt_caching_ttl"]["field_type"] == "Select" + assert fields["anthropic_prompt_caching_ttl"]["field_value"] == "1h" + assert fields["anthropic_prompt_caching_ttl"]["field_options"] == ["5m", "1h"] + finally: + app.dependency_overrides.clear() + + +def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypatch): + """The flag defaults to False rather than None, so a plain 'is not None' check would + report the default as 'In Config' and imply an admin had set it.""" + import types + from unittest.mock import AsyncMock, MagicMock + + from fastapi.testclient import TestClient + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.proxy_server import app + + mock_prisma = MagicMock() + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + client = TestClient(app) + resp = client.get("/config/list", params={"config_type": "general_settings"}) + fields = {item["field_name"]: item for item in resp.json()} + assert fields["enable_anthropic_prompt_caching"]["stored_in_db"] is None + finally: + app.dependency_overrides.clear() + + +@pytest.mark.parametrize( + "field_name, field_value", + [ + ("enable_anthropic_prompt_caching", True), + ("enable_anthropic_prompt_caching", False), + ("anthropic_prompt_caching_ttl", "5m"), + ("anthropic_prompt_caching_ttl", "1h"), + ], +) +@pytest.mark.asyncio +async def test_update_config_field_prompt_caching_persists_to_litellm_settings(monkeypatch, field_name, field_value): + """Toggling either row must set litellm. live and persist under litellm_settings, + so the running proxy caches immediately and still does after a restart.""" + from unittest.mock import MagicMock + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import ( + ConfigFieldUpdate, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.proxy_server import update_config_general_settings + + saved: dict = {} + + async def fake_get_config(): + return {"litellm_settings": {}} + + async def fake_save_config(new_config=None): + saved.update(new_config or {}) + + monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config) + monkeypatch.setattr(ps.proxy_config, "save_config", fake_save_config) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(litellm, "store_audit_logs", False) + monkeypatch.setattr(litellm, field_name, None) + + admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN) + await update_config_general_settings( + data=ConfigFieldUpdate(field_name=field_name, field_value=field_value, config_type="general_settings"), + user_api_key_dict=admin, + ) + + assert getattr(litellm, field_name) == field_value + assert saved["litellm_settings"][field_name] == field_value + + +@pytest.mark.parametrize( + "field_name, bad_value", + [ + ("enable_anthropic_prompt_caching", "yes"), + ("enable_anthropic_prompt_caching", 1), + ("anthropic_prompt_caching_ttl", "10m"), + ("anthropic_prompt_caching_ttl", "1H"), + ("anthropic_prompt_caching_ttl", 3600), + ], +) +@pytest.mark.asyncio +async def test_update_config_field_prompt_caching_rejects_invalid(monkeypatch, field_name, bad_value): + """An unsupported ttl must be refused here rather than reaching Anthropic verbatim.""" + from unittest.mock import MagicMock + + from fastapi import HTTPException + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import ( + ConfigFieldUpdate, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.proxy_server import update_config_general_settings + + async def fake_get_config(): + return {"litellm_settings": {}} + + monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(litellm, field_name, None) + + admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN) + with pytest.raises(HTTPException) as exc: + await update_config_general_settings( + data=ConfigFieldUpdate(field_name=field_name, field_value=bad_value, config_type="general_settings"), + user_api_key_dict=admin, + ) + assert exc.value.status_code == 400 + assert getattr(litellm, field_name) is None + + +@pytest.mark.parametrize( + "field_name, expected_default", + [ + ("enable_anthropic_prompt_caching", False), + ("anthropic_prompt_caching_ttl", None), + ("budget_exceeded_throttle_percentage", None), + ], +) +@pytest.mark.asyncio +async def test_reset_config_field_restores_type_default(monkeypatch, field_name, expected_default): + """Reset must restore each field's own default. Blanket None would leave the boolean flag + set to None, which is not a bool and would read as neither on nor off.""" + from unittest.mock import MagicMock + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import ( + ConfigFieldDelete, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.proxy_server import delete_config_general_settings + + saved: dict = {} + + async def fake_get_config(): + return {"litellm_settings": {field_name: "stale"}} + + async def fake_save_config(new_config=None): + saved.update(new_config or {}) + + monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config) + monkeypatch.setattr(ps.proxy_config, "save_config", fake_save_config) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(litellm, "store_audit_logs", False) + monkeypatch.setattr(litellm, field_name, "stale") + + admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN) + await delete_config_general_settings( + data=ConfigFieldDelete(field_name=field_name, config_type="general_settings"), + user_api_key_dict=admin, + ) + + assert getattr(litellm, field_name) is expected_default + assert field_name not in saved["litellm_settings"] + + @pytest.mark.parametrize("bad_value", [0, -0.1, 1.5, True]) @pytest.mark.asyncio async def test_update_config_field_throttle_rejects_invalid(monkeypatch, bad_value): diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index 3955e80f5e9..a1ef8252afa 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -14,7 +14,7 @@ import { } from "@tremor/react"; import { TabPanel, TabPanels, TabGroup, TabList, Tab } from "@tremor/react"; import { getGeneralSettingsCall, updateConfigFieldSetting, deleteConfigFieldSetting } from "@/components/networking"; -import { InputNumber } from "antd"; +import { InputNumber, Select as AntdSelect } from "antd"; import { TrashIcon } from "@heroicons/react/outline"; import { StatusBadge } from "@/components/shared/table_cells"; @@ -33,6 +33,7 @@ interface generalSettingsItem { field_value: any; field_description: string; stored_in_db: boolean | null; + field_options?: string[] | null; } const GeneralSettings: React.FC = ({ accessToken, userRole, userID }) => { @@ -169,6 +170,18 @@ const GeneralSettings: React.FC = ({ accessToken, user value={value.field_value} onChange={(newValue) => handleInputChange(value.field_name, newValue)} /> + ) : value.field_type == "Select" ? ( + ({ + label: option, + value: option, + }))} + onChange={(newValue) => handleInputChange(value.field_name, newValue ?? "")} + /> ) : null} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0d8f55164f9..45b06c9e44f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22711,6 +22711,8 @@ export interface components { field_description: string; /** Field Name */ field_name: string; + /** Field Options */ + field_options?: string[] | null; /** Field Type */ field_type: string; /** Field Value */ From 9f7f53a82a938b0471e41c89be967d49b3275434 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 16 Jul 2026 13:28:34 -0700 Subject: [PATCH 02/18] refactor(ui): extract the General Settings value editor into a component The value cell was a ternary chain over field_type; adding Select made it a fourth level and tripped no-nested-ternary. Early returns read better than a deeper chain and let the suppression baseline ratchet down. --- ui/litellm-dashboard/eslint-suppressions.json | 2 +- .../_components/general_settings.tsx | 80 +++++++++++-------- 2 files changed, 49 insertions(+), 33 deletions(-) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 32e9a03da95..90b0c84244e 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1079,7 +1079,7 @@ }, "src/app/(dashboard)/router-settings/_components/general_settings.tsx": { "no-nested-ternary": { - "count": 3 + "count": 1 }, "no-restricted-imports": { "count": 2 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index a1ef8252afa..af6bbdde8b9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -36,6 +36,53 @@ interface generalSettingsItem { field_options?: string[] | null; } +const SettingValueEditor: React.FC<{ + setting: generalSettingsItem; + onChange: (fieldName: string, newValue: any) => void; +}> = ({ setting, onChange }) => { + if (setting.field_type === "Integer") { + return ( + onChange(setting.field_name, newValue)} + /> + ); + } + if (setting.field_type === "Boolean") { + return ( + onChange(setting.field_name, checked)} + /> + ); + } + if (setting.field_type === "Float") { + return ( + onChange(setting.field_name, newValue)} + /> + ); + } + if (setting.field_type === "Select") { + return ( + ({ label: option, value: option }))} + onChange={(newValue) => onChange(setting.field_name, newValue ?? "")} + /> + ); + } + return null; +}; + const GeneralSettings: React.FC = ({ accessToken, userRole, userID }) => { const [generalSettings, setGeneralSettings] = useState([]); @@ -151,38 +198,7 @@ const GeneralSettings: React.FC = ({ accessToken, user

- {value.field_type == "Integer" ? ( - handleInputChange(value.field_name, newValue)} - /> - ) : value.field_type == "Boolean" ? ( - handleInputChange(value.field_name, checked)} - /> - ) : value.field_type == "Float" ? ( - handleInputChange(value.field_name, newValue)} - /> - ) : value.field_type == "Select" ? ( - ({ - label: option, - value: option, - }))} - onChange={(newValue) => handleInputChange(value.field_name, newValue ?? "")} - /> - ) : null} + {value.stored_in_db == true ? ( From 16e39542a095c0f4aaf1a7835d8598b664dd6716 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 16 Jul 2026 15:54:53 -0700 Subject: [PATCH 03/18] docs(ui): state that Anthropic prompt caches are shared per upstream credential The provider caches a prefix against the credentials that sent it, not per end user, so turning the flag on makes every caller's prompts cacheable on that shared account. Surface that where the toggle is, since it is the operator's call to make. --- litellm/proxy/proxy_server.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ae91eb70427..ad1617017f9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14826,7 +14826,11 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, GeneralSettingsUILiteLLMFieldSpec "Automatically add Anthropic cache_control breakpoints to the system prompt and the " "trailing turn, for Anthropic and Bedrock Claude models that support prompt caching. " "Lets clients that never set cache_control themselves still get cached prompts. " - "Requests that already carry their own cache_control are left untouched." + "Requests that already carry their own cache_control are left untouched. " + "The provider caches a prefix against the upstream credentials that sent it, not per " + "end user, so this makes every caller's prompts cacheable on that shared account. " + "Leave this off if callers sharing a set of credentials must not learn whether " + "another caller recently sent a given prompt." ), }, "anthropic_prompt_caching_ttl": { From 73cbbdd51defe21a6db5beadf2b3ee73454677be Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 17 Jul 2026 11:38:24 -0700 Subject: [PATCH 04/18] feat(ui): move Anthropic prompt caching to its own Router Settings tab Rather than mixing the flag and its ttl into the generic General settings table (which also surfaced the confusing Not Set / In Config / In DB provenance badges), give prompt caching a dedicated tab with a purpose-built toggle and ttl dropdown. Each registry field gains an optional tab, surfaced as ConfigList.field_tab, so the General tab renders the ungrouped fields and the caching fields render on their own tab. The update, persist and reset endpoints are unchanged. --- litellm/proxy/_types.py | 1 + litellm/proxy/proxy_server.py | 4 + tests/test_litellm/proxy/test_proxy_server.py | 6 ++ .../_components/general_settings.tsx | 78 ++++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 5 files changed, 90 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e07c7b9ae78..b47b43411c5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2123,6 +2123,7 @@ class ConfigList(LiteLLMPydanticObjectBase): premium_field: bool = False nested_fields: Optional[List[FieldDetail]] = None # For nested dictionary or Pydantic fields field_options: Optional[list[str]] = None # Allowed values, for field_type == "Select" + field_tab: Optional[str] = None # Admin UI sub-tab this field renders under; None groups it with the rest class UserHeaderMapping(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ad1617017f9..6725ecdb584 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14809,6 +14809,7 @@ class GeneralSettingsUILiteLLMFieldSpec(TypedDict): type: Literal["Float", "Boolean", "Select"] description: str options: NotRequired[tuple[str, ...]] + tab: NotRequired[str] # Admin UI sub-tab this field renders under; None groups it with the rest _GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, GeneralSettingsUILiteLLMFieldSpec] = { @@ -14822,6 +14823,7 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, GeneralSettingsUILiteLLMFieldSpec }, "enable_anthropic_prompt_caching": { "type": "Boolean", + "tab": "prompt_caching", "description": ( "Automatically add Anthropic cache_control breakpoints to the system prompt and the " "trailing turn, for Anthropic and Bedrock Claude models that support prompt caching. " @@ -14836,6 +14838,7 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, GeneralSettingsUILiteLLMFieldSpec "anthropic_prompt_caching_ttl": { "type": "Select", "options": ("5m", "1h"), + "tab": "prompt_caching", "description": ( "Cache lifetime for the breakpoints added by 'enable_anthropic_prompt_caching'. " "Leave empty for Anthropic's 5m default. 1h suits long agentic sessions but doubles " @@ -15093,6 +15096,7 @@ async def get_config_list( stored_in_db=stored_in_db_litellm, field_default_value=default_value, field_options=list(spec.get("options", ())) or None, + field_tab=spec.get("tab"), nested_fields=None, ) ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 86e807447c1..56cf213f103 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9019,6 +9019,12 @@ def test_get_config_list_includes_anthropic_prompt_caching_fields(monkeypatch): assert fields["anthropic_prompt_caching_ttl"]["field_type"] == "Select" assert fields["anthropic_prompt_caching_ttl"]["field_value"] == "1h" assert fields["anthropic_prompt_caching_ttl"]["field_options"] == ["5m", "1h"] + + # Both caching fields carry their sub-tab so the Admin UI can render them on a + # dedicated Prompt Caching tab, while ungrouped fields stay on General. + assert fields["enable_anthropic_prompt_caching"]["field_tab"] == "prompt_caching" + assert fields["anthropic_prompt_caching_ttl"]["field_tab"] == "prompt_caching" + assert fields["budget_exceeded_throttle_percentage"]["field_tab"] is None finally: app.dependency_overrides.clear() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index af6bbdde8b9..8cea529d25d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -7,6 +7,7 @@ import { TableHeaderCell, TableCell, TableBody, + Title, Text, Button, Icon, @@ -21,6 +22,11 @@ import { StatusBadge } from "@/components/shared/table_cells"; import RouterSettings from "@/components/router_settings"; import Fallbacks from "@/components/Settings/RouterSettings/Fallbacks/Fallbacks"; import RoutingGroups from "@/components/routing_groups"; + +const PROMPT_CACHING_TAB = "prompt_caching"; +const ENABLE_ANTHROPIC_PROMPT_CACHING = "enable_anthropic_prompt_caching"; +const ANTHROPIC_PROMPT_CACHING_TTL = "anthropic_prompt_caching_ttl"; + interface GeneralSettingsPageProps { accessToken: string | null; userRole: string | null; @@ -34,6 +40,7 @@ interface generalSettingsItem { field_description: string; stored_in_db: boolean | null; field_options?: string[] | null; + field_tab?: string | null; } const SettingValueEditor: React.FC<{ @@ -83,6 +90,71 @@ const SettingValueEditor: React.FC<{ return null; }; +const PromptCachingPanel: React.FC<{ + accessToken: string; + settings: generalSettingsItem[]; + onChange: (fieldName: string, newValue: any) => void; +}> = ({ accessToken, settings, onChange }) => { + const enableSetting = settings.find((s) => s.field_name === ENABLE_ANTHROPIC_PROMPT_CACHING); + const ttlSetting = settings.find((s) => s.field_name === ANTHROPIC_PROMPT_CACHING_TTL); + + // The two rows come from the same registry the General tab reads; if they + // are not loaded yet there is nothing to render. + if (!enableSetting) { + return null; + } + + const enabled = enableSetting.field_value === true || enableSetting.field_value === "true"; + + // Apply immediately: a toggle and a dropdown are direct controls, so there is + // no separate Update button. Clearing the ttl resets it to the provider default. + const persist = (fieldName: string, value: any) => { + onChange(fieldName, value); + if (value === "" || value === null || value === undefined) { + deleteConfigFieldSetting(accessToken, fieldName); + } else { + updateConfigFieldSetting(accessToken, fieldName, value); + } + }; + + return ( + + Prompt Caching + + Automatically inject Anthropic prompt caching for every Anthropic and Bedrock Claude model, so clients that + never set cache_control themselves still get cached prompts. This is a single + gateway-wide switch; there is no per-model setup. + + +
+
+ Automatic Anthropic prompt caching +

{enableSetting.field_description}

+
+ persist(ENABLE_ANTHROPIC_PROMPT_CACHING, checked)} /> +
+ + {ttlSetting && ( +
+
+ Cache lifetime (TTL) +

{ttlSetting.field_description}

+
+ ({ label: option, value: option }))} + onChange={(newValue) => persist(ANTHROPIC_PROMPT_CACHING_TTL, newValue ?? "")} + /> +
+ )} + + ); +}; + const GeneralSettings: React.FC = ({ accessToken, userRole, userID }) => { const [generalSettings, setGeneralSettings] = useState([]); @@ -156,6 +228,7 @@ const GeneralSettings: React.FC = ({ accessToken, user Loadbalancing Routing Groups Fallbacks + Prompt Caching General @@ -168,6 +241,9 @@ const GeneralSettings: React.FC = ({ accessToken, user + + + @@ -181,7 +257,7 @@ const GeneralSettings: React.FC = ({ accessToken, user {generalSettings - .filter((value) => value.field_type !== "TypedDictionary") + .filter((value) => value.field_type !== "TypedDictionary" && value.field_tab !== PROMPT_CACHING_TAB) .map((value, index) => ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 45b06c9e44f..6dc63762e5c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22713,6 +22713,8 @@ export interface components { field_name: string; /** Field Options */ field_options?: string[] | null; + /** Field Tab */ + field_tab?: string | null; /** Field Type */ field_type: string; /** Field Value */ From 4e5f4884523ea124c6a252563624d104b4dc394c Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 17 Jul 2026 12:16:09 -0700 Subject: [PATCH 05/18] feat(ui): tighten the Prompt Caching descriptions The toggle and ttl descriptions were a wall of text, with a panel intro that mostly repeated the toggle description. Drop the intro and cut both descriptions to one or two lines, keeping a one-clause note that the cache is shared across callers on the same upstream credentials. --- litellm/proxy/proxy_server.py | 16 +++------------- .../_components/general_settings.tsx | 5 ----- 2 files changed, 3 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6725ecdb584..7d87207f7a9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14825,25 +14825,15 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, GeneralSettingsUILiteLLMFieldSpec "type": "Boolean", "tab": "prompt_caching", "description": ( - "Automatically add Anthropic cache_control breakpoints to the system prompt and the " - "trailing turn, for Anthropic and Bedrock Claude models that support prompt caching. " - "Lets clients that never set cache_control themselves still get cached prompts. " - "Requests that already carry their own cache_control are left untouched. " - "The provider caches a prefix against the upstream credentials that sent it, not per " - "end user, so this makes every caller's prompts cacheable on that shared account. " - "Leave this off if callers sharing a set of credentials must not learn whether " - "another caller recently sent a given prompt." + "Auto-adds cache_control to the system prompt and trailing turn for supported Anthropic " + "and Bedrock Claude models. The cache is shared across callers on the same upstream credentials." ), }, "anthropic_prompt_caching_ttl": { "type": "Select", "options": ("5m", "1h"), "tab": "prompt_caching", - "description": ( - "Cache lifetime for the breakpoints added by 'enable_anthropic_prompt_caching'. " - "Leave empty for Anthropic's 5m default. 1h suits long agentic sessions but doubles " - "the cache write premium." - ), + "description": "Empty uses Anthropic's 5m default. 1h suits long sessions but doubles the cache write cost.", }, } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index 8cea529d25d..1e8658d5104 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -120,11 +120,6 @@ const PromptCachingPanel: React.FC<{ return ( Prompt Caching - - Automatically inject Anthropic prompt caching for every Anthropic and Bedrock Claude model, so clients that - never set cache_control themselves still get cached prompts. This is a single - gateway-wide switch; there is no per-model setup. -
From 47ba9e76121dd7dbf572e112d7df5ebad5414bac Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 17 Jul 2026 19:38:42 -0700 Subject: [PATCH 06/18] fix(proxy): propagate the caching flag across workers via the safe-override allowlist enable_anthropic_prompt_caching and anthropic_prompt_caching_ttl are set as live litellm attributes on the worker that handles the UI save, exactly like budget_exceeded_throttle_percentage, but they were missing from LITELLM_SETTINGS_SAFE_DB_OVERRIDES, so a peer worker's config reload merged the DB value without applying it to the live attribute and stayed stale. Add both to the allowlist so they behave like the sibling field, and add test_general_settings_ui_fields_are_db_overridable so the UI registry and the override allowlist cannot drift again (the exact omission that caused this), plus a regression test that the flag flips on a simulated peer-worker reload. --- litellm/constants.py | 6 +++ tests/test_litellm/proxy/test_proxy_server.py | 45 +++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/litellm/constants.py b/litellm/constants.py index e104c937a9b..6432e2176c7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1517,6 +1517,12 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ "cost_discount_config", "cost_margin_config", "budget_exceeded_throttle_percentage", + # Every field editable from the Admin UI (proxy_server._GENERAL_SETTINGS_UI_LITELLM_FIELDS) + # must be listed here so a DB write from one worker overrides the live litellm attribute on + # the others when config reloads; otherwise peer workers stay on their startup value. + # test_general_settings_ui_fields_are_db_overridable enforces that pairing. + "enable_anthropic_prompt_caching", + "anthropic_prompt_caching_ttl", ] SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 56cf213f103..a100e7837f4 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9029,6 +9029,51 @@ def test_get_config_list_includes_anthropic_prompt_caching_fields(monkeypatch): app.dependency_overrides.clear() +def test_general_settings_ui_fields_are_db_overridable(): + """Every field the Admin UI can edit is a `litellm.` set via setattr on the handling + worker (`_persist_general_settings_ui_litellm_field`). Unless it is also in + LITELLM_SETTINGS_SAFE_DB_OVERRIDES, a config reload on a peer worker merges the DB value but + never applies it to the live attribute, so peer workers stay on their startup value. + + This invariant is the guard against the two registries drifting: adding a UI-editable field + without enrolling it in the DB-override allowlist silently breaks cross-worker propagation. + """ + from litellm.constants import LITELLM_SETTINGS_SAFE_DB_OVERRIDES + from litellm.proxy.proxy_server import _GENERAL_SETTINGS_UI_LITELLM_FIELDS + + missing = set(_GENERAL_SETTINGS_UI_LITELLM_FIELDS) - set(LITELLM_SETTINGS_SAFE_DB_OVERRIDES) + assert not missing, ( + f"UI-editable litellm_settings fields missing from LITELLM_SETTINGS_SAFE_DB_OVERRIDES: {sorted(missing)}. " + "Add them, or they will not propagate to other workers when changed from the UI." + ) + + +@pytest.mark.parametrize( + "field_name, db_value", + [ + ("enable_anthropic_prompt_caching", True), + ("anthropic_prompt_caching_ttl", "1h"), + ], +) +def test_prompt_caching_settings_propagate_on_config_reload(monkeypatch, field_name, db_value): + """A UI toggle on one worker persists to the DB; a peer worker picks it up only when the + config reload applies the safe-override allowlist. Regression for the fields being absent + from that allowlist, which left peer workers stale.""" + import litellm.proxy.proxy_server as ps + + # peer worker booted with the opposite/absent value + monkeypatch.setattr(litellm, field_name, False if isinstance(db_value, bool) else None) + + pc = ps.ProxyConfig() + pc._update_config_fields( + current_config={"litellm_settings": {}}, + param_name="litellm_settings", + db_param_value={field_name: db_value}, + ) + + assert getattr(litellm, field_name) == db_value + + def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypatch): """The flag defaults to False rather than None, so a plain 'is not None' check would report the default as 'In Config' and imply an admin had set it.""" From 99b85a3f2cac8fff501a07e6301274cc387ef245 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 17 Jul 2026 12:10:12 -0700 Subject: [PATCH 07/18] fix(mcp): persist config.yaml DCR clients in a server-scoped store Config.yaml-declared OAuth2 MCP servers using Dynamic Client Registration have no LiteLLM_MCPServerTable row, so the DCR persist path called update_mcp_server, which returns None for a missing row, then update_server(None), which dereferenced .approval_status and raised AttributeError. The exception was swallowed to a warning while /register still returned 200, so the minted client was never stored and every access-token expiry forced a full re-authorization Persist the acquired DCR client (client_id, client_secret, token_endpoint_auth_method, redirect_uris, encrypted at rest) in a dedicated LiteLLM_MCPServerOAuthClient store keyed by server_id when the server has no row, overlay it onto the in-memory config server so the refresh_token grant can authenticate within the process, and rehydrate it when the registry syncs from the database (which runs after the DB connects, unlike config load) so restarts and other pods pick it up. The store is encrypted at rest and is re-encrypted by the master-key rotation path alongside the server rows, through a shared helper so the two sites cannot diverge. The DB-backed server path is unchanged, and guarding the None return removes the swallowed-crash footgun Resolves the config.yaml DCR persistence regression introduced in v1.92.0 by #31912 --- .../migration.sql | 9 + .../litellm_proxy_extras/schema.prisma | 7 + litellm/proxy/_experimental/mcp_server/db.py | 82 +++- .../mcp_server/discoverable_endpoints.py | 144 +++++-- .../mcp_server/mcp_server_manager.py | 34 ++ litellm/proxy/schema.prisma | 7 + litellm/repositories/table_repositories.py | 4 + schema.prisma | 7 + .../mcp_server/test_db_credentials.py | 64 +++ .../mcp_server/test_discoverable_endpoints.py | 370 ++++++++++++++++++ .../mcp_server/test_mcp_sigv4_auth.py | 2 + 11 files changed, 679 insertions(+), 51 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_add_mcp_server_oauth_client_table/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_add_mcp_server_oauth_client_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_add_mcp_server_oauth_client_table/migration.sql new file mode 100644 index 00000000000..7aa6cdb1e33 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_add_mcp_server_oauth_client_table/migration.sql @@ -0,0 +1,9 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_MCPServerOAuthClient" ( + "server_id" TEXT NOT NULL, + "credentials" JSONB, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_MCPServerOAuthClient_pkey" PRIMARY KEY ("server_id") +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index f842bf13da9..a99cec49417 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -396,6 +396,13 @@ model LiteLLM_MCPUserEnvVars { @@index([server_id]) } +model LiteLLM_MCPServerOAuthClient { + server_id String @id + credentials Json? + created_at DateTime @default(now()) @map("created_at") + updated_at DateTime @default(now()) @updatedAt @map("updated_at") +} + // Generate Tokens for Proxy model LiteLLM_VerificationToken { token String @id diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index d55eb3ac014..7129582ff2a 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -33,6 +33,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( from litellm.proxy.utils import PrismaClient from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import ( + MCPServerOAuthClientRepository, MCPServerRepository, MCPUserCredentialsRepository, ) @@ -639,6 +640,7 @@ async def delete_mcp_server( for model, label in ( (prisma_client.db.litellm_mcpusercredentials, "credential"), (prisma_client.db.litellm_mcpuserenvvars, "env var"), + (prisma_client.db.litellm_mcpserveroauthclient, "OAuth client"), ): try: await model.delete_many(where={"server_id": server_id}) @@ -823,26 +825,66 @@ async def update_mcp_server( return updated_mcp_server -async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, touched_by: str, new_master_key: str): +async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, server_id: str) -> object | None: + """Read the persisted (encrypted) DCR OAuth client blob for a server from the + server-scoped store, or None. Config.yaml-declared servers have no + LiteLLM_MCPServerTable row, so their dynamically registered client lives here keyed + by server_id. The returned value is the raw credentials blob for + ``_get_persisted_dcr_credentials`` to parse.""" + row = await MCPServerOAuthClientRepository(prisma_client).table.find_unique(where={"server_id": server_id}) + if row is None: + return None + return row.credentials + + +async def upsert_mcp_server_oauth_client_credentials( + prisma_client: PrismaClient, server_id: str, credentials: MCPCredentials +) -> None: + """Persist a server's dynamically registered OAuth client (RFC 7591 DCR) in the + server-scoped store keyed by server_id, independent of any LiteLLM_MCPServerTable row. + client_id/client_secret are encrypted at rest with the same salt key used for the + server row's credentials blob, so ``_apply_persisted_dcr_credentials`` decrypts them the + same way regardless of which store a server's client came from.""" from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + encrypted = encrypt_credentials(credentials=dict(credentials), encryption_key=_get_salt_key()) + blob = safe_dumps(encrypted) + await MCPServerOAuthClientRepository(prisma_client).table.upsert( + where={"server_id": server_id}, + data={ + "create": {"server_id": server_id, "credentials": blob}, + "update": {"credentials": blob}, + }, + ) + + +def _reencrypt_mcp_credentials_blob(credentials: object, new_master_key: str) -> str | None: + """Decrypt an at-rest MCP credentials blob with the current key and re-encrypt it under + new_master_key, returning the serialized blob or None when there is nothing to rotate. Shared by + every table that stores an encrypted MCP credentials blob so a master-key rotation covers them + uniformly and cannot silently skip one.""" + if not credentials: + return None + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # noqa: PLC0415 # avoids circular import + + creds_dict = json.loads(credentials) if isinstance(credentials, str) else dict(credentials) + decrypted = decrypt_credentials(credentials=cast(MCPCredentials, creds_dict)) + encrypted = encrypt_credentials(credentials=decrypted, encryption_key=new_master_key) + return safe_dumps(encrypted) + + +async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, touched_by: str, new_master_key: str): + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # noqa: PLC0415 # avoids circular import + mcp_servers = await MCPServerRepository(prisma_client).table.find_many() updated = 0 for mcp_server in mcp_servers: update_data: Dict[str, Any] = {} - credentials = mcp_server.credentials - if credentials: - # Decrypt with current key first, then re-encrypt with new key - decrypted_credentials = decrypt_credentials( - credentials=cast(MCPCredentials, dict(credentials)), - ) - encrypted_credentials = encrypt_credentials( - credentials=decrypted_credentials, - encryption_key=new_master_key, - ) - update_data["credentials"] = safe_dumps(encrypted_credentials) + rotated_credentials = _reencrypt_mcp_credentials_blob(mcp_server.credentials, new_master_key) + if rotated_credentials is not None: + update_data["credentials"] = rotated_credentials rotated_env_vars = _reencrypt_global_env_var_values(mcp_server.env_vars, new_master_key) if rotated_env_vars is not None: @@ -857,9 +899,23 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, data=update_data, ) updated += 1 + + oauth_clients = await MCPServerOAuthClientRepository(prisma_client).table.find_many() + oauth_updated = 0 + for oauth_client in oauth_clients: + rotated_credentials = _reencrypt_mcp_credentials_blob(oauth_client.credentials, new_master_key) + if rotated_credentials is None: + continue + await MCPServerOAuthClientRepository(prisma_client).table.update( + where={"server_id": oauth_client.server_id}, + data={"credentials": rotated_credentials}, + ) + oauth_updated += 1 + verbose_proxy_logger.info( - "rotate_mcp_server_credentials_master_key: rotated %d MCP server row(s)", + "rotate_mcp_server_credentials_master_key: rotated %d MCP server row(s) and %d OAuth-client row(s)", updated, + oauth_updated, ) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 54aff86aab2..1af64749304 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -971,43 +971,93 @@ def _apply_persisted_dcr_credentials(mcp_server: MCPServer, credentials: _Persis return True -async def _get_persisted_mcp_server_with_dcr_client_id( - mcp_server: MCPServer, -) -> Optional[tuple["LiteLLM_MCPServerTable", _PersistedDcrCredentials]]: - from litellm.proxy._experimental.mcp_server.db import get_mcp_server # noqa: PLC0415 - from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 +async def _load_store_dcr_credentials(mcp_server: MCPServer) -> _PersistedDcrCredentials | None: + """DCR client persisted in the server-scoped OAuth-client store for a config-declared server + (which has no LiteLLM_MCPServerTable row). Returns None when the store has no usable client_id + or the DB is unreachable.""" + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # avoids circular import + get_mcp_server_oauth_client_credentials, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 # avoids circular import try: prisma_client = get_prisma_client_or_throw("Database not connected. Cannot read MCP OAuth client registration.") - persisted_mcp_server = await get_mcp_server( - prisma_client=prisma_client, - server_id=mcp_server.server_id, + blob = await get_mcp_server_oauth_client_credentials( + prisma_client=prisma_client, server_id=mcp_server.server_id ) - except Exception as exc: # noqa: BLE001 + except Exception as exc: # noqa: BLE001 # best-effort read; DB may be unreachable verbose_logger.debug( - "register_client_with_server: failed to read persisted DCR client registration for server_id=%s: %s", + "register_client_with_server: failed to read stored DCR client for server_id=%s: %s", mcp_server.server_id, exc, ) return None - if persisted_mcp_server is None: - return None - - credentials = _get_persisted_dcr_credentials(persisted_mcp_server.credentials) + credentials = _get_persisted_dcr_credentials(blob) if credentials is None or not credentials.client_id: return None + return credentials - return persisted_mcp_server, credentials + +async def hydrate_config_server_dcr_client(mcp_server: MCPServer) -> bool: + """Overlay a config-declared server's persisted DCR client onto its in-memory object so token + refresh can authenticate. Config.yaml servers have no LiteLLM_MCPServerTable row, so their + minted client lives in the server-scoped store; without this overlay the in-memory server + carries no client_id after a restart. An explicit client_id set in config.yaml wins and is never + overwritten by a persisted store client.""" + if mcp_server.client_id: + return False + credentials = await _load_store_dcr_credentials(mcp_server) + if credentials is None: + return False + return _apply_persisted_dcr_credentials(mcp_server, credentials) + + +async def _resolve_persisted_dcr_client( + mcp_server: MCPServer, +) -> tuple[Optional["LiteLLM_MCPServerTable"], _PersistedDcrCredentials | None]: + """Resolve a server's persisted DCR client using the same two-level rule the write path uses, so + read and write always agree. First, whether the server HAS a LiteLLM_MCPServerTable row: a row is + always resolved to that row and the store is never consulted for a server that has a row, so a + caller-chosen server_id colliding with a config-declared server cannot inherit that config + server's client, and a row that exists but carries no usable client_id yields (row, None) rather + than a store fallback. Second, among rowless servers: a config-declared server keeps its client in + the server-scoped store, while a rowless non-config server is a throwaway temp/session server with + no persisted client. Returns (row_or_None, credentials_or_None); the row is only needed by the + reuse path to refresh the registry for a DB-declared server.""" + from litellm.proxy._experimental.mcp_server.db import get_mcp_server # noqa: PLC0415 # avoids circular import + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # avoids circular import + global_mcp_server_manager, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 # avoids circular import + + try: + prisma_client = get_prisma_client_or_throw("Database not connected. Cannot read MCP OAuth client registration.") + row = await get_mcp_server(prisma_client=prisma_client, server_id=mcp_server.server_id) + except Exception as exc: # noqa: BLE001 # best-effort read; DB may be unreachable + verbose_logger.debug( + "register_client_with_server: failed to read persisted DCR client for server_id=%s: %s", + mcp_server.server_id, + exc, + ) + return None, None + + if row is not None: + credentials = _get_persisted_dcr_credentials(row.credentials) + if credentials is not None and credentials.client_id: + return row, credentials + return row, None + if global_mcp_server_manager.is_config_declared_server(mcp_server.server_id): + return None, await _load_store_dcr_credentials(mcp_server) + return None, None async def _reuse_persisted_dcr_client_if_available( mcp_server: MCPServer, current_redirect_uri: Optional[str] = None ) -> bool: - persisted = await _get_persisted_mcp_server_with_dcr_client_id(mcp_server) - if persisted is None: + persisted_mcp_server, credentials = await _resolve_persisted_dcr_client(mcp_server) + if credentials is None: return False - persisted_mcp_server, credentials = persisted if current_redirect_uri is not None and _redirect_uri_not_registered(credentials, current_redirect_uri): verbose_logger.debug( "register_client_with_server: not reusing persisted DCR client for server_id=%s; its registered " @@ -1021,18 +1071,19 @@ async def _reuse_persisted_dcr_client_if_available( if not _apply_persisted_dcr_credentials(mcp_server, credentials): return False - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 - global_mcp_server_manager, - ) - - try: - await global_mcp_server_manager.update_server(persisted_mcp_server) - except Exception as exc: # noqa: BLE001 - verbose_logger.warning( - "register_client_with_server: failed to refresh persisted DCR client registration for server_id=%s: %s", - mcp_server.server_id, - exc, + if persisted_mcp_server is not None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # avoids circular import + global_mcp_server_manager, ) + + try: + await global_mcp_server_manager.update_server(persisted_mcp_server) + except Exception as exc: # noqa: BLE001 # best-effort registry refresh + verbose_logger.warning( + "register_client_with_server: failed to refresh persisted DCR client registration for server_id=%s: %s", + mcp_server.server_id, + exc, + ) return bool(mcp_server.client_id) @@ -1044,10 +1095,9 @@ async def _persisted_dcr_redirect_uri_is_stale(mcp_server: MCPServer, current_re otherwise short-circuits registration before any redirect check can run. Servers without a persisted DCR recording (admin-configured client_id, or registered before redirect_uris were recorded) are never reported stale.""" - persisted = await _get_persisted_mcp_server_with_dcr_client_id(mcp_server) - if persisted is None: + _, credentials = await _resolve_persisted_dcr_client(mcp_server) + if credentials is None: return False - _, credentials = persisted if not _redirect_uri_not_registered(credentials, current_redirect_uri): return False verbose_logger.warning( @@ -1067,7 +1117,10 @@ DcrRegistrationPersistenceResult = Literal["persisted", "reused", "skipped", "fa async def _persist_dcr_client_registration( mcp_server: MCPServer, registration_response: object, current_redirect_uri: str ) -> DcrRegistrationPersistenceResult: - """Persist the dynamically registered OAuth client (RFC 7591) onto the MCP server row. + """Persist the dynamically registered OAuth client (RFC 7591) to its single home: the server's + ``LiteLLM_MCPServerTable`` row when it has one, otherwise the server-scoped store when the server + is config-declared. A rowless server that is not config-declared is a throwaway temp/session + server, so its client is overlaid in memory only and not persisted. The interactive authorization_code flow mints a ``client_id`` via Dynamic Client Registration that discovery cannot re-derive; without persisting it the autonomous @@ -1106,16 +1159,20 @@ async def _persist_dcr_client_registration( if await _reuse_persisted_dcr_client_if_available(mcp_server, current_redirect_uri=current_redirect_uri): return "reused" + token_endpoint_auth_method = ( + "client_secret_basic" if registration.token_endpoint_auth_method == "client_secret_basic" else None + ) credentials: MCPCredentials = { "client_id": registration.client_id, "client_secret": registration.client_secret, - "token_endpoint_auth_method": ( - "client_secret_basic" if registration.token_endpoint_auth_method == "client_secret_basic" else None - ), + "token_endpoint_auth_method": token_endpoint_auth_method, "redirect_uris": [current_redirect_uri], } - from litellm.proxy._experimental.mcp_server.db import update_mcp_server # noqa: PLC0415 + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # avoids circular import + update_mcp_server, + upsert_mcp_server_oauth_client_credentials, + ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 global_mcp_server_manager, ) @@ -1136,7 +1193,18 @@ async def _persist_dcr_client_registration( ), touched_by="mcp_oauth_dcr", ) - await global_mcp_server_manager.update_server(updated_row) + if updated_row is not None: + await global_mcp_server_manager.update_server(updated_row) + return "persisted" + if global_mcp_server_manager.is_config_declared_server(mcp_server.server_id): + await upsert_mcp_server_oauth_client_credentials( + prisma_client=prisma_client, + server_id=mcp_server.server_id, + credentials=credentials, + ) + mcp_server.client_id = registration.client_id + mcp_server.client_secret = registration.client_secret + mcp_server.token_endpoint_auth_method = token_endpoint_auth_method return "persisted" except Exception as exc: # noqa: BLE001 verbose_logger.warning( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 115ff2e492c..941bd45f98c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1127,6 +1127,14 @@ class MCPServerManager: """ return self.config_mcp_servers | self.registry + def is_config_declared_server(self, server_id: str) -> bool: + """True when server_id was declared in config.yaml (present in the in-memory config map). + Config servers are rowless and persistent, so their DCR client belongs in the server-scoped + store; a rowless server that is NOT config-declared is a throwaway temp/session server whose + client must not be persisted. This never overrides the row-existence check: a server that has + a LiteLLM_MCPServerTable row is always resolved to that row first.""" + return server_id in self.config_mcp_servers + async def load_servers_from_config( self, mcp_servers_config: dict[str, Any], @@ -1367,8 +1375,32 @@ class MCPServerManager: verbose_logger.debug(f"Loaded MCP Servers: {json.dumps(self.config_mcp_servers, indent=4, default=str)}") + await self._hydrate_config_servers_dcr_clients() + self.initialize_tool_name_to_mcp_server_name_mapping() + async def _hydrate_config_servers_dcr_clients(self) -> None: + """Overlay each config-declared server's persisted DCR client (from the server-scoped + store) onto its in-memory object so token refresh authenticates after a restart. A + best-effort no-op when the DB is unreachable at config-load time.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( # noqa: PLC0415 # circular import + hydrate_config_server_dcr_client, + ) + + for server in self.config_mcp_servers.values(): + try: + if await hydrate_config_server_dcr_client(server): + verbose_logger.debug( + "hydrated persisted DCR client onto config MCP server server_id=%s", + server.server_id, + ) + except Exception as exc: # noqa: BLE001 # best-effort hydration; never fail config load + verbose_logger.debug( + "load_servers_from_config: failed to hydrate DCR client for server_id=%s: %s", + server.server_id, + exc, + ) + async def _register_openapi_tools(self, spec_path: str, server: MCPServer, base_url: str): """ Register tools from an OpenAPI specification for a given server. @@ -4968,6 +5000,8 @@ class MCPServerManager: verbose_logger.debug("MCP registry refreshed (%s servers in registry)", len(registered_registry)) + await self._hydrate_config_servers_dcr_clients() + def get_mcp_servers_from_ids(self, server_ids: list[str]) -> list[MCPServer]: servers = [] registry = self.get_registry() diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f842bf13da9..a99cec49417 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -396,6 +396,13 @@ model LiteLLM_MCPUserEnvVars { @@index([server_id]) } +model LiteLLM_MCPServerOAuthClient { + server_id String @id + credentials Json? + created_at DateTime @default(now()) @map("created_at") + updated_at DateTime @default(now()) @updatedAt @map("updated_at") +} + // Generate Tokens for Proxy model LiteLLM_VerificationToken { token String @id diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 7ce4607e1ca..dc2a7d25259 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -77,6 +77,10 @@ class MCPUserCredentialsRepository(PrismaTableRepository): table_name = "litellm_mcpusercredentials" +class MCPServerOAuthClientRepository(PrismaTableRepository): + table_name = "litellm_mcpserveroauthclient" + + class PromptRepository(PrismaTableRepository): table_name = "litellm_prompttable" diff --git a/schema.prisma b/schema.prisma index f842bf13da9..a99cec49417 100644 --- a/schema.prisma +++ b/schema.prisma @@ -396,6 +396,13 @@ model LiteLLM_MCPUserEnvVars { @@index([server_id]) } +model LiteLLM_MCPServerOAuthClient { + server_id String @id + credentials Json? + created_at DateTime @default(now()) @map("created_at") + updated_at DateTime @default(now()) @updatedAt @map("updated_at") +} + // Generate Tokens for Proxy model LiteLLM_VerificationToken { token String @id diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 7269774442b..a245200c4d1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -978,3 +978,67 @@ def test_prepare_mcp_server_data_update_carries_token_exchange_columns(): assert data["audience"] == "https://upstream.example.com" assert data["subject_token_type"] == "urn:ietf:params:oauth:token-type:jwt" assert data["token_exchange_profile"] == "entra_obo" + + +@pytest.mark.asyncio +async def test_master_key_rotation_reencrypts_oauth_client_store(monkeypatch): + """The server-scoped DCR client store (LiteLLM_MCPServerOAuthClient) is encrypted at rest, so a + master-key rotation must re-encrypt it alongside the server rows. Skipping it leaves + config-declared DCR clients under the retired key, where they decrypt back to ciphertext and + force a full re-authorization.""" + import litellm.proxy.common_utils.encrypt_decrypt_utils as enc + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.proxy._experimental.mcp_server.db import ( + decrypt_credentials, + encrypt_credentials, + rotate_mcp_server_credentials_master_key, + ) + + key_old, key_new = "salt-old-key", "salt-new-key" + + blob_old = safe_dumps( + encrypt_credentials( + credentials={"client_id": "cid-123", "client_secret": "sec-456"}, + encryption_key=key_old, + ) + ) + + monkeypatch.setattr(enc, "_get_salt_key", lambda: key_old) + + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_mcpserveroauthclient.find_many = AsyncMock( + return_value=[SimpleNamespace(server_id="config_faros", credentials=blob_old)] + ) + store_update = AsyncMock() + prisma.db.litellm_mcpserveroauthclient.update = store_update + + await rotate_mcp_server_credentials_master_key(prisma, touched_by="test", new_master_key=key_new) + + store_update.assert_awaited_once() + assert store_update.await_args.kwargs["where"] == {"server_id": "config_faros"} + rotated_blob = store_update.await_args.kwargs["data"]["credentials"] + + monkeypatch.setattr(enc, "_get_salt_key", lambda: key_new) + recovered = decrypt_credentials(credentials=json.loads(rotated_blob)) + assert recovered["client_id"] == "cid-123" + assert recovered["client_secret"] == "sec-456" + + +@pytest.mark.asyncio +async def test_delete_mcp_server_cleans_oauth_client_store(): + """Deleting a server must remove its server-scoped DCR client store entry alongside the per-user + credential and env-var rows, or a re-created server reusing the same server_id would inherit the + deleted server's OAuth client.""" + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + prisma = MagicMock() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=SimpleNamespace(server_id="s1")) + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock() + prisma.db.litellm_mcpuserenvvars.delete_many = AsyncMock() + prisma.db.litellm_mcpserveroauthclient.delete_many = AsyncMock() + + await delete_mcp_server(prisma, "s1", invalidate_token_cache=AsyncMock()) + + prisma.db.litellm_mcpserveroauthclient.delete_many.assert_awaited_once_with(where={"server_id": "s1"}) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index f5ac229d119..6f2f24df8fa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7130,3 +7130,373 @@ async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): assert response.status_code == 502 body = json.loads(response.body) assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} + + +@pytest.mark.asyncio +async def test_persist_dcr_client_for_config_server_uses_side_store(): + """A config.yaml-declared OAuth2 DCR server has no LiteLLM_MCPServerTable row, so + update_mcp_server returns None. The minted client must then persist to the server-scoped + OAuth-client store keyed by server_id (never a shadow server row), overlay onto the in-memory + server so refresh can authenticate this process, and never call update_server(None) (which + previously raised AttributeError on .approval_status, was swallowed, and reported a 200 that + persisted nothing).""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _persist_dcr_client_registration, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + config_server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + + mock_upsert = AsyncMock() + mock_update_server = AsyncMock() + + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.update_mcp_server", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", + new=mock_upsert, + ), + patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), + ): + result = await _persist_dcr_client_registration( + mcp_server=config_server, + registration_response={ + "client_id": "minted-client", + "client_secret": "minted-secret", + "token_endpoint_auth_method": "client_secret_basic", + }, + current_redirect_uri="https://proxy.litellm.example/callback", + ) + + assert result == "persisted" + + mock_upsert.assert_called_once() + assert mock_upsert.call_args.kwargs["server_id"] == "config_faros" + stored = mock_upsert.call_args.kwargs["credentials"] + assert stored["client_id"] == "minted-client" + assert stored["client_secret"] == "minted-secret" + assert stored["token_endpoint_auth_method"] == "client_secret_basic" + assert stored["redirect_uris"] == ["https://proxy.litellm.example/callback"] + + assert config_server.client_id == "minted-client" + assert config_server.client_secret == "minted-secret" + assert config_server.token_endpoint_auth_method == "client_secret_basic" + + mock_update_server.assert_not_called() + + +@pytest.mark.asyncio +async def test_hydrate_config_server_applies_stored_dcr_client(monkeypatch): + """On restart a config server's in-memory object has no client_id; hydration overlays the + persisted DCR client from the server-scoped store, decrypting the encrypted-at-rest blob, so the + refresh_token grant can authenticate as the registered client instead of re-authenticating.""" + import litellm.proxy.common_utils.encrypt_decrypt_utils as enc + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.proxy._experimental.mcp_server.db import encrypt_credentials + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + hydrate_config_server_dcr_client, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-hydrate-key") + stored_blob = safe_dumps( + encrypt_credentials( + credentials={ + "client_id": "stored-client", + "client_secret": "stored-secret", + "token_endpoint_auth_method": "client_secret_basic", + "redirect_uris": ["https://proxy.litellm.example/callback"], + }, + encryption_key="salt-hydrate-key", + ) + ) + assert "stored-client" not in stored_blob and "stored-secret" not in stored_blob + + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=AsyncMock(return_value=stored_blob), + ), + ): + applied = await hydrate_config_server_dcr_client(server) + + assert applied is True + assert server.client_id == "stored-client" + assert server.client_secret == "stored-secret" + assert server.token_endpoint_auth_method == "client_secret_basic" + + +@pytest.mark.asyncio +async def test_reuse_config_server_reads_store_with_real_crypto(monkeypatch): + """A config-declared server (rowless) keeps its DCR client in the store, so the reuse read + resolves it from the store and decrypts the encrypted-at-rest client, mirroring the write path so + a re-authorize reuses the client instead of re-minting one.""" + import litellm.proxy.common_utils.encrypt_decrypt_utils as enc + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.proxy._experimental.mcp_server.db import encrypt_credentials + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-reuse-key") + blob = safe_dumps( + encrypt_credentials( + credentials={"client_id": "stored-client", "client_secret": "sec", "redirect_uris": ["https://x/callback"]}, + encryption_key="salt-reuse-key", + ) + ) + assert "stored-client" not in blob + store_lookup = AsyncMock(return_value=blob) + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_lookup, + ), + ): + result = await _reuse_persisted_dcr_client_if_available(server, current_redirect_uri="https://x/callback") + + assert result is True + assert server.client_id == "stored-client" + store_lookup.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_temp_server_is_not_persisted_to_store(): + """A rowless server that is NOT config-declared (a throwaway /server/oauth/session server) must + not leave a permanent store row on persist, and the read must never consult the store for it. Its + minted client is overlaid in memory for the session only.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _persist_dcr_client_registration, + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + temp = MCPServer( + server_id="temp-uuid", + name="temp", + server_name="temp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + authorization_url="https://p.example/authorize", + token_url="https://p.example/token", + registration_url="https://p.example/register", + ) + + upsert = AsyncMock() + store_read = AsyncMock(return_value=None) + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=False), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=AsyncMock(return_value=None)), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), + patch("litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", new=upsert), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_read, + ), + patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()), + ): + result = await _persist_dcr_client_registration( + temp, {"client_id": "temp-client", "client_secret": "s"}, "https://x/callback" + ) + reused = await _reuse_persisted_dcr_client_if_available( + MCPServer( + server_id="temp-uuid", + name="temp", + server_name="temp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ), + current_redirect_uri="https://x/callback", + ) + + assert result == "persisted" + assert temp.client_id == "temp-client" + upsert.assert_not_called() + store_read.assert_not_called() + assert reused is False + + +@pytest.mark.asyncio +async def test_hydrate_does_not_overwrite_explicit_config_client_id(): + """An explicit client_id set in config.yaml wins: hydration must not overwrite it with a stale + persisted store client, and must not even read the store when config already supplied a client.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + hydrate_config_server_dcr_client, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_static", + name="config_static", + server_name="config_static", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="explicit-from-config", + ) + store_read = AsyncMock( + return_value={"client_id": "stale-store-client", "client_secret": "x", "redirect_uris": []} + ) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_read, + ), + ): + applied = await hydrate_config_server_dcr_client(server) + + assert applied is False + assert server.client_id == "explicit-from-config" + store_read.assert_not_called() + + +@pytest.mark.asyncio +async def test_reuse_does_not_inherit_store_client_when_a_row_exists(): + """Security: a server that HAS a LiteLLM_MCPServerTable row reads its DCR client only from that + row, never from the server-scoped store. server_id is caller-settable on create, so a submitted + server whose id collides with a config-declared server must not be able to load that config + server's client from the store and send it to its own token endpoint. A row that exists but has + no client_id yields no reusable client and must not fall back to the store.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + submitted = MCPServer( + server_id="collides_with_config", + name="submitted", + server_name="submitted", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + row_without_client = MagicMock() + row_without_client.credentials = None + row_without_client.server_id = "collides_with_config" + store_lookup = AsyncMock( + return_value={"client_id": "config-secret-client", "client_secret": "leak", "redirect_uris": []} + ) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=AsyncMock(return_value=row_without_client), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_lookup, + ), + ): + result = await _reuse_persisted_dcr_client_if_available(submitted, current_redirect_uri="https://x/callback") + + assert result is False + assert submitted.client_id is None + store_lookup.assert_not_called() + + +@pytest.mark.asyncio +async def test_load_servers_from_config_hydrates_dcr_clients(): + """load_servers_from_config must invoke DCR-client hydration so config servers pick up their + persisted client on startup; deleting the call site leaves a restarted server with no client_id + and forces re-authentication on every token expiry.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + hydrate_spy = AsyncMock() + with patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy): + await global_mcp_server_manager.load_servers_from_config({}) + + hydrate_spy.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_reload_servers_from_database_hydrates_dcr_clients(): + """load_servers_from_config runs before the DB connects at startup, so its hydration no-ops; + reload_servers_from_database runs after the DB connects and must hydrate config servers' persisted + DCR clients too, or a fresh pod has no client_id for a config server and forces re-authentication + on the first token refresh.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + + hydrate_spy = AsyncMock() + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma, + ), + patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy), + ): + await global_mcp_server_manager.reload_servers_from_database() + + hydrate_spy.assert_awaited_once() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py index 5992fd1814f..f6b61c1d9f7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -989,6 +989,7 @@ class TestRotateCredentials: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[server]) mock_prisma.db.litellm_mcpservertable.update = AsyncMock() + mock_prisma.db.litellm_mcpserveroauthclient.find_many = AsyncMock(return_value=[]) with ( patch( @@ -1036,6 +1037,7 @@ class TestRotateCredentials: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[server]) mock_prisma.db.litellm_mcpservertable.update = AsyncMock() + mock_prisma.db.litellm_mcpserveroauthclient.find_many = AsyncMock(return_value=[]) with ( patch( From 07e07e6e2b0dd27f9bd50180ed8d916cc32068f0 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 21:17:49 -0700 Subject: [PATCH 08/18] fix(vertex_ai): exclude Gemini Google Search grounding tokens from input token billing (#33742) * fix(vertex_ai): exclude Google Search grounding tokens from Gemini input token billing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): stub get_configured_token_limits on mocked routers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../vertex_and_google_ai_studio_gemini.py | 32 +++++- ...test_vertex_and_google_ai_studio_gemini.py | 97 +++++++++++++++++++ .../test_model_management_endpoints.py | 2 + .../test_team_model_name_translation.py | 6 ++ 4 files changed, 136 insertions(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 3193b72a7d9..624190a0b61 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1744,6 +1744,30 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) return non_thinking_tokens == usage_metadata.get("totalTokenCount", 0) + @staticmethod + def _response_has_search_grounding( + completion_response: Union[GenerateContentResponseBody, BidiGenerateContentServerMessage], + ) -> bool: + """ + Whether the response used Grounding with Google Search, detected via + groundingMetadata.webSearchQueries (an actual web search was performed). + + Google bills grounding-with-Google-Search retrieved tokens separately (a per-request / + per-query search fee) and excludes them from input token billing, unlike URL context / + File Search / code execution whose tool-use tokens are charged at the input token rate. + URL context also emits groundingMetadata (with groundingChunks but no webSearchQueries), + so presence of groundingMetadata alone is not a sufficient signal. + See https://ai.google.dev/gemini-api/docs/pricing and + https://github.com/BerriAI/litellm/discussions/33198 + """ + if "candidates" not in completion_response: + return False + for candidate in completion_response["candidates"] or []: + grounding_metadata, _, _, _ = VertexGeminiConfig._extract_candidate_metadata(candidate) + if VertexGeminiConfig._calculate_web_search_requests(grounding_metadata): + return True + return False + @staticmethod def _calculate_usage( completion_response: Union[GenerateContentResponseBody, BidiGenerateContentServerMessage], @@ -1899,12 +1923,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): tool_use_tokens=tool_use_prompt_tokens, ) + billable_tool_use_prompt_tokens = ( + 0 + if VertexGeminiConfig._response_has_search_grounding(completion_response) + else (tool_use_prompt_tokens or 0) + ) + completion_tokens = response_tokens or completion_response["usageMetadata"].get("candidatesTokenCount", 0) if not VertexGeminiConfig.is_candidate_token_count_inclusive(usage_metadata) and reasoning_tokens: completion_tokens = reasoning_tokens + completion_tokens ## GET USAGE ## usage = Usage( - prompt_tokens=usage_metadata.get("promptTokenCount", 0) + (tool_use_prompt_tokens or 0), + prompt_tokens=usage_metadata.get("promptTokenCount", 0) + billable_tool_use_prompt_tokens, completion_tokens=completion_tokens, total_tokens=usage_metadata.get("totalTokenCount", 0), prompt_tokens_details=prompt_tokens_details, diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 5adc5b76990..95e8e6561f1 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -547,6 +547,103 @@ def test_vertex_ai_non_grounded_usage_omits_tool_use_tokens(): assert not hasattr(usage.prompt_tokens_details, "tool_use_tokens") +def test_response_has_search_grounding_detection(): + """ + Only groundingMetadata.webSearchQueries signals an actual Google Search. URL context also + emits groundingMetadata (groundingChunks but no webSearchQueries) and must not be treated + as search grounding. + """ + assert ( + VertexGeminiConfig._response_has_search_grounding( + {"candidates": [{"groundingMetadata": {"webSearchQueries": ["latest nobel physics"]}}]} + ) + is True + ) + assert ( + VertexGeminiConfig._response_has_search_grounding( + { + "candidates": [ + { + "urlContextMetadata": {"urlMetadata": []}, + "groundingMetadata": { + "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] + }, + } + ] + } + ) + is False + ) + assert ( + VertexGeminiConfig._response_has_search_grounding({"candidates": [{"groundingMetadata": {"webSearchQueries": []}}]}) + is False + ) + assert VertexGeminiConfig._response_has_search_grounding({"candidates": []}) is False + assert VertexGeminiConfig._response_has_search_grounding({}) is False + + +def test_vertex_ai_search_grounding_tool_use_tokens_excluded_from_prompt_tokens(): + """ + Grounding with Google Search retrieved tokens are not billed at the input token rate + (Google charges a separate per-request / per-query search fee), so toolUsePromptTokenCount + must be surfaced on prompt_tokens_details.tool_use_tokens but excluded from prompt_tokens. + See https://ai.google.dev/gemini-api/docs/pricing and + https://github.com/BerriAI/litellm/discussions/33198 + """ + v = VertexGeminiConfig() + completion_response = { + "candidates": [{"groundingMetadata": {"webSearchQueries": ["latest nobel physics"]}}], + "usageMetadata": UsageMetadata( + promptTokenCount=19, + candidatesTokenCount=304, + thoughtsTokenCount=122, + toolUsePromptTokenCount=142, + totalTokenCount=587, + ), + } + + usage = v._calculate_usage(completion_response=completion_response) + + assert usage.prompt_tokens == 19 + assert usage.completion_tokens == 304 + 122 + assert usage.total_tokens == 587 + assert usage.prompt_tokens_details.tool_use_tokens == 142 + assert usage.total_tokens - usage.prompt_tokens - usage.completion_tokens == 142 + + +def test_vertex_ai_url_context_tool_use_tokens_billed_as_input_tokens(): + """ + URL context / File Search / code execution tool-use tokens are billed as input tokens, so + toolUsePromptTokenCount is folded into prompt_tokens when the response is not search grounded. + """ + v = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "urlContextMetadata": {"urlMetadata": []}, + "groundingMetadata": { + "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] + }, + } + ], + "usageMetadata": UsageMetadata( + promptTokenCount=19, + candidatesTokenCount=304, + thoughtsTokenCount=122, + toolUsePromptTokenCount=142, + totalTokenCount=587, + ), + } + + usage = v._calculate_usage(completion_response=completion_response) + + assert usage.prompt_tokens == 19 + 142 + assert usage.completion_tokens == 304 + 122 + assert usage.total_tokens == 587 + assert usage.prompt_tokens_details.tool_use_tokens == 142 + assert usage.total_tokens - usage.prompt_tokens - usage.completion_tokens == 0 + + def test_streaming_chunk_includes_reasoning_tokens(): from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator, diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 8c6bdefedae..79c5f3ea549 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1727,6 +1727,7 @@ class TestModelInfoEndpoint: "gpt-3.5-turbo", ] mock_router.get_model_access_groups.return_value = {} + mock_router.get_configured_token_limits.return_value = (None, None) mock_get_key_models.return_value = ["gpt-4", "claude-3"] mock_get_team_models.return_value = ["gpt-3.5-turbo"] mock_get_complete_models.return_value = [ @@ -1812,6 +1813,7 @@ class TestModelInfoEndpoint: # Setup mocks mock_router.get_model_names.return_value = ["team-model-1"] mock_router.get_model_access_groups.return_value = {} + mock_router.get_configured_token_limits.return_value = (None, None) mock_get_key_models.return_value = [] mock_get_team_models.return_value = ["team-model-1"] mock_get_complete_models.return_value = ["team-model-1"] diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 25e84fb59a7..577af3dcffc 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -725,6 +725,7 @@ async def test_v1_models_translates_team_model_for_access_group_key(monkeypatch) router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() + router.get_configured_token_limits.return_value = (None, None) router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] @@ -766,6 +767,7 @@ async def test_v1_models_keeps_internal_names_when_public_name_flag_disabled( router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() + router.get_configured_token_limits.return_value = (None, None) router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] @@ -800,6 +802,7 @@ async def test_v1_models_translates_team_model_with_metadata(monkeypatch): router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() + router.get_configured_token_limits.return_value = (None, None) router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] router.get_model_group_info.return_value = None @@ -845,6 +848,7 @@ async def test_v1_models_metadata_fallbacks_use_internal_routing_key(monkeypatch router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() + router.get_configured_token_limits.return_value = (None, None) router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] # Fallbacks are keyed on the internal routing name, as the router stores them. @@ -901,6 +905,7 @@ async def test_v1_models_metadata_does_not_leak_other_team_fallbacks(monkeypatch router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() + router.get_configured_token_limits.return_value = (None, None) router.model_list = [team_x, team_y] router.get_model_list.return_value = [team_x, team_y] router.fallbacks = [ @@ -1155,6 +1160,7 @@ def test_translate_team_model_names_for_listing_respects_legacy_flag(): def _public_named_router(*team_rows: dict) -> MagicMock: router = MagicMock() router.get_model_list.return_value = list(team_rows) + router.get_configured_token_limits.return_value = (None, None) return router From b3d05bd10b9a044ea08a1f1ce0e165ee5ba1ef35 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 21:33:34 -0700 Subject: [PATCH 09/18] feat(fireworks_ai): map litellm session id to x-session-affinity header for prompt caching (#33717) * feat(fireworks_ai): map litellm session id to x-session-affinity header for prompt caching Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): normalize cached usage in spend logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(fireworks_ai): initialize chat config base class Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(fireworks_ai): normalize cached usage for spend logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(fireworks_ai): cover cached usage normalization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): normalize cached usage in spend logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(fireworks_ai): cover session id precedence Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yuneng-jiang Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/fireworks_ai/chat/transformation.py | 14 ++- litellm/llms/fireworks_ai/common_utils.py | 24 ++++- .../spend_tracking/spend_tracking_utils.py | 6 ++ .../test_fireworks_ai_chat_transformation.py | 100 ++++++++++++++++++ .../test_spend_tracking_utils.py | 63 +++++++++++ 5 files changed, 204 insertions(+), 3 deletions(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index d4258557fe7..319f03fea89 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -48,7 +48,7 @@ from ...openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, ) -from ..common_utils import FireworksAIException +from ..common_utils import FireworksAIMixin, FireworksAIException def _extract_fireworks_hidden_params(payload: dict) -> dict: @@ -70,7 +70,7 @@ def _extract_fireworks_hidden_params(payload: dict) -> dict: return {**top_level, **per_choice} -class FireworksAIConfig(OpenAIGPTConfig): +class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): """ Reference: https://docs.fireworks.ai/api-reference/post-chatcompletions @@ -114,6 +114,16 @@ class FireworksAIConfig(OpenAIGPTConfig): prompt_truncate_len: Optional[int] = None, context_length_exceeded_behavior: Optional[Literal["error", "truncate"]] = None, ) -> None: + OpenAIGPTConfig.__init__( + self, + frequency_penalty=frequency_penalty, + max_tokens=max_tokens, + n=n, + stop=stop, + temperature=temperature, + top_p=top_p, + response_format=response_format, + ) locals_ = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index a1b6309d1e0..4e22445bcc0 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -12,6 +12,23 @@ class FireworksAIException(BaseLLMException): pass +def get_fireworks_session_id(litellm_params: dict) -> str | None: + params = litellm_params + for key in ("litellm_session_id", "session_id"): + value = params.get(key) + if value: + return str(value) + metadata = params.get("metadata") + if isinstance(metadata, dict): + value = metadata.get("session_id") + if value: + return str(value) + value = params.get("litellm_trace_id") + if value: + return str(value) + return None + + class FireworksAIMixin: """ Common Base Config functions across Fireworks AI Endpoints @@ -47,4 +64,9 @@ class FireworksAIMixin: if api_key is None: raise ValueError("FIREWORKS_API_KEY is not set") - return {"Authorization": "Bearer {}".format(api_key), **headers} + validated_headers = {"Authorization": "Bearer {}".format(api_key), **headers} + if not any(key.lower() == "x-session-affinity" for key in validated_headers): + session_id = get_fireworks_session_id(litellm_params) + if session_id: + validated_headers["x-session-affinity"] = session_id + return validated_headers diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index b38d5e39800..23e7711b223 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -373,6 +373,12 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs if isinstance(v, BaseModel): v = v.model_dump() additional_usage_values.update({k: v}) + if "cache_read_input_tokens" not in additional_usage_values: + prompt_tokens_details = additional_usage_values.get("prompt_tokens_details") + if isinstance(prompt_tokens_details, dict): + cached_tokens = prompt_tokens_details.get("cached_tokens") + if isinstance(cached_tokens, int) and cached_tokens > 0: + additional_usage_values["cache_read_input_tokens"] = cached_tokens clean_metadata["additional_usage_values"] = additional_usage_values if litellm.cache is not None: 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 03e763a4161..6809799d34f 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 @@ -13,6 +13,7 @@ sys.path.insert( from litellm import get_model_info, supports_reasoning, supports_vision from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig +from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import ( ChatCompletionMessageToolCall, Function, @@ -32,6 +33,105 @@ def force_local_model_cost(monkeypatch): litellm.model_cost = get_model_cost_map(url=litellm.model_cost_map_url) +def test_validate_environment_sets_session_affinity_from_litellm_session_id(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"litellm_session_id": "session-123"}, + api_key="test-key", + ) + + assert headers["x-session-affinity"] == "session-123" + + +def test_validate_environment_sets_session_affinity_from_metadata_session_id(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"metadata": {"session_id": "metadata-session-123"}}, + api_key="test-key", + ) + + assert headers["x-session-affinity"] == "metadata-session-123" + + +def test_validate_environment_sets_session_affinity_from_session_id(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"session_id": "session-id-123"}, + api_key="test-key", + ) + + assert headers["x-session-affinity"] == "session-id-123" + + +def test_validate_environment_sets_session_affinity_from_trace_id(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"litellm_trace_id": "trace-id-123"}, + api_key="test-key", + ) + + assert headers["x-session-affinity"] == "trace-id-123" + + +def test_validate_environment_does_not_set_session_affinity_without_session_id(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-key", + ) + + assert "x-session-affinity" not in headers + + +def test_validate_environment_preserves_explicit_session_affinity_header(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={"x-session-affinity": "explicit-session"}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"litellm_session_id": "session-123"}, + api_key="test-key", + ) + + assert headers["x-session-affinity"] == "explicit-session" + + +def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id(): + assert ( + get_fireworks_session_id( + {"litellm_session_id": "session-123", "litellm_trace_id": "trace-123"} + ) + == "session-123" + ) + + def test_handle_message_content_with_tool_calls(): config = FireworksAIConfig() message = Message( diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9a8f8146d6f..51d72aa2ab2 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -46,6 +46,69 @@ from litellm.types.utils import ( ) +def _get_additional_usage_values_for_usage(usage: litellm.Usage) -> dict: + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj=litellm.ModelResponse( + id="chatcmpl-test", + choices=[], + usage=usage, + ), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + metadata = json.loads(payload["metadata"]) + return metadata["additional_usage_values"] + + +def test_get_logging_payload_maps_openai_cached_tokens_to_cache_read_input_tokens(): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + prompt_tokens_details={"cached_tokens": 123}, + ) + ) + + assert additional_usage_values["cache_read_input_tokens"] == 123 + assert additional_usage_values["prompt_tokens_details"]["cached_tokens"] == 123 + + +def test_get_logging_payload_preserves_anthropic_cache_read_input_tokens(): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + prompt_tokens_details={"cached_tokens": 123}, + cache_read_input_tokens=456, + ) + ) + + assert additional_usage_values["cache_read_input_tokens"] == 456 + + +@pytest.mark.parametrize( + "prompt_tokens_details", + [None, {"cached_tokens": 0}], +) +def test_get_logging_payload_does_not_map_missing_or_zero_cached_tokens(prompt_tokens_details): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + prompt_tokens_details=prompt_tokens_details, + ) + ) + + assert "cache_read_input_tokens" not in additional_usage_values + + def test_sanitize_request_body_for_spend_logs_payload_basic(): request_body = { "messages": [{"role": "user", "content": "Hello, how are you?"}], From 010b20072d20f043650ab654e2c0190b1c9da1fb Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 10:26:48 -0700 Subject: [PATCH 10/18] fix(router): enforce context-window pre-call checks for Responses API input (#33706) * fix(router): enforce context-window pre-call checks for Responses API input * test(router): cover _count_pre_call_check_tokens across API surfaces * fix(router): count Responses instructions and skip pre-call token count when no input * fix(router): forward Responses input into deployment selection for context-window checks --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 56 +++++++++- tests/test_litellm/test_router.py | 179 ++++++++++++++++++++++++++++++ 2 files changed, 229 insertions(+), 6 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b1a5405ebf1..0b1471dc527 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4461,6 +4461,7 @@ class Router: model=model, request_kwargs=kwargs, messages=kwargs.get("messages", None), + input=kwargs.get("input", None), specific_deployment=kwargs.pop("specific_deployment", None), ) except Exception as e: @@ -4608,6 +4609,7 @@ class Router: deployment = self.get_available_deployment( model=model, messages=kwargs.get("messages", None), + input=kwargs.get("input", None), specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, ) @@ -10002,11 +10004,44 @@ class Router: client = self.cache.get_cache(key=cache_key, parent_otel_span=parent_otel_span) return client + def _count_pre_call_check_tokens( + self, + messages: list[dict[str, str]] | None, + input: str | list | None, + instructions: str | None = None, + ) -> int: + """ + Count input tokens for context-window pre-call checks. + + Chat Completions send `messages`; the Responses API sends `input` (a string or + a list of Responses input items) plus an optional `instructions` system prompt. + The Responses payload is normalized to chat messages via the shared + LiteLLMCompletionResponsesConfig transform so the same token_counter path covers + both API surfaces and `instructions` tokens are included in the count. + """ + if messages is not None: + return litellm.token_counter(messages=messages) + if input is not None: + from openai.types.responses.response_create_params import ResponseInputParam + + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + typed_input = cast(str | ResponseInputParam, input) # cast-ok: str | list matches transform input + input_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=typed_input, + responses_api_request={"instructions": instructions} if instructions is not None else {}, + ) + return litellm.token_counter(messages=cast(list, input_messages)) # cast-ok: transformed chat messages + raise ValueError("Either messages or input must be provided to count tokens") + def _pre_call_checks( self, model: str, healthy_deployments: List, - messages: List[Dict[str, str]], + messages: list[dict[str, str]] | None = None, + input: str | list | None = None, request_kwargs: Optional[dict] = None, ): """ @@ -10036,6 +10071,10 @@ class Router: _rate_limit_error = False parent_otel_span = _get_parent_otel_span_from_kwargs(request_kwargs) + raw_instructions = request_kwargs.get("instructions") if request_kwargs else None + instructions = raw_instructions if isinstance(raw_instructions, str) else None + has_countable_input = messages is not None or input is not None + ## get model group RPM ## dt = get_utc_datetime() current_minute = dt.strftime("%H-%M") @@ -10058,10 +10097,12 @@ class Router: _deployment_model = base_model or _litellm_params.get("model", None) max_input_tokens = model_info.get("max_input_tokens") if isinstance(model_info, dict) else None - if isinstance(max_input_tokens, int): + if isinstance(max_input_tokens, int) and has_countable_input: if input_tokens is None: try: - input_tokens = litellm.token_counter(messages=messages) + input_tokens = self._count_pre_call_check_tokens( + messages=messages, input=input, instructions=instructions + ) except Exception as e: verbose_router_logger.error( "litellm.router.py::_pre_call_checks: failed to count tokens. Returning initial list of deployments. Got - {}".format( @@ -10526,11 +10567,12 @@ class Router: parent_otel_span=parent_otel_span, ) - if self.enable_pre_call_checks and messages is not None: + if self.enable_pre_call_checks and (messages is not None or input is not None): healthy_deployments = self._pre_call_checks( model=model, healthy_deployments=cast(List[Dict], healthy_deployments), messages=messages, + input=input, request_kwargs=request_kwargs, ) # check if user wants to do tag based routing @@ -11041,11 +11083,12 @@ class Router: healthy_deployments = self._filter_blocked_deployments(healthy_deployments) # filter pre-call checks - if self.enable_pre_call_checks and messages is not None: + if self.enable_pre_call_checks and (messages is not None or input is not None): healthy_deployments = self._pre_call_checks( model=model, healthy_deployments=healthy_deployments, messages=messages, + input=input, request_kwargs=request_kwargs, ) @@ -11195,11 +11238,12 @@ class Router: pass_through_deployments = self._filter_blocked_deployments(pass_through_deployments) # 5. Apply pre-call checks (if enabled) - if self.enable_pre_call_checks and messages is not None: + if self.enable_pre_call_checks and (messages is not None or input is not None): pass_through_deployments = self._pre_call_checks( model=model, healthy_deployments=pass_through_deployments, messages=messages, + input=input, request_kwargs=request_kwargs, ) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 55c09e6cac4..76a6e3c1bbe 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2864,6 +2864,185 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch assert calls == [1] +def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): + """ + Responses API calls pass `input` (str) instead of `messages`. Context-window + checks must count tokens from `input` and filter deployments over the limit. Uses + the real token_counter so the transform + counting path is a true regression guard. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} + ) + + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="m", + healthy_deployments=deployments, + input="a very long prompt that exceeds the tiny context window", + ) + + +def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): + """ + Responses API `input` can be a list of input items. It must be normalized to + chat messages and counted so oversized requests are filtered out. Uses the real + token_counter (no mock) so the transform + counting path is a true regression guard. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} + ) + + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="m", + healthy_deployments=deployments, + input=[ + {"role": "user", "content": "count these tokens against the one token limit please"}, + ], + ) + + +def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): + """ + Responses API `instructions` become a system message the model receives, so their + tokens must be counted too. A request whose `input` alone fits under the limit but + whose `input` + `instructions` exceeds it must be filtered (regression for the + context-window check under-filtering when instructions were ignored). + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + + short_input = "hi" + long_instructions = "you are a helpful assistant. " * 20 + + input_only_tokens = router._count_pre_call_check_tokens(messages=None, input=short_input) + with_instructions_tokens = router._count_pre_call_check_tokens( + messages=None, input=short_input, instructions=long_instructions + ) + assert with_instructions_tokens > input_only_tokens + + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} + ) + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="m", + healthy_deployments=deployments, + input=short_input, + request_kwargs={"instructions": long_instructions}, + ) + + +def test_count_pre_call_check_tokens_across_api_surfaces(): + """ + _count_pre_call_check_tokens must count tokens from chat `messages`, a Responses + API string `input`, and a Responses API list `input`, and raise when given neither. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + ) + + messages_tokens = router._count_pre_call_check_tokens( + messages=[{"role": "user", "content": "hello world"}], input=None + ) + string_input_tokens = router._count_pre_call_check_tokens(messages=None, input="hello world") + list_input_tokens = router._count_pre_call_check_tokens( + messages=None, input=[{"role": "user", "content": "hello world"}] + ) + + assert messages_tokens > 0 + assert string_input_tokens > 0 + assert list_input_tokens > 0 + + with pytest.raises(ValueError): + router._count_pre_call_check_tokens(messages=None, input=None) + + +def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): + """ + When neither messages nor input is provided (e.g. endpoints without prompt text), + token counting is skipped gracefully and all deployments are returned. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) + + counted: list[dict] = [] + original = router._count_pre_call_check_tokens + monkeypatch.setattr( + router, + "_count_pre_call_check_tokens", + lambda **kwargs: counted.append(kwargs) or original(**kwargs), + ) + + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + result = router._pre_call_checks(model="m", healthy_deployments=deployments) + assert len(result) == 1 + assert counted == [] # token counting skipped entirely, so no misleading error is logged + + +@pytest.mark.asyncio +async def test_aresponses_enforces_context_window_pre_call_check(): + """ + End-to-end router regression: a Responses API call whose `input` exceeds the + deployment's max_input_tokens must be filtered by the pre-call check, raising + ContextWindowExceededError instead of being silently routed. This guards the + wiring that forwards `input` from the generic-call path into deployment selection + (the deployment uses mock_response, so the check must trip before any real call). + """ + router = litellm.Router( + model_list=[ + { + "model_name": "small-ctx", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + "model_info": {"max_input_tokens": 5}, + } + ], + enable_pre_call_checks=True, + ) + with pytest.raises(litellm.ContextWindowExceededError): + await router.aresponses( + model="small-ctx", + input="this responses input is definitely much longer than five tokens for sure", + ) + + def test_get_deployment_model_info_base_model_flow(): """Test that get_deployment_model_info correctly handles the base model flow""" from unittest.mock import patch From 4a297dd6114cdaa1ba6795c33b45bc9b8f5fdd8b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 10:52:27 -0700 Subject: [PATCH 11/18] fix(otel): restore proxy-level error.* attributes on v2 failure spans (LIT-4179) (#33664) * fix(otel): restore proxy-level error.* attributes on v2 failure spans (LIT-4179) * refactor(otel): narrow v2 failure hook return type to drop fastapi import (LIT-4179) --------- Co-authored-by: yucheng-berri --- litellm/integrations/otel/emitter.py | 55 ++++++--- litellm/integrations/otel/logger.py | 79 +++++++++++- litellm/proxy/proxy_server.py | 21 ++-- .../integrations/otel/test_otel_v2_emitter.py | 42 +++++++ .../integrations/otel/test_otel_v2_logger.py | 114 ++++++++++++++++++ .../proxy_server/test_exception_handlers.py | 41 +++++++ 6 files changed, 328 insertions(+), 24 deletions(-) diff --git a/litellm/integrations/otel/emitter.py b/litellm/integrations/otel/emitter.py index f97f8b8394c..8651cf586cd 100644 --- a/litellm/integrations/otel/emitter.py +++ b/litellm/integrations/otel/emitter.py @@ -72,6 +72,42 @@ def _stamp_litellm_error_attributes(span: Span, error: SpanError) -> None: span.set_attribute(LiteLLMError.LLM_PROVIDER, error.llm_provider) +def stamp_error( + span: Span, + error: SpanError, + *, + record_event: bool = True, + set_status: bool = True, +) -> tuple[str, str] | None: + """Stamp the full v2 error attribute set on ``span`` and return the resolved + ``(error_type, message)`` pair, or ``None`` when the error carries neither a + type nor a message. + + Shared by the LLM-call span (``finish_span``) and the proxy-level failure + spans (the FastAPI SERVER span and the ``auth`` phase span) so every v2 error + span carries identical keys. The semconv ``exception`` event rides alongside + the attributes so backends that map unknown string attrs to a truncated + ``keyword`` (e.g. Elasticsearch's 1024-char ``ignore_above``) still see the + full untruncated message on the recognized event field. ``record_event`` and + ``set_status`` are opt-outs for callers whose span lifecycle (``use_span``) or + owner (the FastAPI instrumentor) already records the event or the status. + """ + if not (error.error_type or error.message): + return None + error_type = error.error_type or "error" + message = error.message or error.error_type or "error" + _stamp_otel_error_attributes(span, error_type, message) + _stamp_litellm_error_attributes(span, error) + if set_status: + span.set_status(Status(StatusCode.ERROR, message)) + if record_event: + span.add_event( + ExceptionEvent.NAME, + {ExceptionEvent.TYPE: error_type, ExceptionEvent.MESSAGE: message}, + ) + return error_type, message + + class SpanEmitter: def __init__( self, @@ -212,21 +248,10 @@ class SpanEmitter: ) else None ) - if error and (error.error_type or error.message): - error_type = error.error_type or "error" - message = error.message or error.error_type or "error" - _stamp_otel_error_attributes(span, error_type, message) - _stamp_litellm_error_attributes(span, error) - span.set_status(Status(StatusCode.ERROR, message)) - # Also emit the semconv ``exception`` event so backends that - # dynamic-map unknown string span attrs to ``keyword`` (e.g. - # Elasticsearch with a 1024-char ``ignore_above``) still see the - # full untruncated message on the recognized event field. - span.add_event( - ExceptionEvent.NAME, - {ExceptionEvent.TYPE: error_type, ExceptionEvent.MESSAGE: message}, - ) - if self._event_recorder is not None and role is SpanRole.LLM_CALL: + if error: + stamped = stamp_error(span, error) + if stamped is not None and self._event_recorder is not None and role is SpanRole.LLM_CALL: + error_type, message = stamped self._event_recorder.record_operation_exception( span_context=span.get_span_context(), error_type=error_type, diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index be72fabd387..778f5342e90 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -24,7 +24,7 @@ from litellm.integrations.otel.plumbing.context import ( set_request_baggage, set_request_root_span, ) -from litellm.integrations.otel.emitter import SpanEmitter +from litellm.integrations.otel.emitter import SpanEmitter, stamp_error from litellm.integrations.otel.mappers import resolve_mappers from litellm.integrations.otel.model.metadata import ( LLMCallEvent, @@ -59,6 +59,7 @@ from litellm.integrations.otel.model.spans import SpanRole, span_role_for_servic from litellm.integrations.otel.model.utils import to_ns if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import ( StandardLoggingGuardrailInformation, StandardLoggingPayload, @@ -66,6 +67,33 @@ if TYPE_CHECKING: LITELLM_TRACER_NAME = "litellm" + +def _span_error_from_exception( + exception: "Exception | None", + *, + status_code: int | None = None, + traceback_str: str | None = None, +) -> SpanError: + """A ``SpanError`` for a proxy-level failure that never produced a + ``StandardLoggingPayload`` (auth / validation / malformed-body rejections), + mirroring ``_parse_error``'s field mapping so it stamps the same v2 keys a + failed LLM call does. ``status_code`` pins ``error.code`` to the real response + status, matching v1's SERVER-span behavior.""" + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + info = StandardLoggingPayloadSetup.get_error_information( + original_exception=exception, + traceback_str=traceback_str, + ) + return SpanError( + error_type=info.get("error_class") or info.get("error_code") or None, + message=info.get("error_message") or None, + code=str(status_code) if status_code is not None else (info.get("error_code") or None), + stack_trace=info.get("traceback") or None, + llm_provider=info.get("llm_provider") or None, + ) + + # Any callback whose class belongs to one of these modules is "the OTel # callback" for proxy-global-registration purposes. _OTEL_MODULES = ( @@ -558,7 +586,12 @@ class OpenTelemetryV2(CustomLogger): def start_phase_span(self, name: str) -> "Iterator[Span]": span = self._emitter.start_span(SpanRole.SERVICE, name) with use_span(span, end_on_exit=True): - yield span + try: + yield span + except Exception as exc: + if is_recordable_span(span): + stamp_error(span, _span_error_from_exception(exc), record_event=False, set_status=False) + raise async def async_pre_call_hook( self, @@ -573,6 +606,48 @@ class OpenTelemetryV2(CustomLogger): ) return data + def record_error_attributes_on_span( + self, + span: "Span | None", + exception: "Exception | None", + status_code: int, + ) -> None: + """Stamp the v2 error.* attributes on the FastAPI-owned SERVER span for a + failure that dies before any LLM-call span exists (malformed body, auth / + validation rejection). Called from the proxy's global exception handler via + ``_close_dangling_otel_server_span``. The instrumentor still owns the span's + status and lifecycle, so this only decorates it — never sets status, never + ends it — and emits no exception event, matching v1's SERVER-span behavior + and avoiding a duplicate of the event ``async_post_call_failure_hook`` or + the ``auth`` phase span already records.""" + if span is None or not is_recordable_span(span): + return + stamp_error( + span, + _span_error_from_exception(exception, status_code=status_code), + record_event=False, + set_status=False, + ) + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: "UserAPIKeyAuth", + traceback_str: "str | None" = None, + ) -> None: + """Stamp error.* on the request's root SERVER span for a proxy-level + failure that never reached an LLM call (empty body rejected in the + endpoint, auth failure), so the failed request carries the same error keys + a failed LLM call does. v1's ``OpenTelemetry`` implemented this same hook; + v2 lost it when it stopped subclassing ``OpenTelemetry``, which is the + LIT-4179 regression for pre-call failures.""" + span = request_root_span() or user_api_key_dict.parent_otel_span + if span is None or not is_recordable_span(span): + return None + stamp_error(span, _span_error_from_exception(original_exception, traceback_str=traceback_str)) + return None + def emit_guardrail_span(self, entry: "StandardLoggingGuardrailInformation") -> None: # Emitted by the guardrail-recording code the moment a guardrail finishes, # not from a post-call hook — that hook does not fire on every path (a diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3b640ee54fd..aed345c5db4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1395,19 +1395,25 @@ def _close_dangling_otel_server_span(request: Request, status_code: int, exc: Op if open_telemetry_logger is None: return # Under OTel V2 the FastAPI instrumentor owns the server span (parent_otel_span - # is that same span), and it records the error + ends it itself. Ending it here - # would end it early — losing the http.* attributes the instrumentor stamps on - # completion — and double-end it. Leave it to the instrumentor. + # is that same span) and ends it itself with the http.* attributes stamped on + # completion. The instrumentor only records an error when the exception reaches + # it uncaught, but these handlers swallow it into a JSONResponse, so it never + # does; stamp the error.* attributes here (without ending or re-statusing the + # span, which the instrumentor still owns) so pre-call failures carry the error + # like v1 did. Otherwise close and annotate the dangling span ourselves. try: from litellm.integrations.otel.model.config import is_otel_v2_enabled - if is_otel_v2_enabled(): - return + v2_enabled = is_otel_v2_enabled() except Exception: - pass + v2_enabled = False try: from opentelemetry.trace import Status, StatusCode + if v2_enabled: + if status_code >= 400: + open_telemetry_logger.record_error_attributes_on_span(parent_otel_span, exc, status_code) + return open_telemetry_logger.set_response_status_code_attribute(parent_otel_span, status_code) if status_code >= 400: open_telemetry_logger.record_error_attributes_on_span(parent_otel_span, exc, status_code) @@ -1416,7 +1422,8 @@ def _close_dangling_otel_server_span(request: Request, status_code: int, exc: Op except Exception as e: verbose_proxy_logger.debug("Error closing dangling OTEL SERVER span: %s", str(e)) finally: - request.state.parent_otel_span = None + if not v2_enabled: + request.state.parent_otel_span = None @app.exception_handler(RequestValidationError) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index 48190a798da..6b1da4c2952 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -16,10 +16,12 @@ from litellm.integrations.otel import ( # noqa: E402 from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402 from litellm.integrations.otel.plumbing import providers # noqa: E402 from litellm.integrations.otel.emitter import SpanEmitter # noqa: E402 +from litellm.integrations.otel.emitter import stamp_error # noqa: E402 from litellm.integrations.otel.model.payloads import ( # noqa: E402 GuardrailSpanData, LLMCallSpanData, ServiceSpanData, + SpanError, ) from litellm.integrations.otel.model.spans import SPAN_REGISTRY, SpanRole # noqa: E402 @@ -155,6 +157,46 @@ def test_error_span_sets_status_and_error_type(): assert span.attributes["error.type"] == "RateLimitError" +def test_stamp_error_writes_full_attribute_set_and_event(): + engine, exporter = _engine() + span = engine.start_span(SpanRole.PROXY_REQUEST, "POST /chat/completions") + result = stamp_error( + span, SpanError("ProxyException", "boom", code="401", stack_trace="tb", llm_provider="anthropic") + ) + span.end() + (s,) = exporter.get_finished_spans() + assert result == ("ProxyException", "boom") + assert s.attributes["error.type"] == "ProxyException" + assert s.attributes["error.message"] == "boom" + assert s.attributes["litellm.provider.error.code"] == "401" + assert s.attributes["litellm.provider.error.stack_trace"] == "tb" + assert s.attributes["litellm.provider.error.llm_provider"] == "anthropic" + assert s.status.status_code is StatusCode.ERROR + assert [e.name for e in s.events] == ["exception"] + + +def test_stamp_error_opt_outs_skip_status_and_event(): + engine, exporter = _engine() + span = engine.start_span(SpanRole.PROXY_REQUEST, "POST /chat/completions") + stamp_error(span, SpanError("ProxyException", "boom", code="401"), record_event=False, set_status=False) + span.end() + (s,) = exporter.get_finished_spans() + assert s.attributes["error.type"] == "ProxyException" + assert s.attributes["litellm.provider.error.code"] == "401" + assert s.status.status_code is StatusCode.UNSET + assert s.events == () + + +def test_stamp_error_without_type_or_message_is_noop(): + engine, exporter = _engine() + span = engine.start_span(SpanRole.PROXY_REQUEST, "POST /chat/completions") + assert stamp_error(span, SpanError()) is None + span.end() + (s,) = exporter.get_finished_spans() + assert "error.type" not in s.attributes + assert s.status.status_code is StatusCode.UNSET + + def test_hierarchy_and_kinds_match_registry(): engine, exporter = _engine() data = LLMCallSpanData.from_standard_logging_payload(_payload()) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index b5e077e3561..5f6002f4cdf 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -860,6 +860,120 @@ def test_guardrail_span_anchors_to_root_inside_active_phase_span(): assert guard.parent.span_id != auth_span.get_span_context().span_id +# --------------------------------------------------------------------------- # +# LIT-4179 — proxy-level failures that never reach an LLM call must still stamp +# the structured error.* attributes onto the request's spans, restoring the v1 +# behavior v2 dropped when it stopped subclassing ``OpenTelemetry``. +# --------------------------------------------------------------------------- # + + +def _proxy_exc(message, code): + from litellm.proxy._types import ProxyException + + return ProxyException(message=message, type="bad_request_error", param=None, code=code) + + +def test_async_post_call_failure_hook_stamps_error_on_root_span(): + """PATH B: an endpoint-level failure (empty body rejected before dispatch) + reaches ``async_post_call_failure_hook``; it must stamp error.* + an exception + event on the anchored request root span.""" + from litellm.proxy._types import UserAPIKeyAuth + + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + set_request_root_span(server) + exc = _proxy_exc("litellm.BadRequestError: messages is required", 400) + result = asyncio.run( + logger.async_post_call_failure_hook( + request_data={}, original_exception=exc, user_api_key_dict=UserAPIKeyAuth() + ) + ) + server.end() + assert result is None + (span,) = exporter.get_finished_spans() + assert span.attributes["error.type"] == "ProxyException" + assert "messages is required" in span.attributes["error.message"] + assert span.attributes["litellm.provider.error.code"] == "400" + assert span.status.status_code is StatusCode.ERROR + assert any(e.name == "exception" for e in span.events) + + +def test_async_post_call_failure_hook_falls_back_to_user_api_key_parent_span(): + """With no anchor set (a path that never captured the root), the hook must fall + back to ``user_api_key_dict.parent_otel_span`` rather than dropping the error.""" + from litellm.proxy._types import UserAPIKeyAuth + + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + asyncio.run( + logger.async_post_call_failure_hook( + request_data={}, + original_exception=_proxy_exc("boom", 401), + user_api_key_dict=UserAPIKeyAuth(parent_otel_span=server), + ) + ) + server.end() + (span,) = exporter.get_finished_spans() + assert span.attributes["error.type"] == "ProxyException" + assert span.attributes["litellm.provider.error.code"] == "401" + + +def test_record_error_attributes_on_span_decorates_without_ending(): + """PATH A: a failure that dies before any LLM-call span (malformed body, + validation) is stamped onto the instrumentor-owned SERVER span. The method must + not end the span or emit a duplicate exception event, and must pin error.code + to the real response status (not the exception's own code).""" + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + logger.record_error_attributes_on_span(server, _proxy_exc("Invalid JSON body", 400), 422) + assert server.is_recording() + server.end() + (span,) = exporter.get_finished_spans() + assert span.attributes["error.type"] == "ProxyException" + assert span.attributes["error.message"] == "Invalid JSON body" + assert span.attributes["litellm.provider.error.code"] == "422" + assert all(e.name != "exception" for e in span.events) + + +def test_record_error_attributes_on_span_ignores_below_400_and_missing_span(): + logger, _ = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + logger.record_error_attributes_on_span(None, _proxy_exc("boom", 400), 400) # no span → no-op + logger.record_error_attributes_on_span(server, None, 400) # no exception → no-op + server.end() + assert "error.type" not in (server.attributes or {}) + + +def test_start_phase_span_stamps_error_attributes_on_failure(): + """An ``auth`` phase span that dies (expired key) must carry the structured + error.* attributes, not only the exception event ``use_span`` records.""" + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + set_request_root_span(server) + exc = _proxy_exc("Authentication Error, ExpiredToken", 401) + with trace.use_span(server, end_on_exit=False): + with contextlib.suppress(Exception): + with logger.start_phase_span("auth /chat/completions"): + raise exc + server.end() + by_name = {s.name: s for s in exporter.get_finished_spans()} + auth = by_name["auth /chat/completions"] + assert auth.attributes["error.type"] == "ProxyException" + assert "ExpiredToken" in auth.attributes["error.message"] + assert auth.attributes["litellm.provider.error.code"] == "401" + assert auth.status.status_code is StatusCode.ERROR + assert any(e.name == "exception" for e in auth.events) + + +def test_start_phase_span_success_carries_no_error(): + logger, exporter = _logger() + with logger.start_phase_span("auth /chat/completions"): + pass + (span,) = exporter.get_finished_spans() + assert "error.type" not in span.attributes + assert span.status.status_code is not StatusCode.ERROR + + def test_real_logging_pre_call_opens_span_end_to_end(): """Regression guard: a real ``LiteLLMLoggingObj.pre_call`` must fire ``log_pre_api_call`` on the V2 logger (via ``litellm.input_callback``), so the diff --git a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py index cf92f9cd12b..e4bf06991b4 100644 --- a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py +++ b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py @@ -124,6 +124,47 @@ def test_close_dangling_otel_server_span_records_status_and_ends(monkeypatch): } +def test_close_dangling_otel_server_span_v2_stamps_error_without_ending(monkeypatch): + """LIT-4179: under OTel v2 the FastAPI instrumentor owns the SERVER span, so + the handler must only stamp error.* on it (via record_error_attributes_on_span) + and must NOT set status, end the span, or clear request state — otherwise the + instrumentor's http.* attributes and span close are lost.""" + import litellm.integrations.otel.model.config as otel_config + import litellm.proxy.proxy_server as ps + + span = MagicMock() + fake_logger = MagicMock() + monkeypatch.setattr(ps, "open_telemetry_logger", fake_logger, raising=False) + monkeypatch.setattr(otel_config, "is_otel_v2_enabled", lambda: True) + request = _make_request(parent_otel_span=span) + exc = ProxyException(message="bad", type="bad_request_error", param=None, code=400) + + _close_dangling_otel_server_span(request=request, status_code=422, exc=exc) + + fake_logger.record_error_attributes_on_span.assert_called_once_with(span, exc, 422) + assert not span.end.called + assert not span.set_status.called + assert not fake_logger.set_response_status_code_attribute.called + assert request.state.parent_otel_span is span + + +def test_close_dangling_otel_server_span_v2_success_does_not_stamp(monkeypatch): + """Under v2 a sub-400 status must not stamp an error onto the SERVER span.""" + import litellm.integrations.otel.model.config as otel_config + import litellm.proxy.proxy_server as ps + + span = MagicMock() + fake_logger = MagicMock() + monkeypatch.setattr(ps, "open_telemetry_logger", fake_logger, raising=False) + monkeypatch.setattr(otel_config, "is_otel_v2_enabled", lambda: True) + request = _make_request(parent_otel_span=span) + + _close_dangling_otel_server_span(request=request, status_code=200) + + assert not fake_logger.record_error_attributes_on_span.called + assert not span.end.called + + def test_close_dangling_otel_server_span_missing_span_is_noop_error(): """When parent_otel_span is missing the call short-circuits — no error.""" request = _make_request(parent_otel_span=None) From 6f4f4f69df2e2369e95235eaa8a6c0e1aea5a6aa Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 18 Jul 2026 11:24:03 -0700 Subject: [PATCH 12/18] refactor(ui): consolidate Add/Edit credential modals into one CredentialModal (#32572) * refactor(ui): consolidate Add/Edit credential modals into one CredentialModal AddCredentialModal and EditCredentialModal were ~90% identical: the same provider select, ProviderSpecificFields, and submit/filter logic, differing only in title, button text, edit-mode prefill, and the disabled credential name. Replace both with a single CredentialModal driven by a mode: 'add' | 'edit' prop, and point the two call sites in credentials.tsx at it. Removes ~120 lines of duplication and drops the no-explicit-any and no-restricted-imports baselines. The two per-file tests merge into one CredentialModal.test.tsx covering both modes (add: editable empty name; edit: prefilled, disabled name; provider fields render). * refactor(ui): derive credential name disabled state from mode, not data The disabled flag on the credential name field was tied to whether existingCredential?.credential_name is truthy, an artifact of the old EditCredentialModal. Drive it from the isEdit flag like the rest of the component so mode='add' with a stray existingCredential can't disable the field and mode='edit' with an empty name can't leave it editable. Behavior is unchanged for real call sites; adds a regression test for the edit-with- empty-name case. * refactor(ui): prefill credential form declaratively instead of via useEffect The edit-mode form was seeded with an imperative form.setFieldsValue inside a useEffect that also set React state (setSelectedProvider), an antd anti- pattern carried over from the old EditCredentialModal. Both call sites mount the modal fresh with existingCredential already present (conditional && plus destroyOnHidden), so there is no 'prop arrives after mount' case to handle. Replace it with antd's declarative initialValues on the Form and a lazy useState initializer for the provider. Removes the effect, its react-hooks/set-state-in-effect suppression and exhaustive-deps warning, and one any cast; behavior is unchanged (edit now shows the real provider on first paint instead of flashing the default). Existing tests cover prefill and the disabled name field. --- ui/litellm-dashboard/eslint-suppressions.json | 10 +- .../model_add/AddCredentialModal.test.tsx | 108 ------------- .../model_add/CredentialModal.test.tsx | 140 ++++++++++++++++ ...redentialModal.tsx => CredentialModal.tsx} | 70 ++++---- .../model_add/EditCredentialModal.test.tsx | 123 -------------- .../model_add/EditCredentialModal.tsx | 150 ------------------ .../src/components/model_add/credentials.tsx | 13 +- 7 files changed, 191 insertions(+), 423 deletions(-) delete mode 100644 ui/litellm-dashboard/src/components/model_add/AddCredentialModal.test.tsx create mode 100644 ui/litellm-dashboard/src/components/model_add/CredentialModal.test.tsx rename ui/litellm-dashboard/src/components/model_add/{AddCredentialModal.tsx => CredentialModal.tsx} (71%) delete mode 100644 ui/litellm-dashboard/src/components/model_add/EditCredentialModal.test.tsx delete mode 100644 ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 90b0c84244e..dcf482450e9 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1885,19 +1885,11 @@ "count": 1 } }, - "src/components/model_add/AddCredentialModal.tsx": { + "src/components/model_add/CredentialModal.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/model_add/EditCredentialModal.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/model_add/credentials.tsx": { "no-restricted-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.test.tsx b/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.test.tsx deleted file mode 100644 index aee7a0cdd1d..00000000000 --- a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.test.tsx +++ /dev/null @@ -1,108 +0,0 @@ -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render, screen, waitFor } from "@testing-library/react"; -import { describe, expect, it, vi } from "vitest"; -import { Providers } from "../provider_info_helpers"; -import AddCredentialModal from "./AddCredentialModal"; - -vi.mock("../networking", async () => { - const actual = await vi.importActual("../networking"); - return { - ...actual, - getProviderCreateMetadata: vi.fn().mockResolvedValue([ - { - provider: "OpenAI", - provider_display_name: Providers.OpenAI, - litellm_provider: "openai", - default_model_placeholder: "gpt-3.5-turbo", - credential_fields: [ - { - key: "api_key", - label: "OpenAI API Key", - field_type: "password", - required: true, - }, - { - key: "api_base", - label: "API Base", - field_type: "text", - placeholder: "https://api.openai.com/v1", - }, - ], - }, - { - provider: "Anthropic", - provider_display_name: Providers.Anthropic, - litellm_provider: "anthropic", - default_model_placeholder: "claude-3-opus-20240229", - credential_fields: [ - { - key: "api_key", - label: "Anthropic API Key", - field_type: "password", - required: true, - }, - ], - }, - ]), - }; -}); - -const createQueryClient = () => - new QueryClient({ - defaultOptions: { - queries: { - retry: false, - gcTime: 0, - }, - }, - }); - -const mockUploadProps = { - beforeUpload: vi.fn(), - onChange: vi.fn(), -}; - -describe("AddCredentialModal", () => { - it("should render", () => { - const queryClient = createQueryClient(); - const onCancel = vi.fn(); - const onAddCredential = vi.fn(); - - render( - - - , - ); - - expect(screen.getByText("Add New Credential")).toBeInTheDocument(); - expect(screen.getByLabelText("Credential Name:")).toBeInTheDocument(); - expect(screen.getByLabelText("Provider:")).toBeInTheDocument(); - }); - - it("should show the correct provider fields", async () => { - const queryClient = createQueryClient(); - const onCancel = vi.fn(); - const onAddCredential = vi.fn(); - - render( - - - , - ); - - await waitFor(() => { - expect(screen.getByLabelText("OpenAI API Key")).toBeInTheDocument(); - expect(screen.getByPlaceholderText("https://api.openai.com/v1")).toBeInTheDocument(); - }); - }); -}); diff --git a/ui/litellm-dashboard/src/components/model_add/CredentialModal.test.tsx b/ui/litellm-dashboard/src/components/model_add/CredentialModal.test.tsx new file mode 100644 index 00000000000..6804d0cba92 --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_add/CredentialModal.test.tsx @@ -0,0 +1,140 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen, waitFor } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import { Providers } from "../provider_info_helpers"; +import { CredentialItem } from "../networking"; +import CredentialModal from "./CredentialModal"; + +vi.mock("../networking", async () => { + const actual = await vi.importActual("../networking"); + return { + ...actual, + getProviderCreateMetadata: vi.fn().mockResolvedValue([ + { + provider: "OpenAI", + provider_display_name: Providers.OpenAI, + litellm_provider: "openai", + default_model_placeholder: "gpt-3.5-turbo", + credential_fields: [ + { + key: "api_key", + label: "OpenAI API Key", + field_type: "password", + required: true, + }, + { + key: "api_base", + label: "API Base", + field_type: "text", + placeholder: "https://api.openai.com/v1", + }, + ], + }, + { + provider: "Anthropic", + provider_display_name: Providers.Anthropic, + litellm_provider: "anthropic", + default_model_placeholder: "claude-3-opus-20240229", + credential_fields: [ + { + key: "api_key", + label: "Anthropic API Key", + field_type: "password", + required: true, + }, + ], + }, + ]), + }; +}); + +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + +const mockUploadProps = { + beforeUpload: vi.fn(), + onChange: vi.fn(), +}; + +const mockCredential: CredentialItem = { + credential_name: "test-credential", + credential_values: { + api_key: "test-api-key", + api_base: "https://api.test.com", + }, + credential_info: { + custom_llm_provider: Providers.OpenAI, + }, +}; + +const renderModal = (props: Partial> = {}) => + render( + + + , + ); + +describe("CredentialModal", () => { + describe("add mode", () => { + it("renders the add title and an editable credential name", () => { + renderModal({ mode: "add" }); + + expect(screen.getByText("Add New Credential")).toBeInTheDocument(); + expect(screen.getByText("Add Credential")).toBeInTheDocument(); + const nameInput = screen.getByLabelText("Credential Name:") as HTMLInputElement; + expect(nameInput.value).toBe(""); + expect(nameInput.disabled).toBe(false); + }); + + it("shows provider-specific fields for the selected provider", async () => { + renderModal({ mode: "add" }); + + await waitFor(() => { + expect(screen.getByLabelText("OpenAI API Key")).toBeInTheDocument(); + expect(screen.getByPlaceholderText("https://api.openai.com/v1")).toBeInTheDocument(); + }); + }); + }); + + describe("edit mode", () => { + it("renders the edit title and update button", () => { + renderModal({ mode: "edit", existingCredential: mockCredential }); + + expect(screen.getByText("Edit Credential")).toBeInTheDocument(); + expect(screen.getByText("Update Credential")).toBeInTheDocument(); + }); + + it("prefills the credential name and disables it", async () => { + renderModal({ mode: "edit", existingCredential: mockCredential }); + + await waitFor(() => { + const nameInput = screen.getByLabelText("Credential Name:") as HTMLInputElement; + expect(nameInput.value).toBe("test-credential"); + expect(nameInput.disabled).toBe(true); + }); + }); + + it("disables the name from the mode, not the credential's name value", () => { + renderModal({ + mode: "edit", + existingCredential: { ...mockCredential, credential_name: "" }, + }); + + expect((screen.getByLabelText("Credential Name:") as HTMLInputElement).disabled).toBe(true); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx similarity index 71% rename from ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx rename to ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx index b86a379d3d1..c92a4a90578 100644 --- a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/model_add/CredentialModal.tsx @@ -1,23 +1,47 @@ import { TextInput } from "@tremor/react"; import { Select as AntdSelect, Button, Form, Modal, Tooltip, Typography } from "antd"; import type { UploadProps } from "antd/es/upload"; -import React, { useState } from "react"; +import { useState } from "react"; import ProviderSpecificFields from "../add_model/provider_specific_fields"; +import { CredentialItem } from "../networking"; import { Providers, providerLogoMap } from "../provider_info_helpers"; import { resolveLogoSrc } from "@/lib/assetPaths"; import { resetCredentialFormOnProviderChange } from "./credential_form_helpers"; + const { Link } = Typography; -interface AddCredentialsModalProps { +interface CredentialModalProps { open: boolean; onCancel: () => void; - onAddCredential: (values: any) => void; + onSubmit: (values: any) => void; uploadProps: UploadProps; + mode: "add" | "edit"; + existingCredential?: CredentialItem | null; } -const AddCredentialsModal: React.FC = ({ open, onCancel, onAddCredential, uploadProps }) => { +export default function CredentialModal({ + open, + onCancel, + onSubmit, + uploadProps, + mode, + existingCredential = null, +}: CredentialModalProps) { + const isEdit = mode === "edit"; const [form] = Form.useForm(); - const [selectedProvider, setSelectedProvider] = useState(Providers.OpenAI); + const [selectedProvider, setSelectedProvider] = useState( + (existingCredential?.credential_info.custom_llm_provider as Providers) ?? Providers.OpenAI, + ); + + const initialValues = existingCredential + ? { + credential_name: existingCredential.credential_name, + custom_llm_provider: existingCredential.credential_info.custom_llm_provider, + ...Object.fromEntries( + Object.entries(existingCredential.credential_values || {}).map(([key, value]) => [key, value ?? null]), + ), + } + : undefined; const handleSubmit = (values: any) => { const filteredValues = Object.entries(values).reduce((acc, [key, value]) => { @@ -26,32 +50,33 @@ const AddCredentialsModal: React.FC = ({ open, onCance } return acc; }, {} as any); - onAddCredential(filteredValues); + onSubmit(filteredValues); + form.resetFields(); + }; + + const closeAndReset = () => { + onCancel(); form.resetFields(); }; return ( { - onCancel(); - form.resetFields(); - }} + onCancel={closeAndReset} footer={null} width={600} + destroyOnHidden={isEdit} > -
- {/* Credential Name */} + - + - {/* Provider Selection */} = ({ open, onCance - {/* Modal Footer */}
Need Help?
- - +
); -}; - -export default AddCredentialsModal; +} diff --git a/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.test.tsx b/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.test.tsx deleted file mode 100644 index def3b4f6cd7..00000000000 --- a/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.test.tsx +++ /dev/null @@ -1,123 +0,0 @@ -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render, screen, waitFor } from "@testing-library/react"; -import { describe, expect, it, vi } from "vitest"; -import { Providers } from "../provider_info_helpers"; -import { CredentialItem } from "../networking"; -import EditCredentialModal from "./EditCredentialModal"; - -vi.mock("../networking", async () => { - const actual = await vi.importActual("../networking"); - return { - ...actual, - getProviderCreateMetadata: vi.fn().mockResolvedValue([ - { - provider: "OpenAI", - provider_display_name: Providers.OpenAI, - litellm_provider: "openai", - default_model_placeholder: "gpt-3.5-turbo", - credential_fields: [ - { - key: "api_key", - label: "OpenAI API Key", - field_type: "password", - required: true, - }, - { - key: "api_base", - label: "API Base", - field_type: "text", - placeholder: "https://api.openai.com/v1", - }, - ], - }, - { - provider: "Anthropic", - provider_display_name: Providers.Anthropic, - litellm_provider: "anthropic", - default_model_placeholder: "claude-3-opus-20240229", - credential_fields: [ - { - key: "api_key", - label: "Anthropic API Key", - field_type: "password", - required: true, - }, - ], - }, - ]), - }; -}); - -const createQueryClient = () => - new QueryClient({ - defaultOptions: { - queries: { - retry: false, - gcTime: 0, - }, - }, - }); - -const mockUploadProps = { - beforeUpload: vi.fn(), - onChange: vi.fn(), -}; - -const mockCredential: CredentialItem = { - credential_name: "test-credential", - credential_values: { - api_key: "test-api-key", - api_base: "https://api.test.com", - }, - credential_info: { - custom_llm_provider: Providers.OpenAI, - }, -}; - -describe("EditCredentialModal", () => { - it("should render", () => { - const queryClient = createQueryClient(); - const onCancel = vi.fn(); - const onUpdateCredential = vi.fn(); - - render( - - - , - ); - - expect(screen.getByText("Edit Credential")).toBeInTheDocument(); - expect(screen.getByLabelText("Credential Name:")).toBeInTheDocument(); - expect(screen.getByLabelText("Provider:")).toBeInTheDocument(); - }); - - it("should render initial values", async () => { - const queryClient = createQueryClient(); - const onCancel = vi.fn(); - const onUpdateCredential = vi.fn(); - - render( - - - , - ); - - await waitFor(() => { - const credentialNameInput = screen.getByLabelText("Credential Name:") as HTMLInputElement; - expect(credentialNameInput.value).toBe("test-credential"); - expect(credentialNameInput.disabled).toBe(true); - }); - }); -}); diff --git a/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx deleted file mode 100644 index d087edc1069..00000000000 --- a/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx +++ /dev/null @@ -1,150 +0,0 @@ -import { TextInput } from "@tremor/react"; -import { Select as AntdSelect, Button, Form, Modal, Tooltip, Typography } from "antd"; -import type { UploadProps } from "antd/es/upload"; -import { useEffect, useState } from "react"; -import ProviderSpecificFields from "../add_model/provider_specific_fields"; -import { CredentialItem } from "../networking"; -import { Providers, providerLogoMap } from "../provider_info_helpers"; -import { resolveLogoSrc } from "@/lib/assetPaths"; -import { resetCredentialFormOnProviderChange } from "./credential_form_helpers"; -const { Link } = Typography; - -interface EditCredentialsModalProps { - open: boolean; - onCancel: () => void; - onUpdateCredential: (values: any) => void; - uploadProps: UploadProps; - existingCredential: CredentialItem | null; -} - -export default function EditCredentialsModal({ - open, - onCancel, - onUpdateCredential, - uploadProps, - existingCredential, -}: EditCredentialsModalProps) { - const [form] = Form.useForm(); - const [selectedProvider, setSelectedProvider] = useState(Providers.Anthropic); - - const handleSubmit = (values: any) => { - const filteredValues = Object.entries(values).reduce((acc, [key, value]) => { - if (value !== "" && value !== undefined && value !== null) { - acc[key] = value; - } - return acc; - }, {} as any); - onUpdateCredential(filteredValues); - form.resetFields(); - }; - - useEffect(() => { - if (existingCredential) { - // Spread all credential_values dynamically, converting undefined/null to null for form compatibility - const credentialValues = Object.entries(existingCredential.credential_values || {}).reduce( - (acc, [key, value]) => { - acc[key] = value ?? null; - return acc; - }, - {} as Record, - ); - - form.setFieldsValue({ - credential_name: existingCredential.credential_name, - custom_llm_provider: existingCredential.credential_info.custom_llm_provider, - ...credentialValues, - }); - setSelectedProvider(existingCredential.credential_info.custom_llm_provider as Providers); - } - }, [existingCredential]); - - return ( - { - onCancel(); - form.resetFields(); - }} - footer={null} - width={600} - destroyOnHidden={true} - > -
- {/* Credential Name */} - - - - - {/* Provider Selection */} - - { - resetCredentialFormOnProviderChange(form, value as Providers, setSelectedProvider); - }} - > - {Object.entries(Providers).map(([providerEnum, providerDisplayName]) => ( - -
- {`${providerEnum} { - const target = e.target as HTMLImageElement; - const parent = target.parentElement; - if (parent) { - const fallbackDiv = document.createElement("div"); - fallbackDiv.className = - "w-5 h-5 rounded-full bg-gray-200 flex items-center justify-center text-xs"; - fallbackDiv.textContent = providerDisplayName.charAt(0); - parent.replaceChild(fallbackDiv, target); - } - }} - /> - {providerDisplayName} -
- - ))} - - - - - - {/* Modal Footer */} -
- - Need Help? - - -
- - -
-
- - - ); -} diff --git a/ui/litellm-dashboard/src/components/model_add/credentials.tsx b/ui/litellm-dashboard/src/components/model_add/credentials.tsx index 82320b7ff8d..9289888c1ed 100644 --- a/ui/litellm-dashboard/src/components/model_add/credentials.tsx +++ b/ui/litellm-dashboard/src/components/model_add/credentials.tsx @@ -22,8 +22,7 @@ import { UploadProps } from "antd/es/upload"; import { useState } from "react"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; import NotificationsManager from "../molecules/notifications_manager"; -import AddCredentialsTab from "./AddCredentialModal"; -import EditCredentialsModal from "./EditCredentialModal"; +import CredentialModal from "./CredentialModal"; import { useCredentials } from "@/app/(dashboard)/hooks/credentials/useCredentials"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { isProxyAdminRole } from "@/utils/roles"; @@ -201,18 +200,20 @@ const CredentialsPanel: React.FC = ({ uploadProps }) => { {isAddModalOpen && ( - setIsAddModalOpen(false)} uploadProps={uploadProps} /> )} {isUpdateModalOpen && ( - setIsUpdateModalOpen(false)} /> From e18966625d63847a8c2e476767734bb711a2b88b Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 18 Jul 2026 11:36:25 -0700 Subject: [PATCH 13/18] feat(mcp): add ID-JAG (identity assertion authorization grant) support for MCP egress (#31516) * feat(mcp): add ID-JAG egress auth as a v2 outbound-credentials arm Adds the oauth2_id_jag MCP egress auth mode (draft-ietf-oauth-identity-assertion-authz-grant, shipped by Okta as "AI agent token exchange") as a first-class arm of the v2 outbound_credentials resolver rather than a standalone v1 handler. ID-JAG is a two-leg flow: an RFC 8693 token exchange swaps the caller's id_token for an ID-JAG assertion at the IdP org authorization server, then an RFC 7523 jwt-bearer grant presents that assertion to the MCP's resource authorization server for the access token used to call the upstream. The gateway authenticates to both endpoints with a private-key JWT client_assertion, falling back to client_secret when no key is configured. The mode is modeled as IdJagConfig in the AuthConfig discriminated union, with client auth as a ClientAuth tagged union (private_key_jwt or client_secret) so required fields are enforced at construction and illegal states are unrepresentable. A new token_endpoint collaborator performs the authenticated OAuth token-endpoint call and caches the result with per-key single-flight; the resolver's _id_jag arm runs the two legs and returns an httpx.Auth or a typed CredError. A missing caller identity token fails closed (precondition_required), so an ID-JAG server never falls back to a static credential. The v1->v2 adapter maps oauth2_id_jag servers onto IdJagConfig and the existing live v2 path resolves them, so no standalone handler, has_id_jag_config flag, or resolve_mcp_auth precedence branch is needed. The ID-JAG client_private_key is encrypted at rest alongside client_secret. * fix(mcp): sort token_endpoint imports to satisfy the I001 budget gate * fix(mcp): give token_endpoint pyright suppressions reasons for the LIT004 budget The freshly-merged base ratcheted the LIT004 ceiling down, so the six unexplained pyright suppressions in token_endpoint.py went over budget. Annotate each with why the boundary is untyped (litellm http handler and InMemoryCache are untyped; response.json() is validated by _TokenEndpointResponse in fetch) so the gate counts them as explained. * fix(mcp): enforce ID-JAG exchange over caller auth overrides and redact token endpoint from client errors For oauth2_id_jag servers the v2 resolver mints the upstream assertion from the caller's identity token; a caller-supplied x-mcp-auth / x-mcp--authorization override or a conflicting injected Authorization must not disable that exchange and forward an arbitrary bearer, so IdJagConfig now joins authorization_code and token_exchange as a resolver-owned mode that keeps the v2 spec and ignores the override. The token endpoint error branches previously returned the configured endpoint URL in the client-visible 503 detail. The endpoint now stays in server-side logs and clients get a generic token-exchange failure. * fix(mcp): bind the ID-JAG token cache to the exchange config and map token endpoint network errors to typed CredErrors * fix(mcp): fail closed when an oauth2_id_jag server is half-configured instead of deferring to v1 static credentials * fix(mcp): evict the cached ID-JAG bearer on an upstream 401 so the retry re-exchanges * fix(mcp): map an unsignable client assertion to a typed misconfigured error instead of an unhandled 500 * fix(mcp): redact credential fields from the server-registry debug dump --- litellm/proxy/_experimental/mcp_server/db.py | 7 + .../mcp_server/mcp_server_manager.py | 107 ++++- .../outbound_credentials/__init__.py | 8 + .../outbound_credentials/adapter.py | 61 +++ .../outbound_credentials/resolver.py | 116 ++++- .../outbound_credentials/token_endpoint.py | 225 ++++++++++ .../mcp_server/outbound_credentials/types.py | 45 ++ litellm/types/mcp.py | 27 ++ .../types/mcp_server/mcp_server_manager.py | 9 + .../outbound_credentials/test_adapter.py | 78 ++++ .../outbound_credentials/test_resolver.py | 212 +++++++++ .../test_token_endpoint.py | 408 ++++++++++++++++++ .../outbound_credentials/test_types.py | 81 ++++ .../mcp_server/test_db_credentials.py | 25 ++ .../mcp_server/test_mcp_server_manager.py | 192 +++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 16 files changed, 1582 insertions(+), 21 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 7129582ff2a..9fe970f7fa9 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -375,6 +375,12 @@ def encrypt_credentials(credentials: MCPCredentials, encryption_key: Optional[st value=client_secret, new_encryption_key=encryption_key, ) + client_private_key = credentials.get("client_private_key") + if client_private_key is not None: + credentials["client_private_key"] = encrypt_value_helper( + value=client_private_key, + new_encryption_key=encryption_key, + ) # AWS SigV4 credential fields aws_access_key_id = credentials.get("aws_access_key_id") if aws_access_key_id is not None: @@ -406,6 +412,7 @@ def decrypt_credentials( "auth_value", "client_id", "client_secret", + "client_private_key", "aws_access_key_id", "aws_secret_access_key", "aws_session_token", diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index ed6dde23d9e..1ba608b9510 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -93,6 +93,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_ ) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( AuthorizationCodeConfig, + CredError, + IdJagConfig, PassthroughConfig, ServerSpec, TokenExchangeConfig, @@ -621,6 +623,47 @@ def _consumes_caller_authorization(server: MCPServer) -> bool: ) +_REGISTRY_DUMP_SECRET_FIELDS = frozenset( + {"authentication_token", "client_secret", "client_private_key", "aws_secret_access_key", "aws_session_token"} +) + + +def _redacted_registry_dump(servers: dict[str, MCPServer]) -> dict[str, dict[str, str]]: + """A JSON-safe view of the server registry with credential fields masked, for debug logging. + + The registry holds long-lived secrets as plain strings (the static token, OAuth client secret, + the ID-JAG signing key, AWS keys); dumping them verbatim hands the gateway's client identity to + anyone who can read debug logs. + """ + dumps: dict[str, dict[str, object]] = {server_id: server.model_dump() for server_id, server in servers.items()} + return { + server_id: { + field: ("**REDACTED**" if field in _REGISTRY_DUMP_SECRET_FIELDS and value is not None else str(value)) + for field, value in dump.items() + } + for server_id, dump in dumps.items() + } + + +def _to_server_spec_fail_closed(server: MCPServer) -> Optional[ServerSpec]: + """`to_server_spec`, except a half-configured `oauth2_id_jag` server refuses instead of deferring. + + ID-JAG has no v1 arm, so deferring to v1 would let `resolve_mcp_auth` honor a caller x-mcp-* + override or fall through to the static `authentication_token`, both of which bypass the per-user + identity assertion the mode promises. That is an operator misconfiguration, not a fallback. + """ + spec = to_server_spec(server) + if spec is None and server.auth_type == MCPAuth.oauth2_id_jag: + raise_public( + CredError.of_misconfigured( + "oauth2_id_jag requires token_exchange_endpoint, id_jag_resource_token_endpoint, " + "client_id, and a client_secret or client_private_key; refusing to fall back to " + "a static credential." + ) + ) + return spec + + def _caller_authorization_fans_out( server: MCPServer, scope_servers: Optional[list[MCPServer]], @@ -1326,6 +1369,12 @@ class MCPServerManager: "subject_token_type", DEFAULT_SUBJECT_TOKEN_TYPE, ), + # ID-JAG fields + id_jag_resource_token_endpoint=server_config.get("id_jag_resource_token_endpoint", None), + id_jag_resource=server_config.get("id_jag_resource", None), + client_private_key=server_config.get("client_private_key", None), + client_private_key_id=server_config.get("client_private_key_id", None), + client_assertion_signing_alg=server_config.get("client_assertion_signing_alg", "RS256"), token_exchange_profile=server_config.get("token_exchange_profile", "rfc8693"), allow_sampling=bool(server_config.get("allow_sampling", False)), allow_elicitation=bool(server_config.get("allow_elicitation", False)), @@ -1346,7 +1395,9 @@ class MCPServerManager: base_url=server_config.get("url", ""), ) - verbose_logger.debug(f"Loaded MCP Servers: {json.dumps(self.config_mcp_servers, indent=4, default=str)}") + verbose_logger.debug( + f"Loaded MCP Servers: {json.dumps(_redacted_registry_dump(self.config_mcp_servers), indent=4)}" + ) await self._hydrate_config_servers_dcr_clients() @@ -1797,6 +1848,21 @@ class MCPServerManager: subject_token_type=mcp_server.subject_token_type or (credentials_dict.get("subject_token_type") if credentials_dict else None) or DEFAULT_SUBJECT_TOKEN_TYPE, + # ID-JAG fields — read from credentials JSON blob + id_jag_resource_token_endpoint=( + credentials_dict.get("id_jag_resource_token_endpoint") if credentials_dict else None + ), + id_jag_resource=(credentials_dict.get("id_jag_resource") if credentials_dict else None), + client_private_key=self._decrypt_credential_field( + credentials_dict.get("client_private_key") if credentials_dict else None, + "client_private_key", + credentials_are_encrypted, + ), + client_private_key_id=(credentials_dict.get("client_private_key_id") if credentials_dict else None), + client_assertion_signing_alg=( + credentials_dict.get("client_assertion_signing_alg") if credentials_dict else None + ) + or "RS256", token_exchange_profile=mcp_server.token_exchange_profile or (credentials_dict.get("token_exchange_profile") if credentials_dict else None) or "rfc8693", @@ -2673,9 +2739,10 @@ class MCPServerManager: ) if not conflicts: return auth, extra_headers - if isinstance(spec.config, (TokenExchangeConfig, AuthorizationCodeConfig)): + if isinstance(spec.config, (TokenExchangeConfig, AuthorizationCodeConfig, IdJagConfig)): # The resolver owns the per-user credential here (token_exchange's exchanged - # token, authorization_code's stored token). It is authoritative: a guardrail such + # token, authorization_code's stored token, id_jag's minted assertion). It is + # authoritative: a guardrail such # as MCPJWTSigner, static_headers, or any other injected Authorization must NOT # shadow it (otherwise the upstream gets e.g. the signer's JWT instead of the # exchanged token and rejects it). Drop the conflicting header so the resolved @@ -2766,20 +2833,23 @@ class MCPServerManager: Configured MCP client instance. """ transport = server.transport or MCPTransport.sse - spec = None if transport == MCPTransport.stdio else to_server_spec(server) + spec = None if transport == MCPTransport.stdio else _to_server_spec_fail_closed(server) provider = cred_provider or self._cred_provider # A caller-supplied per-request override (mcp_auth_header / x-mcp-*) defers to the v1 path # so it wins - except for the modes the v2 resolver owns per-caller (authorization_code's - # stored token, token_exchange's RFC 8693 minted token, and the passthrough modes' - # forwarded caller token). A caller must not be able to substitute another user's stored - # credential, nor silently disable the OBO exchange and forward an arbitrary bearer - # upstream, so we keep the v2 spec and ignore the override for these; the REST tools - # preview supplies its not-yet-persisted token through the resolver (cred_provider), - # never this path. + # stored token, token_exchange's RFC 8693 minted token, id_jag's minted assertion, and the + # passthrough modes' forwarded caller token). A caller must not be able to substitute another + # user's stored credential, nor silently disable the OBO / ID-JAG exchange and forward an + # arbitrary bearer upstream, so we keep the v2 spec and ignore the override for these; the + # REST tools preview supplies its not-yet-persisted token through the resolver + # (cred_provider), never this path. if ( spec is not None and mcp_auth_header - and not isinstance(spec.config, (AuthorizationCodeConfig, PassthroughConfig, TokenExchangeConfig)) + and not isinstance( + spec.config, + (AuthorizationCodeConfig, IdJagConfig, PassthroughConfig, TokenExchangeConfig), + ) ): spec = None auth_value = ( @@ -4308,10 +4378,13 @@ class MCPServerManager: if server_auth_header is None: server_auth_header = mcp_auth_header - # Extract subject token for OAuth2 Token Exchange (OBO) flow + # Extract subject token for OAuth2 Token Exchange (OBO) and ID-JAG flows subject_token: Optional[str] = None extra_headers: Optional[dict[str, str]] = None - if mcp_server.auth_type == MCPAuth.oauth2_token_exchange: + if mcp_server.auth_type in ( + MCPAuth.oauth2_token_exchange, + MCPAuth.oauth2_id_jag, + ): subject_token = self._extract_bearer_token(oauth2_headers, raw_headers) elif mcp_server.auth_type == MCPAuth.oauth2: if mcp_server.has_client_credentials: @@ -4413,10 +4486,10 @@ class MCPServerManager: arguments=arguments, ) - if mcp_server.auth_type == MCPAuth.oauth2_token_exchange and subject_token: - # OBO: the exchanged token may have been revoked/rotated upstream since it was cached, so - # an upstream 401 gets one re-mint + retry. Gated to this mode; all others keep the plain - # single call below. + if mcp_server.auth_type in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag) and subject_token: + # OBO / ID-JAG: the exchanged token may have been revoked/rotated upstream since it was + # cached, so an upstream 401 gets one invalidate + re-mint + retry. Gated to these modes; + # all others keep the plain single call below. async def _obo_call_tool_limited(): async with self._limit_outbound_concurrency(mcp_server): return await self._obo_call_tool_with_retry( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py index 73166a45d6e..2bdb8770e4e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py @@ -31,10 +31,14 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( AwsCredentialSource, AwsSigV4Config, Byok, + ClientAuth, ClientCredentialsConfig, + ClientSecretAuth, CredError, + IdJagConfig, NoneConfig, PassthroughConfig, + PrivateKeyJwtAuth, ServerSpec, SharedKey, StaticKeys, @@ -59,6 +63,10 @@ __all__ = [ "AuthorizationCodeConfig", "ClientCredentialsConfig", "TokenExchangeConfig", + "IdJagConfig", + "ClientAuth", + "PrivateKeyJwtAuth", + "ClientSecretAuth", "ApiKeyConfig", "ApiKeySource", "SharedKey", diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index e87e8081ced..6631e38f524 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -21,9 +21,13 @@ from typing_extensions import assert_never from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, + ClientAuth, + ClientSecretAuth, CredError, + IdJagConfig, NoneConfig, PassthroughConfig, + PrivateKeyJwtAuth, ServerSpec, SharedKey, Subject, @@ -35,6 +39,9 @@ if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer +_TOKEN_EXCHANGE_SUBJECT_TOKEN_DEFAULT = "urn:ietf:params:oauth:token-type:access_token" +_ID_JAG_SUBJECT_TOKEN_DEFAULT = "urn:ietf:params:oauth:token-type:id_token" + def to_subject(user_api_key_auth: Optional[UserAPIKeyAuth], subject_token: Optional[str]) -> Subject: """Map v1's authenticated principal onto the resolver's Subject. @@ -96,6 +103,8 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: ) # client_credentials (M2M) and delegate/passthrough oauth2 stay on v1 return None + case MCPAuth.oauth2_id_jag: + return _id_jag_spec(server, resource) case MCPAuth.true_passthrough | MCPAuth.oauth_delegate: return ServerSpec(server_id=server.server_id, resource=resource, config=PassthroughConfig()) case MCPAuth.oauth2_token_exchange: @@ -167,6 +176,58 @@ def _shared_key_spec( ) +def _id_jag_spec(server: MCPServer, resource: str) -> Optional[ServerSpec]: + """Build an ID-JAG spec from the v1 server's raw fields, or defer (None) if half-configured. + + The enum already routes here, but a server missing an endpoint, ``client_id``, or any client-auth + secret would make ``IdJagConfig`` raise at construction; returning None instead defers to v1 so a + partially configured server does not 500. ``token_exchange_endpoint`` is leg 1 (the IdP org AS); + leg 2 is ``id_jag_resource_token_endpoint`` (the upstream resource AS). + """ + org_token_endpoint = server.token_exchange_endpoint + resource_token_endpoint = server.id_jag_resource_token_endpoint + client_id = server.client_id + client_auth = _id_jag_client_auth(server) + if not org_token_endpoint or not resource_token_endpoint or not client_id or client_auth is None: + return None + return ServerSpec( + server_id=server.server_id, + resource=resource, + config=IdJagConfig( + org_token_endpoint=org_token_endpoint, + resource_token_endpoint=resource_token_endpoint, + client_id=client_id, + client_auth=client_auth, + subject_token_type=_id_jag_subject_token_type(server), + audience=server.audience, + resource=server.id_jag_resource, + scopes=tuple(server.scopes or ()), + ), + ) + + +def _id_jag_client_auth(server: MCPServer) -> Optional[ClientAuth]: + """Private-key JWT when a key is configured, else client_secret, else None (defer to v1).""" + if server.client_private_key: + return PrivateKeyJwtAuth( + private_key=SecretStr(server.client_private_key), + key_id=server.client_private_key_id, + signing_alg=server.client_assertion_signing_alg, + ) + if server.client_secret: + return ClientSecretAuth(client_secret=SecretStr(server.client_secret)) + return None + + +def _id_jag_subject_token_type(server: MCPServer) -> str: + """ID-JAG asserts the user's id_token, so the token-exchange access_token default maps to id_token; + an explicitly configured value (e.g. a SAML2 assertion type) is honored verbatim.""" + configured = server.subject_token_type + if configured and configured != _TOKEN_EXCHANGE_SUBJECT_TOKEN_DEFAULT: + return configured + return _ID_JAG_SUBJECT_TOKEN_DEFAULT + + def raise_public(error: CredError) -> NoReturn: """Map a resolver CredError onto the proxy's public HTTP contract. The one edge that raises.""" match error.tag: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index ecfd471190c..7e5c073870a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -16,6 +16,8 @@ follow-up PR with their seam. Pure v2: no imports from v1. from __future__ import annotations +import hashlib + import httpx from typing_extensions import assert_never @@ -33,6 +35,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( Ok, Result, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( + ExchangedToken, + ExchangedTokenCache, + TokenEndpointClient, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( TokenExchanger, ) @@ -42,16 +49,24 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( AuthSpecKind, AwsSigV4Config, Byok, + ClientAuth, ClientCredentialsConfig, + ClientSecretAuth, CredError, + IdJagConfig, NoneConfig, PassthroughConfig, + PrivateKeyJwtAuth, ServerSpec, SharedKey, Subject, TokenExchangeConfig, ) +_TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" +_JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer" +_ID_JAG_REQUESTED_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:id-jag" + class _NullOAuthTokenStore: """Fail-closed default: with no token store wired, every user reads as not authorized.""" @@ -87,9 +102,13 @@ class UpstreamCredentialProvider: self, oauth_token_store: OAuthTokenStore | None = None, token_exchanger: TokenExchanger | None = None, + token_endpoint: TokenEndpointClient | None = None, + exchanged_tokens: ExchangedTokenCache | None = None, ) -> None: self._oauth_token_store: OAuthTokenStore = oauth_token_store or _NullOAuthTokenStore() self._token_exchanger: TokenExchanger = token_exchanger or _NullTokenExchanger() + self._token_endpoint: TokenEndpointClient = token_endpoint or TokenEndpointClient() + self._exchanged_tokens: ExchangedTokenCache = exchanged_tokens or ExchangedTokenCache() async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: @@ -103,6 +122,8 @@ class UpstreamCredentialProvider: return _not_implemented(AuthSpecKind.client_credentials) case TokenExchangeConfig() as config: return await self._token_exchange(subject, server, config) + case IdJagConfig() as config: + return await self._id_jag(subject, server, config) case AuthorizationCodeConfig(): return await self._authorization_code(subject, server) case AwsSigV4Config(): @@ -141,6 +162,53 @@ class UpstreamCredentialProvider: return Error(CredError.of_not_implemented("api_key BYOK source not implemented yet")) assert_never(config.key_source) + async def _id_jag(self, subject: Subject, server: ServerSpec, config: IdJagConfig) -> Result[httpx.Auth, CredError]: + if subject.inbound_token is None: + return Error( + CredError.of_precondition_required( + "ID-JAG requires a caller identity token; it asserts the calling " + "user's identity upstream and cannot use a static credential." + ) + ) + token = subject.inbound_token.get_secret_value() + cache_key = _id_jag_cache_key(token, server.server_id, config) + + async def _exchange() -> Result[ExchangedToken, CredError]: + leg1_params = { + "grant_type": _TOKEN_EXCHANGE_GRANT_TYPE, + "requested_token_type": _ID_JAG_REQUESTED_TOKEN_TYPE, + "subject_token": token, + "subject_token_type": config.subject_token_type, + **({"audience": config.audience} if config.audience else {}), + **({"resource": config.resource} if config.resource else {}), + **({"scope": " ".join(config.scopes)} if config.scopes else {}), + } + match await self._token_endpoint.fetch( + config.org_token_endpoint, + config.client_id, + leg1_params, + config.client_auth, + ): + case Error(err): + return Error(err) + case Ok(id_jag): + leg2_params = { + "grant_type": _JWT_BEARER_GRANT_TYPE, + "assertion": id_jag.access_token, + } + return await self._token_endpoint.fetch( + config.resource_token_endpoint, + config.client_id, + leg2_params, + config.client_auth, + ) + + match await self._exchanged_tokens.get_or_compute(cache_key, _exchange): + case Ok(access_token): + return Ok(StaticHeaderAuth(f"Bearer {access_token}")) + case Error(err): + return Error(err) + async def _authorization_code(self, subject: Subject, server: ServerSpec) -> Result[StaticHeaderAuth, CredError]: token = await self._authz_token(subject, server) if token is None: @@ -176,13 +244,19 @@ class UpstreamCredentialProvider: """Drop any cached credential the resolver owns for this `(subject, server)`. Used after an upstream rejects the injected credential, so the next resolve re-mints rather - than serving the same rejected token until TTL. Only `token_exchange` holds a re-mintable - cached credential here; other modes are a no-op. + than serving the same rejected token until TTL. `token_exchange` and `id_jag` hold a + re-mintable cached credential here; other modes are a no-op. """ - if isinstance(server.config, TokenExchangeConfig) and subject.inbound_token is not None: + if subject.inbound_token is None: + return + if isinstance(server.config, TokenExchangeConfig): await self._token_exchanger.invalidate( subject.inbound_token.get_secret_value(), server, server.config, tenant_id=subject.tenant_id ) + if isinstance(server.config, IdJagConfig): + self._exchanged_tokens.invalidate( + _id_jag_cache_key(subject.inbound_token.get_secret_value(), server.server_id, server.config) + ) async def _authz_token(self, subject: Subject, server: ServerSpec) -> OAuthToken | None: """The user's authorization_code token, or None when absent or the store is unreachable. @@ -196,5 +270,41 @@ class UpstreamCredentialProvider: return None +def _id_jag_cache_key(subject_token: str, server_id: str, config: IdJagConfig) -> str: + """Bind the cached leg-2 bearer to the caller token, the server, AND the config that minted it. + + Every exchange parameter derives from the config (endpoints, audience, resource, scopes, client + auth), so a server update that changes any of them must change the key; otherwise the old bearer, + authorized under the old policy, keeps being served until its TTL. Everything is hashed, so no + secret is held in the key. + """ + material = "\x00".join( + ( + subject_token, + server_id, + config.org_token_endpoint, + config.resource_token_endpoint, + config.client_id, + _client_auth_fingerprint(config.client_auth), + config.subject_token_type, + config.audience or "", + config.resource or "", + " ".join(config.scopes), + ) + ) + return hashlib.sha256(material.encode()).hexdigest() + + +def _client_auth_fingerprint(client_auth: ClientAuth) -> str: + match client_auth: + case PrivateKeyJwtAuth() as auth: + return "\x00".join( + ("private_key_jwt", auth.private_key.get_secret_value(), auth.key_id or "", auth.signing_alg) + ) + case ClientSecretAuth() as auth: + return "\x00".join(("client_secret", auth.client_secret.get_secret_value())) + assert_never(client_auth) + + def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: return Error(CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet")) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py new file mode 100644 index 00000000000..4bc5732ec0e --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py @@ -0,0 +1,225 @@ +"""An authenticated OAuth token-endpoint call plus a short-lived-token cache. + +`TokenEndpointClient.fetch` POSTs one grant to a token endpoint, authenticating the gateway as +an OAuth client via `client_auth` (RFC 7523 private-key JWT, or `client_secret_post`), and returns +the minted token or a typed `CredError`. `ExchangedTokenCache` memoizes the final token string per +opaque cache key with per-key single-flight, so concurrent callers share one round-trip and a hit +skips the endpoint entirely. + +Pure v2: no imports from the v1 MCP auth handlers. The multi-leg flows that compose these (ID-JAG, +and later token_exchange / client_credentials) live in the resolver arms; this collaborator owns +only the single authenticated call and the cache. +""" + +from __future__ import annotations + +import asyncio +import json +import time +import uuid +import weakref +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass + +import httpx +import jwt +from pydantic import BaseModel, ValidationError +from typing_extensions import assert_never + +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import ( + MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, + MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, + MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, +) +from litellm.exceptions import Timeout +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # litellm http handler is untyped +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientAuth, + ClientSecretAuth, + CredError, + PrivateKeyJwtAuth, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + +CLIENT_ASSERTION_TYPE = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer" +CLIENT_ASSERTION_LIFETIME_SECONDS = 60 + + +@dataclass(frozen=True, slots=True) +class ExchangedToken: + access_token: str + expires_in: int | None + + +class _TokenEndpointResponse(BaseModel): + access_token: str + expires_in: int | None = None + + +class TokenEndpointClient: + """One authenticated POST to an OAuth token endpoint, returning the minted token as a value.""" + + async def fetch( + self, + endpoint: str, + client_id: str, + grant_params: Mapping[str, str], + client_auth: ClientAuth, + ) -> Result[ExchangedToken, CredError]: + try: + data = {**grant_params, **_client_auth_params(endpoint, client_id, client_auth)} + except (ValueError, TypeError, NotImplementedError, jwt.PyJWTError): + verbose_proxy_logger.warning("MCP token endpoint %s: could not sign the client assertion", endpoint) + return Error( + CredError.of_misconfigured( + "token exchange failed: could not sign the client assertion; " + "check client_private_key and client_assertion_signing_alg" + ) + ) + try: + raw = await _post_form(endpoint, data) + except httpx.HTTPStatusError as exc: + verbose_proxy_logger.warning( + "MCP token endpoint %s failed with status %s", endpoint, exc.response.status_code + ) + return Error( + CredError.of_upstream_unavailable(f"token exchange failed with status {exc.response.status_code}") + ) + except (httpx.RequestError, Timeout) as exc: + verbose_proxy_logger.warning("MCP token endpoint %s unreachable: %s", endpoint, type(exc).__name__) + return Error( + CredError.of_upstream_unavailable( + f"token exchange failed: token endpoint unreachable ({type(exc).__name__})" + ) + ) + except json.JSONDecodeError: + verbose_proxy_logger.warning("MCP token endpoint %s returned a non-JSON response", endpoint) + return Error( + CredError.of_upstream_unavailable("token exchange failed: token endpoint returned a non-JSON response") + ) + if raw is None: + verbose_proxy_logger.warning("MCP token endpoint %s returned no response", endpoint) + return Error(CredError.of_upstream_unavailable("token exchange failed: no response from token endpoint")) + try: + parsed = _TokenEndpointResponse.model_validate(raw) + except ValidationError: + verbose_proxy_logger.warning("MCP token endpoint %s response missing access_token", endpoint) + return Error( + CredError.of_upstream_unavailable("token exchange failed: token endpoint response missing access_token") + ) + return Ok(ExchangedToken(access_token=parsed.access_token, expires_in=parsed.expires_in)) + + +class ExchangedTokenCache: + """Memoizes the final token string per key, single-flighting concurrent misses on one lock.""" + + def __init__(self) -> None: + self._cache = InMemoryCache( + 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() + + async def get_or_compute( + self, + cache_key: str, + compute: Callable[[], Awaitable[Result[ExchangedToken, CredError]]], + ) -> Result[str, CredError]: + cached = self._get(cache_key) + if cached is not None: + return Ok(cached) + async with self._lock(cache_key): + cached = self._get(cache_key) + if cached is not None: + return Ok(cached) + match await compute(): + case Ok(token): + self._cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + cache_key, + token.access_token, + ttl=_cache_ttl_seconds(token.expires_in), + ) + 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).""" + self._cache.delete_cache(cache_key) # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + + def _get(self, cache_key: str) -> str | None: + value = self._cache.get_cache(cache_key) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # InMemoryCache is untyped; narrowed by isinstance below + return value if isinstance(value, str) 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 _cache_ttl_seconds(expires_in: int | None) -> int: + lifetime = expires_in if expires_in is not None else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL + return max( + lifetime - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, + ) + + +async def _post_form(endpoint: str, data: dict[str, str]) -> object | None: + # litellm's httpx handler and httpx.Response are only partially typed; the token endpoint + # returns a JSON object that `_TokenEndpointResponse` validates, so the untyped boundary is + # contained here. A non-2xx raises `httpx.HTTPStatusError`, an unreachable endpoint raises + # `httpx.RequestError` (or litellm's `Timeout`, which the handler substitutes for + # `httpx.TimeoutException`), and a non-JSON body raises `json.JSONDecodeError`; `fetch` maps + # each to a CredError. + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) # pyright: ignore[reportUnknownVariableType] # litellm http handler is untyped + response = await client.post(endpoint, data=data) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # litellm http handler is untyped + if response is None: + return None + response.raise_for_status() + return response.json() # pyright: ignore[reportAny] # untyped JSON; validated by _TokenEndpointResponse in fetch + + +def _client_auth_params(endpoint: str, client_id: str, client_auth: ClientAuth) -> dict[str, str]: + match client_auth: + case PrivateKeyJwtAuth() as auth: + return { + "client_id": client_id, + "client_assertion_type": CLIENT_ASSERTION_TYPE, + "client_assertion": _client_assertion(endpoint, client_id, auth), + } + case ClientSecretAuth() as auth: + return { + "client_id": client_id, + "client_secret": auth.client_secret.get_secret_value(), + } + assert_never(client_auth) + + +def _client_assertion(endpoint: str, client_id: str, auth: PrivateKeyJwtAuth) -> str: + now = int(time.time()) + return jwt.encode( + { + "iss": client_id, + "sub": client_id, + "aud": endpoint, + "jti": uuid.uuid4().hex, + "iat": now, + "exp": now + CLIENT_ASSERTION_LIFETIME_SECONDS, + }, + auth.private_key.get_secret_value(), + algorithm=auth.signing_alg, + headers={"kid": auth.key_id} if auth.key_id else None, + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 7e04be4f045..64a20255ab2 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -56,6 +56,7 @@ class AuthSpecKind(str, Enum): authorization_code = "authorization_code" # per-user 3LO; gateway-stored token client_credentials = "client_credentials" # gateway service account (M2M) token_exchange = "token_exchange" # RFC 8693: token endpoint + subject_token (OBO) + id_jag = "id_jag" # draft-ietf-oauth-identity-assertion-authz-grant: two-leg exchange then jwt-bearer api_key = "api_key" # static header, any scheme (BYOK = per-user-seeded source) passthrough = "passthrough" # client forwards an upstream-audience token none = "none" # no upstream credential; resolve yields a no-op auth, never an error @@ -225,6 +226,49 @@ class TokenExchangeConfig(BaseModel): scopes: tuple[str, ...] = () +class PrivateKeyJwtAuth(BaseModel): + """RFC 7523 private-key-JWT client authentication: the gateway signs a `client_assertion`.""" + + model_config = ConfigDict(frozen=True) + source: Literal["private_key_jwt"] = "private_key_jwt" + private_key: SecretStr + key_id: str | None = None + signing_alg: str = "RS256" + + +class ClientSecretAuth(BaseModel): + """`client_secret_post` client authentication: the gateway posts `client_id` + `client_secret`.""" + + model_config = ConfigDict(frozen=True) + source: Literal["client_secret"] = "client_secret" + client_secret: SecretStr + + +ClientAuth = Annotated[PrivateKeyJwtAuth | ClientSecretAuth, Field(discriminator="source")] + + +class IdJagConfig(BaseModel): + """draft-ietf-oauth-identity-assertion-authz-grant (Okta "AI agent token exchange"). + + Two legs: leg 1 is an RFC 8693 token exchange at the IdP org AS (`org_token_endpoint`) that + swaps the caller's identity token for an ID-JAG assertion; leg 2 is an RFC 7523 jwt-bearer at + the upstream resource AS (`resource_token_endpoint`) that swaps the assertion for the access + token. The gateway authenticates to both endpoints as `client_id` via `client_auth`. Required + fields are enforced at construction so a half-configured server cannot reach the arm. + """ + + model_config = ConfigDict(frozen=True) + kind: Literal[AuthSpecKind.id_jag] = AuthSpecKind.id_jag + org_token_endpoint: str + resource_token_endpoint: str + client_id: str + client_auth: ClientAuth + subject_token_type: str = "urn:ietf:params:oauth:token-type:id_token" + audience: str | None = None + resource: str | None = None + scopes: tuple[str, ...] = () + + class SharedKey(BaseModel): """A fixed key configured on the server, identical for every caller.""" @@ -323,6 +367,7 @@ AuthConfig = Annotated[ AuthorizationCodeConfig | ClientCredentialsConfig | TokenExchangeConfig + | IdJagConfig | ApiKeyConfig | PassthroughConfig | NoneConfig diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index ac411ad9d9a..377ba669082 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -38,6 +38,7 @@ class MCPAuth(str, enum.Enum): aws_sigv4 = "aws_sigv4" token = "token" oauth2_token_exchange = "oauth2_token_exchange" + oauth2_id_jag = "oauth2_id_jag" true_passthrough = "true_passthrough" oauth_delegate = "oauth_delegate" @@ -62,6 +63,7 @@ MCPAuthType = Optional[ MCPAuth.aws_sigv4, MCPAuth.token, MCPAuth.oauth2_token_exchange, + MCPAuth.oauth2_id_jag, MCPAuth.true_passthrough, MCPAuth.oauth_delegate, ] @@ -159,6 +161,31 @@ class MCPCredentials(TypedDict, total=False): the top-level request field. """ + id_jag_resource_token_endpoint: Optional[str] + """ + Resource authorization server JWT-bearer (RFC 7523) endpoint for ID-JAG leg 2 + """ + + id_jag_resource: Optional[str] + """ + Optional RFC 8707 resource indicator sent on ID-JAG leg 1 + """ + + client_private_key: Optional[str] + """ + PEM private key used to sign the private-key-JWT client_assertion (RFC 7523) + """ + + client_private_key_id: Optional[str] + """ + Key id (kid) advertised in the client_assertion JWT header + """ + + client_assertion_signing_alg: Optional[str] + """ + Signing algorithm for the client_assertion JWT. Default: RS256 + """ + token_endpoint_auth_method: Optional[MCPTokenEndpointAuthMethod] """ How the gateway authenticates to the upstream token endpoint. "client_secret_basic" diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index d0d8cc4cb28..8ae974b19a6 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -87,6 +87,15 @@ class MCPServer(BaseModel): token_exchange_endpoint: Optional[str] = None audience: Optional[str] = None subject_token_type: str = DEFAULT_SUBJECT_TOKEN_TYPE + # ID-JAG fields (draft-ietf-oauth-identity-assertion-authz-grant). + # Leg 1 reuses token_exchange_endpoint (IdP org-AS), audience (resource-AS + # identifier), scopes, subject_token_type, client_id/client_secret. Leg 2 + # posts the ID-JAG assertion to id_jag_resource_token_endpoint. + id_jag_resource_token_endpoint: Optional[str] = None + id_jag_resource: Optional[str] = None + client_private_key: Optional[str] = None + client_private_key_id: Optional[str] = None + client_assertion_signing_alg: str = "RS256" # Wire dialect: "rfc8693" (standard token-exchange grant) or "entra_obo" (Microsoft Entra # On-Behalf-Of, the RFC 7523 jwt-bearer grant + requested_token_use extension) token_exchange_profile: str = "rfc8693" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 17960e917a4..707374e7061 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -21,9 +21,12 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, + ClientSecretAuth, CredError, + IdJagConfig, NoneConfig, PassthroughConfig, + PrivateKeyJwtAuth, SharedKey, TokenExchangeConfig, ) @@ -35,6 +38,21 @@ def _server(**kwargs) -> MCPServer: return MCPServer(server_id="s", name="n", transport=MCPTransport.http, **kwargs) +def _id_jag_server(**overrides) -> MCPServer: + defaults = dict( + auth_type=MCPAuth.oauth2_id_jag, + url="https://mcp.example.com/mcp", + client_id="litellm-client-id", + client_secret="litellm-client-secret", + token_exchange_endpoint="https://idp.example.com/token", + id_jag_resource_token_endpoint="https://mcp-as.example.com/token", + audience="api://mcp-server", + scopes=["mcp.read", "mcp.write"], + ) + defaults.update(overrides) + return _server(**defaults) + + def test_none_maps_to_none_config(): spec = to_server_spec(_server(auth_type=None)) assert spec is not None @@ -413,3 +431,63 @@ def test_raise_token_exchange_challenge_uses_insufficient_claims_with_claims_pre assert 'error="invalid_token"' not in www assert f'claims="{base64.b64encode(claims.encode()).decode()}"' in www assert claims not in www # raw JSON never appears; only the base64 form + + +def test_id_jag_client_secret_maps_to_config(): + spec = to_server_spec(_id_jag_server()) + assert spec is not None and isinstance(spec.config, IdJagConfig) + assert spec.config.org_token_endpoint == "https://idp.example.com/token" + assert spec.config.resource_token_endpoint == "https://mcp-as.example.com/token" + assert spec.config.client_id == "litellm-client-id" + assert spec.config.audience == "api://mcp-server" + assert spec.config.scopes == ("mcp.read", "mcp.write") + # ID-JAG asserts the user's id_token; the access_token default maps to id_token. + assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:id_token" + assert isinstance(spec.config.client_auth, ClientSecretAuth) + assert spec.config.client_auth.client_secret.get_secret_value() == ( + "litellm-client-secret" + ) + + +def test_id_jag_private_key_maps_to_private_key_jwt_auth(): + spec = to_server_spec( + _id_jag_server( + client_secret=None, + client_private_key="PEM-DATA", + client_private_key_id="kid-1", + client_assertion_signing_alg="RS384", + ) + ) + assert spec is not None and isinstance(spec.config, IdJagConfig) + assert isinstance(spec.config.client_auth, PrivateKeyJwtAuth) + assert spec.config.client_auth.private_key.get_secret_value() == "PEM-DATA" + assert spec.config.client_auth.key_id == "kid-1" + assert spec.config.client_auth.signing_alg == "RS384" + + +def test_id_jag_private_key_wins_over_client_secret(): + spec = to_server_spec(_id_jag_server(client_private_key="PEM-DATA")) + assert spec is not None and isinstance(spec.config, IdJagConfig) + assert isinstance(spec.config.client_auth, PrivateKeyJwtAuth) + + +def test_id_jag_honors_explicit_subject_token_type(): + spec = to_server_spec( + _id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2") + ) + assert spec is not None and isinstance(spec.config, IdJagConfig) + assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:saml2" + + +@pytest.mark.parametrize( + "server", + [ + _id_jag_server(token_exchange_endpoint=None), + _id_jag_server(id_jag_resource_token_endpoint=None), + _id_jag_server(client_id=None), + _id_jag_server(client_secret=None, client_private_key=None), + ], +) +def test_id_jag_half_configured_defers_to_v1(server): + # A half-configured server must defer (None) rather than 500 at IdJagConfig construction. + assert to_server_spec(server) is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index c88027abcd4..ba7720ffd51 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -16,8 +16,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( AwsSigV4Config, Byok, ClientCredentialsConfig, + ClientSecretAuth, CredError, Error, + IdJagConfig, NoneConfig, NoOpAuth, Ok, @@ -34,10 +36,40 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto OAuthToken, TokenStoreUnavailable, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( + ExchangedToken, +) _SUBJECT = Subject(tenant_id="", subject_id="") +def _id_jag_config() -> IdJagConfig: + return IdJagConfig( + org_token_endpoint="https://idp.example.com/token", + resource_token_endpoint="https://mcp-as.example.com/token", + client_id="litellm", + client_auth=ClientSecretAuth(client_secret=SecretStr("s")), + audience="api://mcp", + scopes=("mcp.read",), + ) + + +class _FakeTokenEndpoint: + """Records each fetch and returns the next canned Result, leg by leg.""" + + def __init__(self, results: list[Result[ExchangedToken, CredError]]) -> None: + self._results = list(results) + self.calls: list[tuple[str, str, dict[str, str]]] = [] + + async def fetch(self, endpoint, client_id, grant_params, client_auth): + self.calls.append((endpoint, client_id, dict(grant_params))) + return self._results.pop(0) + + +def _with_inbound(token: str) -> Subject: + return Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr(token)) + + def _spec(config): return ServerSpec(server_id="s", resource="https://upstream.example.com", config=config) @@ -305,3 +337,183 @@ async def test_unbuilt_arms_fail_closed_with_not_implemented(label, config): result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(config)) assert isinstance(result, Error) assert result.error.tag == "not_implemented" + + +@pytest.mark.asyncio +async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer(): + endpoint = _FakeTokenEndpoint( + [ + Ok(ExchangedToken(access_token="the-id-jag", expires_in=300)), + Ok(ExchangedToken(access_token="final-access", expires_in=3600)), + ] + ) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + result = await provider.resolve_credentials( + _with_inbound("user-id-token"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Ok) + assert _emitted(result.ok)["Authorization"] == "Bearer final-access" + + leg1_endpoint, _, leg1_params = endpoint.calls[0] + assert leg1_endpoint == "https://idp.example.com/token" + assert ( + leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + ) + assert ( + leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" + ) + assert leg1_params["subject_token"] == "user-id-token" + + leg2_endpoint, _, leg2_params = endpoint.calls[1] + assert leg2_endpoint == "https://mcp-as.example.com/token" + assert leg2_params["grant_type"] == "urn:ietf:params:oauth:grant-type:jwt-bearer" + # The leg-1 token is forwarded verbatim as the leg-2 assertion. + assert leg2_params["assertion"] == "the-id-jag" + + +@pytest.mark.asyncio +async def test_id_jag_without_inbound_token_is_precondition_required_no_http(): + endpoint = _FakeTokenEndpoint([]) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_id_jag_propagates_a_leg1_error_without_calling_leg2(): + endpoint = _FakeTokenEndpoint( + [Error(CredError.of_upstream_unavailable("leg1 down"))] + ) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + result = await provider.resolve_credentials( + _with_inbound("user-id-token"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert "leg1 down" in result.error.summary + assert len(endpoint.calls) == 1 + + +@pytest.mark.asyncio +async def test_id_jag_propagates_a_leg2_error(): + endpoint = _FakeTokenEndpoint( + [ + Ok(ExchangedToken(access_token="the-id-jag", expires_in=300)), + Error(CredError.of_upstream_unavailable("leg2 forbidden")), + ] + ) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + result = await provider.resolve_credentials( + _with_inbound("user-id-token"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert "leg2 forbidden" in result.error.summary + assert len(endpoint.calls) == 2 + + +def _two_leg_ok(bearer: str) -> list: + return [ + Ok(ExchangedToken(access_token="the-id-jag", expires_in=300)), + Ok(ExchangedToken(access_token=bearer, expires_in=3600)), + ] + + +@pytest.mark.asyncio +async def test_id_jag_reuses_the_cached_bearer_for_an_unchanged_config(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("first-bearer")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + + first = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) + second = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(second.ok)["Authorization"] == "Bearer first-bearer" + assert len(endpoint.calls) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "changed", + [ + _id_jag_config().model_copy(update={"audience": "api://other"}), + _id_jag_config().model_copy(update={"resource": "https://other.example.com/mcp"}), + _id_jag_config().model_copy(update={"scopes": ("mcp.read", "mcp.write")}), + _id_jag_config().model_copy(update={"org_token_endpoint": "https://idp.example.com/v2/token"}), + _id_jag_config().model_copy(update={"resource_token_endpoint": "https://mcp-as.example.com/v2/token"}), + _id_jag_config().model_copy(update={"client_id": "litellm-rotated"}), + _id_jag_config().model_copy(update={"client_auth": ClientSecretAuth(client_secret=SecretStr("rotated"))}), + _id_jag_config().model_copy(update={"subject_token_type": "urn:ietf:params:oauth:token-type:saml2"}), + ], + ids=[ + "audience", + "resource", + "scopes", + "org_token_endpoint", + "resource_token_endpoint", + "client_id", + "client_auth", + "subject_token_type", + ], +) +async def test_id_jag_config_change_forces_a_fresh_exchange(changed): + endpoint = _FakeTokenEndpoint(_two_leg_ok("old-policy-bearer") + _two_leg_ok("new-policy-bearer")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + + before = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) + after = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(changed)) + + assert isinstance(before, Ok) and isinstance(after, Ok) + assert _emitted(after.ok)["Authorization"] == "Bearer new-policy-bearer" + assert len(endpoint.calls) == 4 + + +@pytest.mark.asyncio +async def test_id_jag_does_not_share_the_cached_bearer_across_caller_tokens(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("alice-bearer") + _two_leg_ok("bob-bearer")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + + alice = await provider.resolve_credentials(_with_inbound("alice-id-token"), _spec(_id_jag_config())) + bob = await provider.resolve_credentials(_with_inbound("bob-id-token"), _spec(_id_jag_config())) + + assert isinstance(alice, Ok) and isinstance(bob, Ok) + assert _emitted(bob.ok)["Authorization"] == "Bearer bob-bearer" + assert len(endpoint.calls) == 4 + + +@pytest.mark.asyncio +async def test_invalidate_credentials_evicts_the_id_jag_bearer_so_the_next_resolve_re_exchanges(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("rejected-bearer") + _two_leg_ok("fresh-bearer")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + subject = _with_inbound("user-id-token") + + first = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + await provider.invalidate_credentials(subject, _spec(_id_jag_config())) + second = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(second.ok)["Authorization"] == "Bearer fresh-bearer" + assert len(endpoint.calls) == 4 + + +@pytest.mark.asyncio +async def test_invalidate_credentials_for_id_jag_is_a_noop_without_a_caller_token(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("cached-bearer")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + subject = _with_inbound("user-id-token") + + first = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + await provider.invalidate_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) + second = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(second.ok)["Authorization"] == "Bearer cached-bearer" + assert len(endpoint.calls) == 2 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 new file mode 100644 index 00000000000..f100bd56f8f --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py @@ -0,0 +1,408 @@ +"""Tests for the v2 token-endpoint collaborator. + +`TokenEndpointClient.fetch` makes one authenticated POST and returns the minted token as a value; +`ExchangedTokenCache` memoizes it with per-key single-flight. These pin the grant/client-auth wire +shape, the private-key-JWT vs client_secret authentication, the error-as-value mapping, and the +cache's hit/single-flight behavior. Each assertion fails under a real mutation of the feature. +""" + +import asyncio +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import jwt +import litellm +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa + +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( + CLIENT_ASSERTION_TYPE, + ExchangedToken, + ExchangedTokenCache, + TokenEndpointClient, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientSecretAuth, + CredError, + PrivateKeyJwtAuth, +) +from pydantic import SecretStr + +_PATCH_TARGET = ( + "litellm.proxy._experimental.mcp_server.outbound_credentials." + "token_endpoint.get_async_httpx_client" +) + +_ENDPOINT = "https://idp.example.com/oauth2/token" +_CLIENT_ID = "litellm-client-id" + +_RSA_KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048) +_PRIVATE_PEM = _RSA_KEY.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), +).decode() +_PUBLIC_PEM = ( + _RSA_KEY.public_key() + .public_bytes( + serialization.Encoding.PEM, + serialization.PublicFormat.SubjectPublicKeyInfo, + ) + .decode() +) + + +def _resp(token="access", expires_in=3600): + resp = MagicMock() + resp.json.return_value = {"access_token": token, "expires_in": expires_in} + resp.raise_for_status = MagicMock() + return resp + + +def _client(response): + client = AsyncMock() + client.post.return_value = response + return client + + +def _posted_data(client): + return client.post.call_args.kwargs["data"] + + +@pytest.mark.asyncio +async def test_fetch_forwards_grant_params_and_client_secret(): + client = _client(_resp("the-token", expires_in=1200)) + with patch(_PATCH_TARGET, return_value=client): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g", "subject_token": "user-jwt"}, + ClientSecretAuth(client_secret=SecretStr("shhh")), + ) + + assert isinstance(result, Ok) + assert result.ok == ExchangedToken(access_token="the-token", expires_in=1200) + assert client.post.call_args.args[0] == _ENDPOINT + data = _posted_data(client) + assert data["grant_type"] == "g" + assert data["subject_token"] == "user-jwt" + assert data["client_id"] == _CLIENT_ID + assert data["client_secret"] == "shhh" + assert "client_assertion" not in data + + +@pytest.mark.asyncio +async def test_fetch_private_key_jwt_client_assertion(): + client = _client(_resp()) + with patch(_PATCH_TARGET, return_value=client): + await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + PrivateKeyJwtAuth( + private_key=SecretStr(_PRIVATE_PEM), + key_id="kid-1", + signing_alg="RS256", + ), + ) + + data = _posted_data(client) + assert data["client_assertion_type"] == CLIENT_ASSERTION_TYPE + assert "client_secret" not in data + decoded = jwt.decode( + data["client_assertion"], + _PUBLIC_PEM, + algorithms=["RS256"], + audience=_ENDPOINT, + ) + assert decoded["iss"] == _CLIENT_ID + assert decoded["sub"] == _CLIENT_ID + assert decoded["aud"] == _ENDPOINT + assert "exp" in decoded + assert jwt.get_unverified_header(data["client_assertion"])["kid"] == "kid-1" + + +@pytest.mark.asyncio +async def test_fetch_http_error_maps_to_upstream_unavailable_with_status(): + error_resp = MagicMock() + error_resp.status_code = 403 + error_resp.raise_for_status.side_effect = httpx.HTTPStatusError( + "Forbidden", request=MagicMock(), response=error_resp + ) + with patch(_PATCH_TARGET, return_value=_client(error_resp)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert "403" in result.error.summary + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raised", + [ + httpx.ConnectError("connection refused", request=MagicMock()), + httpx.ReadTimeout("timed out", request=MagicMock()), + litellm.Timeout( + message="Connection timed out", + model="default-model-name", + llm_provider="litellm-httpx-handler", + ), + ], +) +async def test_fetch_network_error_maps_to_upstream_unavailable(raised): + client = AsyncMock() + client.post.side_effect = raised + with patch(_PATCH_TARGET, return_value=client): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth", + [ + PrivateKeyJwtAuth(private_key=SecretStr("not-a-pem-key"), signing_alg="RS256"), + PrivateKeyJwtAuth(private_key=SecretStr(_PRIVATE_PEM), signing_alg="XX999"), + ], + ids=["garbage-key", "unknown-alg"], +) +async def test_fetch_unsignable_client_assertion_is_misconfigured_not_a_crash(auth): + client = AsyncMock() + with patch(_PATCH_TARGET, return_value=client): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + auth, + ) + + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + client.post.assert_not_called() + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +@pytest.mark.asyncio +async def test_fetch_invalid_json_maps_to_upstream_unavailable(): + bad = MagicMock() + bad.raise_for_status = MagicMock() + bad.json.side_effect = json.JSONDecodeError("Expecting value", "", 0) + with patch(_PATCH_TARGET, return_value=_client(bad)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +@pytest.mark.asyncio +async def test_fetch_none_response_is_upstream_unavailable(): + with patch(_PATCH_TARGET, return_value=_client(None)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_fetch_missing_access_token_is_upstream_unavailable(): + bad = MagicMock() + bad.json.return_value = {"token_type": "Bearer"} + bad.raise_for_status = MagicMock() + with patch(_PATCH_TARGET, return_value=_client(bad)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_fetch_http_error_does_not_leak_endpoint_url(): + error_resp = MagicMock() + error_resp.status_code = 403 + error_resp.raise_for_status.side_effect = httpx.HTTPStatusError( + "Forbidden", request=MagicMock(), response=error_resp + ) + with patch(_PATCH_TARGET, return_value=_client(error_resp)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +@pytest.mark.asyncio +async def test_fetch_none_response_does_not_leak_endpoint_url(): + with patch(_PATCH_TARGET, return_value=_client(None)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +@pytest.mark.asyncio +async def test_fetch_missing_access_token_does_not_leak_endpoint_url(): + bad = MagicMock() + bad.json.return_value = {"token_type": "Bearer"} + bad.raise_for_status = MagicMock() + with patch(_PATCH_TARGET, return_value=_client(bad)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +def _ok_token(value="cached") -> Result[ExchangedToken, CredError]: + return Ok(ExchangedToken(access_token=value, expires_in=3600)) + + +@pytest.mark.asyncio +async def test_cache_hit_skips_the_second_compute(): + cache = ExchangedTokenCache() + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + return _ok_token("tok") + + first = await cache.get_or_compute("k", compute) + second = await cache.get_or_compute("k", compute) + + assert isinstance(first, Ok) and first.ok == "tok" + assert isinstance(second, Ok) and second.ok == "tok" + assert calls == 1 + + +@pytest.mark.asyncio +async def test_cache_single_flights_concurrent_misses(): + cache = ExchangedTokenCache() + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + await asyncio.sleep(0.01) + return _ok_token("shared") + + results = await asyncio.gather( + cache.get_or_compute("k", compute), + cache.get_or_compute("k", compute), + ) + + assert [r.ok for r in results] == ["shared", "shared"] + assert calls == 1 + + +@pytest.mark.asyncio +async def test_cache_invalidate_forces_the_next_compute(): + cache = ExchangedTokenCache() + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + return _ok_token(f"tok-{calls}") + + first = await cache.get_or_compute("k", compute) + cache.invalidate("k") + second = await cache.get_or_compute("k", compute) + + assert isinstance(first, Ok) and first.ok == "tok-1" + assert isinstance(second, Ok) and second.ok == "tok-2" + assert calls == 2 + + +@pytest.mark.asyncio +async def test_cache_invalidate_only_evicts_the_named_key(): + cache = ExchangedTokenCache() + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + return _ok_token(f"tok-{calls}") + + await cache.get_or_compute("keep", compute) + await cache.get_or_compute("evict", compute) + cache.invalidate("evict") + kept = await cache.get_or_compute("keep", compute) + + assert isinstance(kept, Ok) and kept.ok == "tok-1" + assert calls == 2 + + +@pytest.mark.asyncio +async def test_cache_does_not_store_a_failed_compute(): + cache = ExchangedTokenCache() + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + if calls == 1: + return Error(CredError.of_upstream_unavailable("down")) + return _ok_token("recovered") + + first = await cache.get_or_compute("k", compute) + second = await cache.get_or_compute("k", compute) + + assert isinstance(first, Error) + assert isinstance(second, Ok) and second.ok == "recovered" + assert calls == 2 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py index 43b3612a5f2..bb25ab6bd3c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py @@ -17,10 +17,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( AuthSpecKind, AwsSigV4Config, Byok, + ClientSecretAuth, CredError, Error, + IdJagConfig, NoneConfig, Ok, + PrivateKeyJwtAuth, ServerSpec, SharedKey, StaticKeys, @@ -29,6 +32,14 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( _AUTH_CONFIG = TypeAdapter(AuthConfig) +_ID_JAG_MINIMAL = { + "kind": "id_jag", + "org_token_endpoint": "https://idp.example.com/token", + "resource_token_endpoint": "https://mcp-as.example.com/token", + "client_id": "litellm", + "client_auth": {"source": "client_secret", "client_secret": "s"}, +} + def test_parse_auth_spec_kind_accepts_known_mode(): result = parse_auth_spec_kind("token_exchange") @@ -148,3 +159,73 @@ def test_secrets_do_not_leak_in_repr(): key = SharedKey(value=SecretStr("super-secret")) assert "super-secret" not in repr(key) assert key.value.get_secret_value() == "super-secret" + + +@pytest.mark.parametrize( + "missing", + ["org_token_endpoint", "resource_token_endpoint", "client_id", "client_auth"], +) +def test_id_jag_config_requires_each_endpoint_client_and_auth(missing): + payload = {k: v for k, v in _ID_JAG_MINIMAL.items() if k != missing} + with pytest.raises(ValidationError): + _AUTH_CONFIG.validate_python(payload) + + +def test_id_jag_client_auth_discriminates_on_source(): + by_secret = _AUTH_CONFIG.validate_python(_ID_JAG_MINIMAL) + assert isinstance(by_secret, IdJagConfig) + assert isinstance(by_secret.client_auth, ClientSecretAuth) + assert by_secret.client_auth.client_secret.get_secret_value() == "s" + + by_key = _AUTH_CONFIG.validate_python( + { + **_ID_JAG_MINIMAL, + "client_auth": { + "source": "private_key_jwt", + "private_key": "PEM", + "key_id": "kid-1", + "signing_alg": "RS384", + }, + } + ) + assert isinstance(by_key, IdJagConfig) + assert isinstance(by_key.client_auth, PrivateKeyJwtAuth) + assert by_key.client_auth.private_key.get_secret_value() == "PEM" + assert by_key.client_auth.key_id == "kid-1" + assert by_key.client_auth.signing_alg == "RS384" + + +def test_id_jag_client_auth_rejects_unknown_source(): + with pytest.raises(ValidationError): + _AUTH_CONFIG.validate_python( + {**_ID_JAG_MINIMAL, "client_auth": {"source": "mystery"}} + ) + + +def test_id_jag_config_defaults_id_token_subject_and_empty_optionals(): + config = _AUTH_CONFIG.validate_python(_ID_JAG_MINIMAL) + assert isinstance(config, IdJagConfig) + assert config.subject_token_type == "urn:ietf:params:oauth:token-type:id_token" + assert config.audience is None + assert config.resource is None + assert config.scopes == () + + +def test_id_jag_secrets_do_not_leak_in_repr(): + config = IdJagConfig( + org_token_endpoint="https://idp.example.com/token", + resource_token_endpoint="https://mcp-as.example.com/token", + client_id="litellm", + client_auth=PrivateKeyJwtAuth(private_key=SecretStr("super-secret-pem")), + ) + assert "super-secret-pem" not in repr(config) + + +def test_id_jag_server_spec_derives_auth_spec_kind(): + config = _AUTH_CONFIG.validate_python(_ID_JAG_MINIMAL) + spec = ServerSpec( + server_id="s", + resource="https://mcp.example.com/mcp", + config=config, + ) + assert spec.auth_spec_kind is AuthSpecKind.id_jag diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index a245200c4d1..56ca855c814 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -19,6 +19,8 @@ import pytest from litellm.proxy._experimental.mcp_server.db import ( _decode_user_credential, _prepare_mcp_server_data, + decrypt_credentials, + encrypt_credentials, get_user_credential, get_user_oauth_credential, is_oauth_credential_expired, @@ -332,6 +334,29 @@ def _stored_value(prisma) -> str: return create_value +# ── MCP server credentials at rest ────────────────────────────────────────────── + + +def test_client_private_key_encrypted_at_rest(): + """An ID-JAG client_private_key is a secret and must be encrypted in the stored + credentials blob, never persisted in plaintext, and must round-trip back. The + pre-fix code left client_private_key out of encrypt_credentials, so it was stored + verbatim.""" + private_key = ( + "-----BEGIN PRIVATE KEY-----\nsensitive-rsa-material\n-----END PRIVATE KEY-----" + ) + credentials = {"client_secret": "shh", "client_private_key": private_key} + + encrypted = encrypt_credentials(dict(credentials), encryption_key=None) + assert encrypted["client_private_key"] != private_key + assert private_key not in encrypted["client_private_key"] + assert encrypted["client_secret"] != "shh" + + decrypted = decrypt_credentials(dict(encrypted)) + assert decrypted["client_private_key"] == private_key + assert decrypted["client_secret"] == "shh" + + # ── BYOK round-trip ─────────────────────────────────────────────────────────── 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 55b6bbbdbc2..491fa023031 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 @@ -176,6 +176,89 @@ class TestMCPServerManager: assert calls == [("", "authz-srv")] assert client is not None + @pytest.mark.asyncio + async def test_caller_auth_header_cannot_bypass_id_jag_exchange(self): + """A caller-supplied per-request override must not disable the ID-JAG exchange and forward an + arbitrary bearer upstream: _create_mcp_client keeps the v2 spec and resolves through the + injected provider rather than deferring to the v1 caller-override path.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Ok, + ) + from litellm.types.mcp import MCPAuth + + calls = [] + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + calls.append((subject.subject_id, server.server_id)) + return Ok(StaticHeaderAuth("Bearer minted-id-jag-token")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = MCPServer( + server_id="id-jag-srv", + name="id-jag", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.example/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token", + ) + + client = await manager._create_mcp_client( + server, + mcp_auth_header="Bearer caller-supplied-token", + subject_token="caller-id-token", + ) + + assert calls == [("", "id-jag-srv")] + assert client is not None + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "missing_field", + ["token_exchange_endpoint", "id_jag_resource_token_endpoint", "client_id", "client_secret"], + ) + async def test_half_configured_id_jag_fails_closed_instead_of_deferring_to_v1(self, missing_field): + """ID-JAG has no v1 arm, so a half-configured oauth2_id_jag server must not silently fall + through to resolve_mcp_auth, where a caller x-mcp-* override or the static + authentication_token would bypass the per-user identity assertion. It must be refused as an + operator misconfiguration (HTTP 500) before any client is built.""" + from fastapi import HTTPException + + from litellm.types.mcp import MCPAuth + + fields = { + "client_id": "gateway-client", + "client_secret": "gateway-secret", + "token_exchange_endpoint": "https://org-idp.example/oauth2/token", + "id_jag_resource_token_endpoint": "https://resource-as.example/oauth2/token", + } + fields.pop(missing_field) + server = MCPServer( + server_id="id-jag-srv", + name="id-jag", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.oauth2_id_jag, + authentication_token="static-server-secret", + **fields, + ) + + with pytest.raises(HTTPException) as exc_info: + await MCPServerManager()._create_mcp_client( + server, + mcp_auth_header="Bearer caller-supplied-token", + subject_token="caller-id-token", + ) + + assert exc_info.value.status_code == 500 + assert "oauth2_id_jag" in str(exc_info.value.detail) + async def test_create_mcp_client_stdio_injects_npm_config_cache(self): """Test that _create_mcp_client injects NPM_CONFIG_CACHE when not already set, and preserves user-provided NPM_CONFIG_CACHE when present.""" @@ -257,6 +340,35 @@ class TestMCPServerManager: assert env == {} @pytest.mark.asyncio + async def test_load_servers_from_config_debug_dump_redacts_secrets(self, caplog): + """The registry debug dump must not leak long-lived credentials: the ID-JAG signing key, + client secret, and static token are masked while non-secret fields stay readable.""" + + manager = MCPServerManager() + config = { + "idjag": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2_id_jag, + "client_id": "gateway-client", + "client_secret": "SECRET-CLIENT-SECRET", + "client_private_key": "-----BEGIN PRIVATE KEY-----SECRET-PEM-----END PRIVATE KEY-----", + "token_exchange_endpoint": "https://org-idp.example/oauth2/token", + "id_jag_resource_token_endpoint": "https://resource-as.example/oauth2/token", + "authentication_token": "SECRET-STATIC-TOKEN", + } + } + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + await manager.load_servers_from_config(config) + + dump = next(m for m in caplog.messages if "Loaded MCP Servers" in m) + assert "SECRET-PEM" not in dump + assert "SECRET-CLIENT-SECRET" not in dump + assert "SECRET-STATIC-TOKEN" not in dump + assert "gateway-client" in dump + assert "https://org-idp.example/oauth2/token" in dump + async def test_load_servers_from_config_warns_on_invalid_alias(self, caplog): """Invalid aliases from config should emit warnings during load.""" @@ -8122,6 +8234,86 @@ class TestOBOCallToolRetry: manager._create_mcp_client.assert_awaited_once() assert first.attempts == 1 and retry.attempts == 1 + @pytest.mark.asyncio + async def test_upstream_401_on_id_jag_evicts_the_cached_bearer_and_retries(self): + """The retry path must invalidate the ID-JAG leg-2 bearer too: without eviction the rebuilt + client resolves the same rejected token from the cache and the retry 401s identically.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + IdJagConfig, + ) + + manager = self._manager() + success = CallToolResult(content=[], isError=False) + first = _RetryFakeClient(raises=_UpstreamAuthError(401)) + retry = _RetryFakeClient(result=success) + manager._create_mcp_client = AsyncMock(return_value=retry) + server = MCPServer( + server_id="id-jag-srv", + name="id-jag", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.example/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token", + ) + + result = await manager._obo_call_tool_with_retry( + client=first, + call_tool_params=MagicMock(), + host_progress_callback=None, + mcp_server=server, + server_auth_header=None, + extra_headers=None, + stdio_env=None, + subject_token="caller-id-token", + user_api_key_auth=None, + ) + + assert result is success + manager._cred_provider.invalidate_credentials.assert_awaited_once() + invalidated_spec = manager._cred_provider.invalidate_credentials.await_args.args[1] + assert isinstance(invalidated_spec.config, IdJagConfig) + assert first.attempts == 1 and retry.attempts == 1 + + @pytest.mark.asyncio + async def test_call_regular_routes_id_jag_through_the_retry_path(self): + """An oauth2_id_jag tool call with a subject token must take the invalidate-and-retry branch + of _call_regular_mcp_tool, not the plain single call, so an upstream 401 re-exchanges.""" + manager = self._manager() + success = CallToolResult(content=[], isError=False) + first = _RetryFakeClient(raises=_UpstreamAuthError(401)) + retry = _RetryFakeClient(result=success) + manager._create_mcp_client = AsyncMock(side_effect=[first, retry]) + server = MCPServer( + server_id="id-jag-srv", + name="id-jag", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.example/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token", + ) + + result = 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={"Authorization": "Bearer caller-id-token"}, + raw_headers=None, + proxy_logging_obj=None, + ) + + assert result is success + manager._cred_provider.invalidate_credentials.assert_awaited_once() + assert first.attempts == 1 and retry.attempts == 1 + @pytest.mark.asyncio async def test_non_auth_error_does_not_retry(self): manager = self._manager() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6dc63762e5c..daba5639a4f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27288,7 +27288,7 @@ export interface components { /** Alias */ alias?: string | null; /** Auth Type */ - auth_type?: ("none" | "api_key" | "bearer_token" | "basic" | "authorization" | "oauth2" | "aws_sigv4" | "token" | "oauth2_token_exchange" | "true_passthrough" | "oauth_delegate") | null; + auth_type?: ("none" | "api_key" | "bearer_token" | "basic" | "authorization" | "oauth2" | "aws_sigv4" | "token" | "oauth2_token_exchange" | "oauth2_id_jag" | "true_passthrough" | "oauth_delegate") | null; /** Mcp Info */ mcp_info?: { [key: string]: unknown; From 377d54e6946fff87a76a9d30c188ed10a4e1b896 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 18 Jul 2026 11:37:39 -0700 Subject: [PATCH 14/18] refactor(ui): migrate policy attachments table onto shared DataTable (#33827) * refactor(ui): migrate policy attachments table onto shared DataTable * refactor(ui): pass a specific success message to the attachment copy action --- ui/litellm-dashboard/eslint-suppressions.json | 13 - ...able.test.tsx => AttachmentTable.test.tsx} | 113 ++++--- .../policies/_components/AttachmentTable.tsx | 66 ++++ .../_components/AttachmentTableColumns.tsx | 186 +++++++++++ .../policies/_components/attachment_table.tsx | 291 ------------------ .../policies/_components/index.test.tsx | 23 +- .../policies/_components/index.tsx | 2 +- 7 files changed, 309 insertions(+), 385 deletions(-) rename ui/litellm-dashboard/src/app/(dashboard)/policies/_components/{attachment_table.test.tsx => AttachmentTable.test.tsx} (55%) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/policies/_components/attachment_table.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index dcf482450e9..c775af81ba8 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -881,19 +881,6 @@ "count": 1 } }, - "src/app/(dashboard)/policies/_components/attachment_table.test.tsx": { - "react/display-name": { - "count": 1 - } - }, - "src/app/(dashboard)/policies/_components/attachment_table.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx": { "no-nested-ternary": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/attachment_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx similarity index 55% rename from ui/litellm-dashboard/src/app/(dashboard)/policies/_components/attachment_table.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx index c53881e5cce..b544c44d190 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/attachment_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx @@ -1,63 +1,17 @@ import React from "react"; -import { screen } from "@testing-library/react"; +import { screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "@/../tests/test-utils"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import AttachmentTable from "./attachment_table"; +import AttachmentTable from "./AttachmentTable"; import { PolicyAttachment } from "@/components/policies/types"; vi.mock("./impact_popover", () => ({ - default: () =>
- - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - -
-
- {header.isPlaceholder ? null : flexRender(header.column.columnDef.header, header.getContext())} -
- {header.id !== "actions" && ( -
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
- )} -
- - ))} - - ))} - - - {isLoading ? ( - - -
-

Loading...

-
- - - ) : attachments.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - - {row.getVisibleCells().map((cell) => ( - - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No attachments found

-
- - - )} - -
- - - ); -}; - -export default AttachmentTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/index.test.tsx index 85d8c408032..3b6534ab0f0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/index.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/index.test.tsx @@ -56,21 +56,6 @@ vi.mock("./impact_popover", () => ({ default: () =>