From 4e9df553926347412a85e6b3bb22d4e5baae2afd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 12:21:03 -0700 Subject: [PATCH 01/26] fix(policy_engine): preserve config-defined policies across DB sync and expose them via list APIs --- litellm/proxy/_lazy_openapi_snapshot.json | 68 +++++- .../policy_engine/attachment_registry.py | 29 ++- .../proxy/policy_engine/policy_endpoints.py | 74 ++++++- .../proxy/policy_engine/policy_registry.py | 48 ++++- .../proxy/policy_engine/resolver_types.py | 10 +- .../policy_engine/test_attachment_registry.py | 64 ++++++ .../test_policy_engine_endpoints.py | 195 ++++++++++++++++++ .../policy_engine/test_policy_versioning.py | 79 +++++++ .../_components/AttachmentTableColumns.tsx | 7 + .../_components/PolicyTableColumns.tsx | 43 ++-- .../src/components/policies/types.ts | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 35 +++- 12 files changed, 611 insertions(+), 43 deletions(-) create mode 100644 tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 96f6ee89d56..12da0a26708 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -20402,6 +20402,16 @@ "description": "Who created the attachment.", "title": "Created By" }, + "definition_location": { + "default": "db", + "description": "Where this attachment is defined: 'db' (database) or 'config' (config.yaml).", + "enum": [ + "db", + "config" + ], + "title": "Definition Location", + "type": "string" + }, "keys": { "description": "Key patterns.", "items": { @@ -20658,6 +20668,16 @@ "description": "Who created the policy.", "title": "Created By" }, + "definition_location": { + "default": "db", + "description": "Where this policy is defined: 'db' (database) or 'config' (config.yaml).", + "enum": [ + "db", + "config" + ], + "title": "Definition Location", + "type": "string" + }, "description": { "anyOf": [ { @@ -21129,12 +21149,45 @@ "title": "PolicyVersionStatusUpdateRequest", "type": "object" }, + "UsageChartPoint": { + "properties": { + "blocked": { + "title": "Blocked", + "type": "integer" + }, + "date": { + "title": "Date", + "type": "string" + }, + "passed": { + "title": "Passed", + "type": "integer" + }, + "score": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Score" + } + }, + "required": [ + "date", + "passed", + "blocked" + ], + "title": "UsageChartPoint", + "type": "object" + }, "UsageOverviewResponse": { "properties": { "chart": { "items": { - "additionalProperties": true, - "type": "object" + "$ref": "#/components/schemas/UsageChartPoint" }, "title": "Chart", "type": "array" @@ -21243,6 +21296,13 @@ }, "ValidationError": { "properties": { + "ctx": { + "title": "Context", + "type": "object" + }, + "input": { + "title": "Input" + }, "loc": { "items": { "anyOf": [ @@ -21420,7 +21480,7 @@ }, "/policies/attachments/list": { "get": { - "description": "List all policy attachments from the database.\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/attachments/list\" \\\n -H \"Authorization: Bearer \"\n```\n\nExample Response:\n```json\n{\n \"attachments\": [\n {\n \"attachment_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"scope\": \"*\",\n \"teams\": [],\n \"keys\": [],\n \"models\": [],\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```", + "description": "List all policy attachments from the database and config.yaml.\n\nConfig-defined attachments are returned with definition_location \"config\" and a\nsynthetic attachment_id (\"config-\").\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/attachments/list\" \\\n -H \"Authorization: Bearer \"\n```\n\nExample Response:\n```json\n{\n \"attachments\": [\n {\n \"attachment_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"scope\": \"*\",\n \"teams\": [],\n \"keys\": [],\n \"models\": [],\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```", "operationId": "list_policy_attachments_policies_attachments_list_get", "responses": { "200": { @@ -21596,7 +21656,7 @@ }, "/policies/list": { "get": { - "description": "List all policies from the database. Optionally filter by version_status.\n\nQuery params:\n- version_status: Optional. One of \"draft\", \"published\", \"production\".\n If omitted, all versions are returned.\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/list\" \\\n -H \"Authorization: Bearer \"\ncurl -X GET \"http://localhost:4000/policies/list?version_status=production\" \\\n -H \"Authorization: Bearer \"\n```\n\nExample Response:\n```json\n{\n \"policies\": [\n {\n \"policy_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"version_number\": 1,\n \"version_status\": \"production\",\n \"inherit\": null,\n \"description\": \"Base guardrails for all requests\",\n \"guardrails_add\": [\"pii_masking\"],\n \"guardrails_remove\": [],\n \"condition\": null,\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```", + "description": "List all policies from the database and config.yaml. Optionally filter by version_status.\n\nConfig-defined policies are returned with definition_location \"config\" and are treated\nas production versions. On a name conflict with a DB policy, only the DB policy is returned.\n\nQuery params:\n- version_status: Optional. One of \"draft\", \"published\", \"production\".\n If omitted, all versions are returned.\n\nExample Request:\n```bash\ncurl -X GET \"http://localhost:4000/policies/list\" \\\n -H \"Authorization: Bearer \"\ncurl -X GET \"http://localhost:4000/policies/list?version_status=production\" \\\n -H \"Authorization: Bearer \"\n```\n\nExample Response:\n```json\n{\n \"policies\": [\n {\n \"policy_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"policy_name\": \"global-baseline\",\n \"version_number\": 1,\n \"version_status\": \"production\",\n \"inherit\": null,\n \"description\": \"Base guardrails for all requests\",\n \"guardrails_add\": [\"pii_masking\"],\n \"guardrails_remove\": [],\n \"condition\": null,\n \"created_at\": \"2024-01-01T00:00:00Z\",\n \"updated_at\": \"2024-01-01T00:00:00Z\"\n }\n ],\n \"total_count\": 1\n}\n```", "operationId": "list_policies_policies_list_get", "parameters": [ { diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 8ef509810ba..fe9ad3bef6d 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -42,6 +42,7 @@ class AttachmentRegistry: def __init__(self): self._attachments: List[PolicyAttachment] = [] + self._config_attachments: tuple[PolicyAttachment, ...] = () self._initialized: bool = False def load_attachments(self, attachments_config: List[Dict[str, Any]]) -> None: @@ -62,6 +63,7 @@ class AttachmentRegistry: verbose_proxy_logger.error(f"Error loading attachment: {str(e)}") raise ValueError(f"Invalid attachment: {str(e)}") from e + self._config_attachments = tuple(self._attachments) self._initialized = True verbose_proxy_logger.info(f"Loaded {len(self._attachments)} policy attachments") @@ -173,6 +175,15 @@ class AttachmentRegistry: """ return self._attachments.copy() + def get_config_attachments(self) -> tuple[PolicyAttachment, ...]: + """ + Get the attachments loaded from config.yaml. + + Returns: + Tuple of config-defined PolicyAttachment objects + """ + return self._config_attachments + def get_attachments_for_policy(self, policy_name: str) -> List[PolicyAttachment]: """ Get all attachments for a specific policy. @@ -199,6 +210,7 @@ class AttachmentRegistry: Clear all attachments from the registry. """ self._attachments = [] + self._config_attachments = () self._initialized = False def add_attachment(self, attachment: PolicyAttachment) -> None: @@ -428,6 +440,7 @@ class AttachmentRegistry: ) -> None: """ Sync policy attachments from the database to in-memory registry. + Config-loaded attachments are preserved. Args: prisma_client: The Prisma client instance @@ -435,11 +448,8 @@ class AttachmentRegistry: try: attachments = await self.get_all_attachments_from_db(prisma_client) - # Clear existing attachments and reload from DB - self._attachments = [] - - for attachment_response in attachments: - attachment = PolicyAttachment( + db_attachments = [ + PolicyAttachment( policy=attachment_response.policy_name, scope=attachment_response.scope, teams=(attachment_response.teams if attachment_response.teams else None), @@ -447,10 +457,15 @@ class AttachmentRegistry: models=(attachment_response.models if attachment_response.models else None), tags=attachment_response.tags if attachment_response.tags else None, ) - self._attachments.append(attachment) + for attachment_response in attachments + ] + self._attachments = [*self._config_attachments, *db_attachments] self._initialized = True - verbose_proxy_logger.info(f"Synced {len(attachments)} attachments from DB to in-memory registry") + verbose_proxy_logger.info( + f"Synced {len(attachments)} attachments from DB to in-memory registry " + f"({len(self._config_attachments)} config-defined attachments preserved)" + ) except Exception as e: verbose_proxy_logger.exception(f"Error syncing attachments from DB: {e}") raise Exception(f"Error syncing attachments from DB: {str(e)}") diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index a879f6b6f7e..787f7069996 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -17,6 +17,8 @@ from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.types.proxy.policy_engine import ( GuardrailPipeline, PipelineTestRequest, + Policy, + PolicyAttachment, PolicyAttachmentCreateRequest, PolicyAttachmentDBResponse, PolicyAttachmentListResponse, @@ -33,6 +35,35 @@ from litellm.types.proxy.policy_engine import ( router = APIRouter() +def _config_policy_to_db_response(policy_name: str, policy: Policy) -> PolicyDBResponse: + return PolicyDBResponse( + policy_id=policy_name, + policy_name=policy_name, + version_number=1, + version_status="production", + inherit=policy.inherit, + description=policy.description, + guardrails_add=policy.guardrails.get_add(), + guardrails_remove=policy.guardrails.get_remove(), + condition=policy.condition.model_dump() if policy.condition else None, + pipeline=policy.pipeline.model_dump() if policy.pipeline else None, + definition_location="config", + ) + + +def _config_attachment_to_db_response(index: int, attachment: PolicyAttachment) -> PolicyAttachmentDBResponse: + return PolicyAttachmentDBResponse( + attachment_id=f"config-{index}", + policy_name=attachment.policy, + scope=attachment.scope, + teams=attachment.teams or [], + keys=attachment.keys or [], + models=attachment.models or [], + tags=attachment.tags or [], + definition_location="config", + ) + + # ───────────────────────────────────────────────────────────────────────────── # Policy CRUD Endpoints # ───────────────────────────────────────────────────────────────────────────── @@ -46,7 +77,10 @@ router = APIRouter() ) async def list_policies(version_status: Optional[str] = None): """ - List all policies from the database. Optionally filter by version_status. + List all policies from the database and config.yaml. Optionally filter by version_status. + + Config-defined policies are returned with definition_location "config" and are treated + as production versions. On a name conflict with a DB policy, only the DB policy is returned. Query params: - version_status: Optional. One of "draft", "published", "production". @@ -84,11 +118,25 @@ async def list_policies(version_status: Optional[str] = None): """ from litellm.proxy.proxy_server import prisma_client - if prisma_client is None: - raise HTTPException(status_code=500, detail="Database not connected") - try: - policies = await get_policy_registry().get_all_policies_from_db(prisma_client, version_status=version_status) + registry = get_policy_registry() + db_policies = ( + await registry.get_all_policies_from_db(prisma_client, version_status=version_status) + if prisma_client is not None + else [] + ) + db_policy_names = {db_policy.policy_name for db_policy in db_policies} + include_config = version_status in (None, "production") + config_policies = ( + [ + _config_policy_to_db_response(policy_name, policy) + for policy_name, policy in registry.list_config_policies().items() + if policy_name not in db_policy_names and registry.get_source(policy_name) != "db" + ] + if include_config + else [] + ) + policies = db_policies + config_policies return PolicyListDBResponse(policies=policies, total_count=len(policies)) except Exception as e: verbose_proxy_logger.exception(f"Error listing policies: {e}") @@ -606,7 +654,10 @@ async def test_pipeline( ) async def list_policy_attachments(): """ - List all policy attachments from the database. + List all policy attachments from the database and config.yaml. + + Config-defined attachments are returned with definition_location "config" and a + synthetic attachment_id ("config-"). Example Request: ```bash @@ -635,11 +686,14 @@ async def list_policy_attachments(): """ from litellm.proxy.proxy_server import prisma_client - if prisma_client is None: - raise HTTPException(status_code=500, detail="Database not connected") - try: - attachments = await get_attachment_registry().get_all_attachments_from_db(prisma_client) + registry = get_attachment_registry() + db_attachments = await registry.get_all_attachments_from_db(prisma_client) if prisma_client is not None else [] + config_attachments = [ + _config_attachment_to_db_response(index, attachment) + for index, attachment in enumerate(registry.get_config_attachments()) + ] + attachments = db_attachments + config_attachments return PolicyAttachmentListResponse(attachments=attachments, total_count=len(attachments)) except Exception as e: verbose_proxy_logger.exception(f"Error listing policy attachments: {e}") diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index e1afbf2f5f2..9456afaa1b9 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -13,6 +13,7 @@ from datetime import datetime, timezone from typing import ( TYPE_CHECKING, Any, + Literal, Optional, Protocol, TypedDict, @@ -162,6 +163,8 @@ class PolicyRegistry: def __init__(self): self._policies: dict[str, Policy] = {} + self._config_policies: Mapping[str, Policy] = {} + self._sources: Mapping[str, Literal["db", "config"]] = {} self._policies_by_id: dict[str, tuple[str, Policy]] = {} self._initialized: bool = False @@ -174,6 +177,8 @@ class PolicyRegistry: This is the raw config from the YAML file. """ self._policies = {} + self._config_policies = {} + self._sources = {} self._policies_by_id = {} for policy_name, policy_data in policies_config.items(): @@ -185,6 +190,8 @@ class PolicyRegistry: verbose_proxy_logger.error(f"Error loading policy '{policy_name}': {str(e)}") raise ValueError(f"Invalid policy '{policy_name}': {str(e)}") from e + self._config_policies = dict(self._policies) + self._sources = {policy_name: "config" for policy_name in self._policies} self._initialized = True verbose_proxy_logger.info(f"Loaded {len(self._policies)} policies") @@ -299,17 +306,35 @@ class PolicyRegistry: Clear all policies from the registry. """ self._policies = {} + self._config_policies = {} + self._sources = {} self._initialized = False - def add_policy(self, policy_name: str, policy: Policy) -> None: + def get_source(self, policy_name: str) -> Optional[Literal["db", "config"]]: + """ + Return the provenance of an in-memory policy, or None if unknown. + """ + return self._sources.get(policy_name) + + def list_config_policies(self) -> Mapping[str, Policy]: + """ + Return the policies loaded from config.yaml, keyed by policy name. + """ + return dict(self._config_policies) + + def add_policy(self, policy_name: str, policy: Policy, source: Literal["db", "config"] = "db") -> None: """ Add or update a single policy. Args: policy_name: Name of the policy policy: Policy object to add + source: Provenance of the policy ("db" or "config") """ self._policies[policy_name] = policy + self._sources = {**self._sources, policy_name: source} + if source == "config": + self._config_policies = {**self._config_policies, policy_name: policy} self._initialized = True verbose_proxy_logger.debug(f"Added/updated policy: {policy_name}") @@ -325,6 +350,7 @@ class PolicyRegistry: """ if policy_name in self._policies: del self._policies[policy_name] + self._sources = {name: source for name, source in self._sources.items() if name != policy_name} verbose_proxy_logger.debug(f"Removed policy: {policy_name}") return True return False @@ -591,14 +617,14 @@ class PolicyRegistry: """ Sync policies from the database to in-memory registry. - Production versions are loaded into _policies (by policy name) for resolution. + - Config-loaded policies are preserved; on a name conflict the DB version wins. - Draft and published versions are loaded into _policies_by_id so request-body policy_ overrides can be resolved without DB access in the hot path. """ try: - self._policies = {} production = await self.get_all_policies_from_db(prisma_client, version_status="production") - for policy_response in production: - policy = self._parse_policy( + db_policies = { + policy_response.policy_name: self._parse_policy( policy_response.policy_name, { "inherit": policy_response.inherit, @@ -611,7 +637,16 @@ class PolicyRegistry: "pipeline": policy_response.pipeline, }, ) - self.add_policy(policy_response.policy_name, policy) + for policy_response in production + } + for policy_name in set(db_policies) & set(self._config_policies): + verbose_proxy_logger.warning( + f"Policy '{policy_name}' is defined in both config.yaml and the DB; the DB version takes precedence" + ) + config_sources: Mapping[str, Literal["db", "config"]] = {name: "config" for name in self._config_policies} + db_sources: Mapping[str, Literal["db", "config"]] = {name: "db" for name in db_policies} + self._policies = {**self._config_policies, **db_policies} + self._sources = {**config_sources, **db_sources} self._policies_by_id = {} non_production = await _policy_table(prisma_client).find_many( @@ -637,7 +672,8 @@ class PolicyRegistry: self._initialized = True verbose_proxy_logger.info( f"Synced {len(production)} production policies and {len(non_production)} " - "draft/published (by ID) from DB to in-memory registry" + "draft/published (by ID) from DB to in-memory registry " + f"({len(self._config_policies)} config-defined policies preserved)" ) except Exception as e: verbose_proxy_logger.exception(f"Error syncing policies from DB: {e}") diff --git a/litellm/types/proxy/policy_engine/resolver_types.py b/litellm/types/proxy/policy_engine/resolver_types.py index 2c7e8d5afc9..b4096cd2044 100644 --- a/litellm/types/proxy/policy_engine/resolver_types.py +++ b/litellm/types/proxy/policy_engine/resolver_types.py @@ -6,7 +6,7 @@ the final guardrails list. """ from datetime import datetime -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Literal, Optional from pydantic import BaseModel, ConfigDict, Field @@ -220,6 +220,10 @@ class PolicyDBResponse(BaseModel): updated_at: Optional[datetime] = Field(default=None, description="When the policy was last updated.") created_by: Optional[str] = Field(default=None, description="Who created the policy.") updated_by: Optional[str] = Field(default=None, description="Who last updated the policy.") + definition_location: Literal["db", "config"] = Field( + default="db", + description="Where this policy is defined: 'db' (database) or 'config' (config.yaml).", + ) class PolicyListDBResponse(BaseModel): @@ -317,6 +321,10 @@ class PolicyAttachmentDBResponse(BaseModel): updated_at: Optional[datetime] = Field(default=None, description="When the attachment was last updated.") created_by: Optional[str] = Field(default=None, description="Who created the attachment.") updated_by: Optional[str] = Field(default=None, description="Who last updated the attachment.") + definition_location: Literal["db", "config"] = Field( + default="db", + description="Where this attachment is defined: 'db' (database) or 'config' (config.yaml).", + ) class PolicyAttachmentListResponse(BaseModel): diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index 1ae0b4d3d48..cf470c66000 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -4,6 +4,9 @@ Unit tests for AttachmentRegistry - tests policy attachment matching. Tests the main entry point: get_attached_policies() """ +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + import pytest from litellm.proxy.policy_engine.attachment_registry import ( @@ -389,3 +392,64 @@ class TestAttachmentRegistrySingleton: registry1 = get_attachment_registry() registry2 = get_attachment_registry() assert registry1 is registry2 + + +def _make_db_attachment_row(attachment_id="att-1", policy_name="db-policy", scope=None, teams=None): + row = MagicMock() + row.attachment_id = attachment_id + row.policy_name = policy_name + row.scope = scope + row.teams = teams or [] + row.keys = [] + row.models = [] + row.tags = [] + row.created_at = datetime.now(timezone.utc) + row.updated_at = datetime.now(timezone.utc) + row.created_by = None + row.updated_by = None + return row + + +def _prisma_with_attachment_rows(rows): + prisma = MagicMock() + prisma.db.litellm_policyattachmenttable.find_many = AsyncMock(return_value=rows) + return prisma + + +class TestConfigAttachmentsPreservedAcrossDbSync: + """Config-defined attachments must survive sync_attachments_from_db (regression for issue #35255).""" + + @pytest.mark.asyncio + async def test_sync_with_empty_db_preserves_config_attachments(self): + registry = AttachmentRegistry() + registry.load_attachments([{"policy": "config-policy", "scope": "*"}]) + + await registry.sync_attachments_from_db(_prisma_with_attachment_rows([])) + + context = PolicyMatchContext(team_alias="any-team", key_alias="any-key", model="gpt-5.2") + assert registry.get_attached_policies(context) == ["config-policy"] + + @pytest.mark.asyncio + async def test_sync_merges_db_attachments_with_config_attachments(self): + registry = AttachmentRegistry() + registry.load_attachments([{"policy": "config-policy", "scope": "*"}]) + db_row = _make_db_attachment_row(policy_name="db-policy", teams=["db-team"]) + + await registry.sync_attachments_from_db(_prisma_with_attachment_rows([db_row])) + + assert len(registry.get_all_attachments()) == 2 + assert len(registry.get_config_attachments()) == 1 + context = PolicyMatchContext(team_alias="db-team", key_alias="k", model="gpt-5.2") + attached = registry.get_attached_policies(context) + assert "config-policy" in attached + assert "db-policy" in attached + + @pytest.mark.asyncio + async def test_repeated_syncs_do_not_duplicate_config_attachments(self): + registry = AttachmentRegistry() + registry.load_attachments([{"policy": "config-policy", "scope": "*"}]) + + await registry.sync_attachments_from_db(_prisma_with_attachment_rows([])) + await registry.sync_attachments_from_db(_prisma_with_attachment_rows([])) + + assert len(registry.get_all_attachments()) == 1 diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py b/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py new file mode 100644 index 00000000000..78126508c2d --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py @@ -0,0 +1,195 @@ +""" +Unit tests for policy_engine/policy_endpoints.py list endpoints. + +Regression tests for issue #35255: config-defined policies and attachments must be +returned by the list endpoints (marked definition_location="config"), DB rows must keep +their exact shape, and the endpoints must not 500 when no database is connected. +""" + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm.proxy.policy_engine.policy_endpoints as policy_endpoints +from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry +from litellm.proxy.policy_engine.policy_registry import PolicyRegistry + + +def _make_policy_row( + policy_id="uuid-1", + policy_name="db-policy", + version_status="production", + guardrails_add=None, +): + row = MagicMock() + row.policy_id = policy_id + row.policy_name = policy_name + row.version_number = 1 + row.version_status = version_status + row.parent_version_id = None + row.is_latest = True + row.published_at = None + row.production_at = None + row.inherit = None + row.description = "db description" + row.guardrails_add = guardrails_add or [] + row.guardrails_remove = [] + row.condition = None + row.pipeline = None + row.created_at = datetime.now(timezone.utc) + row.updated_at = datetime.now(timezone.utc) + row.created_by = "admin" + row.updated_by = "admin" + return row + + +def _make_attachment_row(attachment_id="att-1", policy_name="db-policy", scope="*"): + row = MagicMock() + row.attachment_id = attachment_id + row.policy_name = policy_name + row.scope = scope + row.teams = [] + row.keys = [] + row.models = [] + row.tags = [] + row.created_at = datetime.now(timezone.utc) + row.updated_at = datetime.now(timezone.utc) + row.created_by = "admin" + row.updated_by = "admin" + return row + + +@pytest.fixture +def policy_registry(monkeypatch): + registry = PolicyRegistry() + monkeypatch.setattr(policy_endpoints, "get_policy_registry", lambda: registry) + return registry + + +@pytest.fixture +def attachment_registry(monkeypatch): + registry = AttachmentRegistry() + monkeypatch.setattr(policy_endpoints, "get_attachment_registry", lambda: registry) + return registry + + +def _set_prisma(monkeypatch, prisma): + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + + +class TestListPoliciesIncludesConfig: + @pytest.mark.asyncio + async def test_returns_config_policies_without_prisma(self, policy_registry, monkeypatch): + _set_prisma(monkeypatch, None) + policy_registry.load_policies( + {"config-policy": {"description": "from config", "guardrails": {"add": ["tooling"]}}} + ) + + response = await policy_endpoints.list_policies() + + assert response.total_count == 1 + entry = response.policies[0] + assert entry.policy_name == "config-policy" + assert entry.policy_id == "config-policy" + assert entry.definition_location == "config" + assert entry.version_status == "production" + assert entry.guardrails_add == ["tooling"] + assert entry.description == "from config" + assert entry.created_at is None + + @pytest.mark.asyncio + async def test_merges_db_rows_with_config_and_keeps_db_row_shape(self, policy_registry, monkeypatch): + row = _make_policy_row(policy_id="uuid-1", policy_name="db-policy", guardrails_add=["db-guard"]) + prisma = MagicMock() + prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[row]) + _set_prisma(monkeypatch, prisma) + policy_registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}}) + + response = await policy_endpoints.list_policies() + + assert response.total_count == 2 + db_entry = next(p for p in response.policies if p.policy_name == "db-policy") + assert db_entry.definition_location == "db" + assert db_entry.policy_id == "uuid-1" + assert db_entry.guardrails_add == ["db-guard"] + assert db_entry.description == "db description" + assert db_entry.created_at == row.created_at + assert db_entry.created_by == "admin" + config_entry = next(p for p in response.policies if p.policy_name == "config-policy") + assert config_entry.definition_location == "config" + + @pytest.mark.asyncio + async def test_db_policy_shadows_config_policy_with_same_name(self, policy_registry, monkeypatch): + row = _make_policy_row(policy_id="uuid-1", policy_name="shared-name", guardrails_add=["db-guard"]) + prisma = MagicMock() + prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[row]) + _set_prisma(monkeypatch, prisma) + policy_registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}}) + + response = await policy_endpoints.list_policies() + + assert response.total_count == 1 + assert response.policies[0].definition_location == "db" + assert response.policies[0].guardrails_add == ["db-guard"] + + @pytest.mark.asyncio + async def test_version_status_filter_excludes_config_policies(self, policy_registry, monkeypatch): + row = _make_policy_row(policy_id="uuid-1", policy_name="db-policy", version_status="draft") + prisma = MagicMock() + prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[row]) + _set_prisma(monkeypatch, prisma) + policy_registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}}) + + response = await policy_endpoints.list_policies(version_status="draft") + + assert response.total_count == 1 + assert response.policies[0].policy_name == "db-policy" + assert response.policies[0].definition_location == "db" + + @pytest.mark.asyncio + async def test_production_filter_includes_config_policies(self, policy_registry, monkeypatch): + _set_prisma(monkeypatch, None) + policy_registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}}) + + response = await policy_endpoints.list_policies(version_status="production") + + assert response.total_count == 1 + assert response.policies[0].definition_location == "config" + + +class TestListAttachmentsIncludesConfig: + @pytest.mark.asyncio + async def test_returns_config_attachments_without_prisma(self, attachment_registry, monkeypatch): + _set_prisma(monkeypatch, None) + attachment_registry.load_attachments([{"policy": "config-policy", "scope": "*"}]) + + response = await policy_endpoints.list_policy_attachments() + + assert response.total_count == 1 + entry = response.attachments[0] + assert entry.attachment_id == "config-0" + assert entry.policy_name == "config-policy" + assert entry.scope == "*" + assert entry.definition_location == "config" + assert entry.created_at is None + + @pytest.mark.asyncio + async def test_merges_db_attachments_with_config_and_keeps_db_row_shape(self, attachment_registry, monkeypatch): + row = _make_attachment_row(attachment_id="att-1", policy_name="db-policy") + prisma = MagicMock() + prisma.db.litellm_policyattachmenttable.find_many = AsyncMock(return_value=[row]) + _set_prisma(monkeypatch, prisma) + attachment_registry.load_attachments([{"policy": "config-policy", "scope": "*"}]) + + response = await policy_endpoints.list_policy_attachments() + + assert response.total_count == 2 + db_entry = next(a for a in response.attachments if a.policy_name == "db-policy") + assert db_entry.attachment_id == "att-1" + assert db_entry.definition_location == "db" + assert db_entry.created_at == row.created_at + assert db_entry.created_by == "admin" + config_entry = next(a for a in response.attachments if a.policy_name == "config-policy") + assert config_entry.attachment_id == "config-0" + assert config_entry.definition_location == "config" diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py index dd20021d0e1..5a840979c1b 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py @@ -450,3 +450,82 @@ class TestGetPolicyRegistrySingleton: a = get_policy_registry() b = get_policy_registry() assert a is b + + +def _prisma_with_policy_rows(production_rows, non_production_rows=None): + prisma = MagicMock() + prisma.db.litellm_policytable.find_many = AsyncMock(side_effect=[production_rows, non_production_rows or []]) + return prisma + + +class TestConfigPoliciesPreservedAcrossDbSync: + """Config-defined policies must survive sync_policies_from_db (regression for issue #35255).""" + + @pytest.mark.asyncio + async def test_sync_with_empty_db_preserves_config_policies(self): + registry = PolicyRegistry() + registry.load_policies({"config-policy": {"description": "from config", "guardrails": {"add": ["tooling"]}}}) + + await registry.sync_policies_from_db(_prisma_with_policy_rows([])) + + assert registry.has_policy("config-policy") + policy = registry.get_policy("config-policy") + assert policy is not None + assert policy.guardrails.add == ["tooling"] + assert registry.get_source("config-policy") == "config" + + @pytest.mark.asyncio + async def test_sync_merges_db_policies_with_config_policies(self): + registry = PolicyRegistry() + registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}}) + db_row = _make_row(policy_id="db-1", policy_name="db-policy", guardrails_add=["db-guard"]) + + await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row])) + + assert registry.has_policy("config-policy") + assert registry.has_policy("db-policy") + assert registry.get_source("config-policy") == "config" + assert registry.get_source("db-policy") == "db" + + @pytest.mark.asyncio + async def test_db_wins_on_policy_name_conflict(self): + registry = PolicyRegistry() + registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}}) + db_row = _make_row(policy_id="db-1", policy_name="shared-name", guardrails_add=["db-guard"]) + + await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row])) + + policy = registry.get_policy("shared-name") + assert policy is not None + assert policy.guardrails.add == ["db-guard"] + assert registry.get_source("shared-name") == "db" + + @pytest.mark.asyncio + async def test_config_policy_restored_after_conflicting_db_row_deleted(self): + registry = PolicyRegistry() + registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}}) + db_row = _make_row(policy_id="db-1", policy_name="shared-name", guardrails_add=["db-guard"]) + + await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row])) + await registry.sync_policies_from_db(_prisma_with_policy_rows([])) + + policy = registry.get_policy("shared-name") + assert policy is not None + assert policy.guardrails.add == ["config-guard"] + assert registry.get_source("shared-name") == "config" + + @pytest.mark.asyncio + async def test_config_policy_resolves_guardrails_after_sync(self): + from litellm.proxy.policy_engine.policy_resolver import PolicyResolver + + registry = PolicyRegistry() + registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}}) + + await registry.sync_policies_from_db(_prisma_with_policy_rows([])) + + resolved = PolicyResolver.resolve_policy_guardrails( + policy_name="config-policy", + policies=registry.get_all_policies(), + context=None, + ) + assert resolved.guardrails == ["tooling"] diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx index 9e0f8d6715d..ded9e3a1e6d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTableColumns.tsx @@ -41,7 +41,12 @@ interface AttachmentRowActionsProps { onDeleteClick: (attachmentId: string) => void; } +const CONFIG_ATTACHMENT_HINT = + "Config attachments are defined in the config file and cannot be deleted from the dashboard."; + function AttachmentRowActions({ attachment, isAdmin, onDeleteClick }: AttachmentRowActionsProps) { + const isConfigAttachment = attachment.definition_location === "config"; + return ( onDeleteClick(attachment.attachment_id)} > diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx index dd036d83283..488de728bad 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTableColumns.tsx @@ -22,6 +22,9 @@ export interface PolicyRow { versionCount: number; } +const CONFIG_POLICY_HINT = + "Config policies are defined in the config file and cannot be edited or deleted from the dashboard."; + function GuardrailChips({ guardrails, tone }: { guardrails: string[]; tone: "success" | "error" }) { if (guardrails.length === 0) { return -; @@ -45,6 +48,8 @@ interface PolicyRowActionsProps { } function PolicyRowActions({ policy, onEditClick, onDeleteClick }: PolicyRowActionsProps) { + const isConfigPolicy = policy.definition_location === "config"; + return ( - onEditClick(policy)}> + onEditClick(policy)} + > Edit policy @@ -63,6 +73,8 @@ function PolicyRowActions({ policy, onEditClick, onDeleteClick }: PolicyRowActio onDeleteClick(policy.policy_id, policy.policy_name || "Unnamed Policy")} > @@ -93,18 +105,23 @@ export const getPolicyTableColumns = ({ header: ({ column }) => , size: 220, enableSorting: true, - cell: ({ row }) => ( - 1 ? ( - - ) : undefined - } - onClick={() => onViewClick(row.original.primaryPolicy.policy_id)} - /> - ), + cell: ({ row }) => { + const isConfigPolicy = row.original.primaryPolicy.definition_location === "config"; + const versionBadge = + row.original.versionCount > 1 ? ( + + ) : undefined; + return ( + : versionBadge + } + onClick={isConfigPolicy ? undefined : () => onViewClick(row.original.primaryPolicy.policy_id)} + /> + ); + }, }, { id: "description", diff --git a/ui/litellm-dashboard/src/components/policies/types.ts b/ui/litellm-dashboard/src/components/policies/types.ts index 887781ff943..6ac110e3c0a 100644 --- a/ui/litellm-dashboard/src/components/policies/types.ts +++ b/ui/litellm-dashboard/src/components/policies/types.ts @@ -14,6 +14,7 @@ export interface Policy { updated_at?: string; created_by?: string; updated_by?: string; + definition_location?: "db" | "config"; } export interface PolicyCondition { @@ -47,6 +48,7 @@ export interface PolicyAttachment { updated_at?: string; created_by?: string; updated_by?: string; + definition_location?: "db" | "config"; } export interface PolicyCreateRequest { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ed975c6be0a..6543e437c07 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9379,7 +9379,10 @@ export interface paths { }; /** * List Policy Attachments - * @description List all policy attachments from the database. + * @description List all policy attachments from the database and config.yaml. + * + * Config-defined attachments are returned with definition_location "config" and a + * synthetic attachment_id ("config-"). * * Example Request: * ```bash @@ -9487,7 +9490,10 @@ export interface paths { }; /** * List Policies - * @description List all policies from the database. Optionally filter by version_status. + * @description List all policies from the database and config.yaml. Optionally filter by version_status. + * + * Config-defined policies are returned with definition_location "config" and are treated + * as production versions. On a name conflict with a DB policy, only the DB policy is returned. * * Query params: * - version_status: Optional. One of "draft", "published", "production". @@ -29379,6 +29385,13 @@ export interface components { * @description Who created the attachment. */ created_by?: string | null; + /** + * Definition Location + * @description Where this attachment is defined: 'db' (database) or 'config' (config.yaml). + * @default db + * @enum {string} + */ + definition_location: "db" | "config"; /** * Keys * @description Key patterns. @@ -29510,6 +29523,13 @@ export interface components { * @description Who created the policy. */ created_by?: string | null; + /** + * Definition Location + * @description Where this policy is defined: 'db' (database) or 'config' (config.yaml). + * @default db + * @enum {string} + */ + definition_location: "db" | "config"; /** * Description * @description Policy description. @@ -33117,6 +33137,17 @@ export interface components { */ model?: string | null; }; + /** UsageChartPoint */ + UsageChartPoint: { + /** Blocked */ + blocked: number; + /** Date */ + date: string; + /** Passed */ + passed: number; + /** Score */ + score?: number | null; + }; /** UsageDetailResponse */ UsageDetailResponse: { /** Avglatency */ From c1f5abf817624610df6e3ec806f87a1d0e2483e9 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 30 Jul 2026 13:06:41 -0700 Subject: [PATCH 02/26] chore: placeholder commit to open the gpt pricing change PR From f1b781d06b6155df7c8979110ddc45938c3b81fb Mon Sep 17 00:00:00 2001 From: lihugang <94808542+lihugang@users.noreply.github.com> Date: Fri, 31 Jul 2026 04:08:01 +0800 Subject: [PATCH 03/26] fix(pricing): adjust gpt-5.6-terra and gpt-5.6-luna prices according to OpenAI's latest article (#35258) Adjust the price of gpt-5.6-terra to 80% of its original rate (2/12), and gpt-5.6-luna to 20% of its original rate (0.2/1.2). References: https://openai.com/index/advancing-the-price-performance-frontier-with-gpt-5-6/ https://developers.openai.com/api/docs/pricing --- ...odel_prices_and_context_window_backup.json | 72 +++++++++---------- model_prices_and_context_window.json | 72 +++++++++---------- .../llm_cost_calc/test_llm_cost_calc_utils.py | 4 +- .../test_gpt_5_6_model_metadata.py | 15 ++-- 4 files changed, 85 insertions(+), 78 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cc418c9c428..a431c4b29fc 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23782,29 +23782,29 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-terra": { - "cache_creation_input_token_cost": 3.125e-06, - "cache_creation_input_token_cost_above_272k_tokens": 6.25e-06, - "cache_creation_input_token_cost_flex": 1.5625e-06, - "cache_creation_input_token_cost_priority": 6.25e-06, - "cache_read_input_token_cost": 2.5e-07, - "cache_read_input_token_cost_above_272k_tokens": 5e-07, - "cache_read_input_token_cost_flex": 1.25e-07, - "cache_read_input_token_cost_priority": 5e-07, - "input_cost_per_token": 2.5e-06, - "input_cost_per_token_above_272k_tokens": 5e-06, - "input_cost_per_token_batches": 1.25e-06, - "input_cost_per_token_flex": 1.25e-06, - "input_cost_per_token_priority": 5e-06, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_flex": 1.25e-06, + "cache_creation_input_token_cost_priority": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 4e-06, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.5e-05, - "output_cost_per_token_above_272k_tokens": 2.25e-05, - "output_cost_per_token_batches": 7.5e-06, - "output_cost_per_token_flex": 7.5e-06, - "output_cost_per_token_priority": 3e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.4e-05, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ @@ -23835,29 +23835,29 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-luna": { - "cache_creation_input_token_cost": 1.25e-06, - "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, - "cache_creation_input_token_cost_flex": 6.25e-07, - "cache_creation_input_token_cost_priority": 2.5e-06, - "cache_read_input_token_cost": 1e-07, - "cache_read_input_token_cost_above_272k_tokens": 2e-07, - "cache_read_input_token_cost_flex": 5e-08, - "cache_read_input_token_cost_priority": 2e-07, - "input_cost_per_token": 1e-06, - "input_cost_per_token_above_272k_tokens": 2e-06, - "input_cost_per_token_batches": 5e-07, - "input_cost_per_token_flex": 5e-07, - "input_cost_per_token_priority": 2e-06, + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_flex": 1.25e-07, + "cache_creation_input_token_cost_priority": 5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "cache_read_input_token_cost_flex": 1e-08, + "cache_read_input_token_cost_priority": 4e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "input_cost_per_token_batches": 1e-07, + "input_cost_per_token_flex": 1e-07, + "input_cost_per_token_priority": 4e-07, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 6e-06, - "output_cost_per_token_above_272k_tokens": 9e-06, - "output_cost_per_token_batches": 3e-06, - "output_cost_per_token_flex": 3e-06, - "output_cost_per_token_priority": 1.2e-05, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "output_cost_per_token_batches": 6e-07, + "output_cost_per_token_flex": 6e-07, + "output_cost_per_token_priority": 2.4e-06, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c4628fecdb8..184bc69c9e8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23857,29 +23857,29 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-terra": { - "cache_creation_input_token_cost": 3.125e-06, - "cache_creation_input_token_cost_above_272k_tokens": 6.25e-06, - "cache_creation_input_token_cost_flex": 1.5625e-06, - "cache_creation_input_token_cost_priority": 6.25e-06, - "cache_read_input_token_cost": 2.5e-07, - "cache_read_input_token_cost_above_272k_tokens": 5e-07, - "cache_read_input_token_cost_flex": 1.25e-07, - "cache_read_input_token_cost_priority": 5e-07, - "input_cost_per_token": 2.5e-06, - "input_cost_per_token_above_272k_tokens": 5e-06, - "input_cost_per_token_batches": 1.25e-06, - "input_cost_per_token_flex": 1.25e-06, - "input_cost_per_token_priority": 5e-06, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_flex": 1.25e-06, + "cache_creation_input_token_cost_priority": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "cache_read_input_token_cost_flex": 1e-07, + "cache_read_input_token_cost_priority": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 4e-06, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.5e-05, - "output_cost_per_token_above_272k_tokens": 2.25e-05, - "output_cost_per_token_batches": 7.5e-06, - "output_cost_per_token_flex": 7.5e-06, - "output_cost_per_token_priority": 3e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.4e-05, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ @@ -23910,29 +23910,29 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-luna": { - "cache_creation_input_token_cost": 1.25e-06, - "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, - "cache_creation_input_token_cost_flex": 6.25e-07, - "cache_creation_input_token_cost_priority": 2.5e-06, - "cache_read_input_token_cost": 1e-07, - "cache_read_input_token_cost_above_272k_tokens": 2e-07, - "cache_read_input_token_cost_flex": 5e-08, - "cache_read_input_token_cost_priority": 2e-07, - "input_cost_per_token": 1e-06, - "input_cost_per_token_above_272k_tokens": 2e-06, - "input_cost_per_token_batches": 5e-07, - "input_cost_per_token_flex": 5e-07, - "input_cost_per_token_priority": 2e-06, + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_flex": 1.25e-07, + "cache_creation_input_token_cost_priority": 5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "cache_read_input_token_cost_flex": 1e-08, + "cache_read_input_token_cost_priority": 4e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "input_cost_per_token_batches": 1e-07, + "input_cost_per_token_flex": 1e-07, + "input_cost_per_token_priority": 4e-07, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 6e-06, - "output_cost_per_token_above_272k_tokens": 9e-06, - "output_cost_per_token_batches": 3e-06, - "output_cost_per_token_flex": 3e-06, - "output_cost_per_token_priority": 1.2e-05, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "output_cost_per_token_batches": 6e-07, + "output_cost_per_token_flex": 6e-07, + "output_cost_per_token_priority": 2.4e-06, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index d282e656ce8..1ad271900f6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -611,8 +611,8 @@ def test_generic_cost_per_token_gpt55_pro(): [ ("gpt-5.6", 5e-6, 3e-5, 5e-7, 6.25e-6), ("gpt-5.6-sol", 5e-6, 3e-5, 5e-7, 6.25e-6), - ("gpt-5.6-terra", 2.5e-6, 1.5e-5, 2.5e-7, 3.125e-6), - ("gpt-5.6-luna", 1e-6, 6e-6, 1e-7, 1.25e-6), + ("gpt-5.6-terra", 2e-6, 1.2e-5, 2e-7, 2.5e-6), + ("gpt-5.6-luna", 2e-7, 1.2e-6, 2e-8, 2.5e-7), ], ) def test_generic_cost_per_token_gpt56( diff --git a/tests/test_litellm/test_gpt_5_6_model_metadata.py b/tests/test_litellm/test_gpt_5_6_model_metadata.py index 5a7b621d521..b17ef32a1dc 100644 --- a/tests/test_litellm/test_gpt_5_6_model_metadata.py +++ b/tests/test_litellm/test_gpt_5_6_model_metadata.py @@ -7,7 +7,14 @@ from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider GPT_5_6_MODELS = ("gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna") -STANDARD_PRICING = { +OPENAI_STANDARD_PRICING = { + "gpt-5.6": (5e-06, 3e-05, 5e-07, 6.25e-06), + "gpt-5.6-sol": (5e-06, 3e-05, 5e-07, 6.25e-06), + "gpt-5.6-terra": (2e-06, 1.2e-05, 2e-07, 2.5e-06), + "gpt-5.6-luna": (2e-07, 1.2e-06, 2e-08, 2.5e-07), +} + +AZURE_STANDARD_PRICING = { "gpt-5.6": (5e-06, 3e-05, 5e-07, 6.25e-06), "gpt-5.6-sol": (5e-06, 3e-05, 5e-07, 6.25e-06), "gpt-5.6-terra": (2.5e-06, 1.5e-05, 2.5e-07, 3.125e-06), @@ -27,7 +34,7 @@ def test_openai_gpt_5_6_model_info(model): assert info["litellm_provider"] == "openai" assert info["mode"] == "chat" - input_cost, output_cost, cache_read_cost, cache_write_cost = STANDARD_PRICING[model] + input_cost, output_cost, cache_read_cost, cache_write_cost = OPENAI_STANDARD_PRICING[model] assert info["input_cost_per_token"] == input_cost assert info["output_cost_per_token"] == output_cost assert info["cache_read_input_token_cost"] == cache_read_cost @@ -95,7 +102,7 @@ def test_azure_gpt_5_6_global_model_info(model): assert info["litellm_provider"] == "azure" assert info["mode"] == "chat" - input_cost, output_cost, cache_read_cost, _ = STANDARD_PRICING[_tier_key(model)] + input_cost, output_cost, cache_read_cost, _ = AZURE_STANDARD_PRICING[_tier_key(model)] assert info["input_cost_per_token"] == input_cost assert info["output_cost_per_token"] == output_cost assert info["cache_read_input_token_cost"] == cache_read_cost @@ -124,7 +131,7 @@ def test_azure_gpt_5_6_regional_model_info(model): assert info["litellm_provider"] == "azure" assert info["mode"] == "chat" - input_cost, output_cost, cache_read_cost, _ = STANDARD_PRICING[_tier_key(model)] + input_cost, output_cost, cache_read_cost, _ = AZURE_STANDARD_PRICING[_tier_key(model)] assert info["input_cost_per_token"] == pytest.approx(input_cost * 1.1) assert info["output_cost_per_token"] == pytest.approx(output_cost * 1.1) From 6aea561319e1b26aace283473ef5ab543d06c703 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 30 Jul 2026 13:29:44 -0700 Subject: [PATCH 04/26] fix(pricing): correct bedrock_mantle gpt-5.6 terra/luna prices after OpenAI's cut AWS rolled out the 2026-07-30 GPT-5.6 price cut the same day, but the bedrock_mantle entries still carried values derived from the pre-cut OpenAI base, so Terra billed 1.25x and Luna 5x over the published rate. Re-derive both from the AWS Bedrock pricing page, which prices in-region inference at parity with OpenAI's data residency tier (1.1x base). Sol was not cut and is unchanged. Also drop tests/test_litellm/test_gpt_5_6_model_metadata.py; its Azure and openai pricing assertions are covered by test_llm_cost_calc_utils.py. --- ...odel_prices_and_context_window_backup.json | 16 +- model_prices_and_context_window.json | 16 +- ...bedrock_mantle_responses_transformation.py | 4 +- .../test_gpt_5_6_model_metadata.py | 166 ------------------ 4 files changed, 18 insertions(+), 184 deletions(-) delete mode 100644 tests/test_litellm/test_gpt_5_6_model_metadata.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a431c4b29fc..d5f412b9279 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45120,10 +45120,10 @@ "supports_vision": true }, "bedrock_mantle/openai.gpt-5.6-terra": { - "input_cost_per_token": 2.75e-06, - "cache_creation_input_token_cost": 3.4375e-06, - "cache_read_input_token_cost": 2.75e-07, - "output_cost_per_token": 1.65e-05, + "input_cost_per_token": 2.2e-06, + "cache_creation_input_token_cost": 2.75e-06, + "cache_read_input_token_cost": 2.2e-07, + "output_cost_per_token": 1.32e-05, "litellm_provider": "bedrock_mantle", "max_input_tokens": 272000, "max_output_tokens": 128000, @@ -45148,10 +45148,10 @@ "supports_vision": true }, "bedrock_mantle/openai.gpt-5.6-luna": { - "input_cost_per_token": 1.1e-06, - "cache_creation_input_token_cost": 1.375e-06, - "cache_read_input_token_cost": 1.1e-07, - "output_cost_per_token": 6.6e-06, + "input_cost_per_token": 2.2e-07, + "cache_creation_input_token_cost": 2.75e-07, + "cache_read_input_token_cost": 2.2e-08, + "output_cost_per_token": 1.32e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 272000, "max_output_tokens": 128000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 184bc69c9e8..e0917ed87b4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45241,10 +45241,10 @@ "supports_vision": true }, "bedrock_mantle/openai.gpt-5.6-terra": { - "input_cost_per_token": 2.75e-06, - "cache_creation_input_token_cost": 3.4375e-06, - "cache_read_input_token_cost": 2.75e-07, - "output_cost_per_token": 1.65e-05, + "input_cost_per_token": 2.2e-06, + "cache_creation_input_token_cost": 2.75e-06, + "cache_read_input_token_cost": 2.2e-07, + "output_cost_per_token": 1.32e-05, "litellm_provider": "bedrock_mantle", "max_input_tokens": 272000, "max_output_tokens": 128000, @@ -45269,10 +45269,10 @@ "supports_vision": true }, "bedrock_mantle/openai.gpt-5.6-luna": { - "input_cost_per_token": 1.1e-06, - "cache_creation_input_token_cost": 1.375e-06, - "cache_read_input_token_cost": 1.1e-07, - "output_cost_per_token": 6.6e-06, + "input_cost_per_token": 2.2e-07, + "cache_creation_input_token_cost": 2.75e-07, + "cache_read_input_token_cost": 2.2e-08, + "output_cost_per_token": 1.32e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 272000, "max_output_tokens": 128000, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 4dd1c663de0..e808023debe 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -1519,8 +1519,8 @@ class TestBedrockMantleResponsesPricing: "model, input_cost, cache_creation_cost, cache_read_cost, output_cost", [ ("openai.gpt-5.6-sol", 5.5e-06, 6.875e-06, 5.5e-07, 3.3e-05), - ("openai.gpt-5.6-terra", 2.75e-06, 3.4375e-06, 2.75e-07, 1.65e-05), - ("openai.gpt-5.6-luna", 1.1e-06, 1.375e-06, 1.1e-07, 6.6e-06), + ("openai.gpt-5.6-terra", 2.2e-06, 2.75e-06, 2.2e-07, 1.32e-05), + ("openai.gpt-5.6-luna", 2.2e-07, 2.75e-07, 2.2e-08, 1.32e-06), ], ) def test_gpt_5_6_pricing_and_mode( diff --git a/tests/test_litellm/test_gpt_5_6_model_metadata.py b/tests/test_litellm/test_gpt_5_6_model_metadata.py deleted file mode 100644 index b17ef32a1dc..00000000000 --- a/tests/test_litellm/test_gpt_5_6_model_metadata.py +++ /dev/null @@ -1,166 +0,0 @@ -import json -from pathlib import Path - -import pytest - -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - -GPT_5_6_MODELS = ("gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna") - -OPENAI_STANDARD_PRICING = { - "gpt-5.6": (5e-06, 3e-05, 5e-07, 6.25e-06), - "gpt-5.6-sol": (5e-06, 3e-05, 5e-07, 6.25e-06), - "gpt-5.6-terra": (2e-06, 1.2e-05, 2e-07, 2.5e-06), - "gpt-5.6-luna": (2e-07, 1.2e-06, 2e-08, 2.5e-07), -} - -AZURE_STANDARD_PRICING = { - "gpt-5.6": (5e-06, 3e-05, 5e-07, 6.25e-06), - "gpt-5.6-sol": (5e-06, 3e-05, 5e-07, 6.25e-06), - "gpt-5.6-terra": (2.5e-06, 1.5e-05, 2.5e-07, 3.125e-06), - "gpt-5.6-luna": (1e-06, 6e-06, 1e-07, 1.25e-06), -} - - -@pytest.mark.parametrize("model", GPT_5_6_MODELS) -def test_openai_gpt_5_6_model_info(model): - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - info = model_cost.get(model) - assert info is not None, f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "openai" - assert info["mode"] == "chat" - - input_cost, output_cost, cache_read_cost, cache_write_cost = OPENAI_STANDARD_PRICING[model] - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cache_read_cost - assert info["cache_creation_input_token_cost"] == cache_write_cost - assert info["cache_creation_input_token_cost"] == pytest.approx(input_cost * 1.25) - - assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2) - assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(output_cost * 1.5) - assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2) - - assert info["max_input_tokens"] == 1050000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_web_search"] is True - assert info["supports_none_reasoning_effort"] is True - assert info["supports_xhigh_reasoning_effort"] is True - assert info["supports_minimal_reasoning_effort"] is False - - assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/batch", "/v1/responses"] - assert info["supported_modalities"] == ["text", "image"] - assert info["supported_output_modalities"] == ["text"] - - routed_model, provider, _, _ = get_llm_provider(model=f"openai/{model}") - assert routed_model == model - assert provider == "openai" - - -AZURE_GLOBAL_MODELS = ( - "azure/gpt-5.6", - "azure/gpt-5.6-sol", - "azure/gpt-5.6-terra", - "azure/gpt-5.6-luna", -) - -AZURE_REGIONAL_MODELS = tuple( - f"azure/{region}/{tier}" - for region in ("us", "eu") - for tier in ("gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna") -) - - -def _tier_key(azure_model): - return azure_model.split("/")[-1] - - -def _load_main(): - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - return json.load(f) - - -@pytest.mark.parametrize("model", AZURE_GLOBAL_MODELS) -def test_azure_gpt_5_6_global_model_info(model): - model_cost = _load_main() - info = model_cost.get(model) - assert info is not None, f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "azure" - assert info["mode"] == "chat" - - input_cost, output_cost, cache_read_cost, _ = AZURE_STANDARD_PRICING[_tier_key(model)] - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cache_read_cost - - assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2) - assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(output_cost * 1.5) - assert info["input_cost_per_token_priority"] == pytest.approx(input_cost * 2) - assert info["output_cost_per_token_priority"] == pytest.approx(output_cost * 2) - assert info["input_cost_per_token_above_272k_tokens_priority"] == pytest.approx(input_cost * 4) - assert info["output_cost_per_token_above_272k_tokens_priority"] == pytest.approx(output_cost * 3) - - assert info["max_input_tokens"] == 1050000 - assert info["max_output_tokens"] == 128000 - assert info["supports_reasoning"] is True - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert provider == "azure" - - -@pytest.mark.parametrize("model", AZURE_REGIONAL_MODELS) -def test_azure_gpt_5_6_regional_model_info(model): - model_cost = _load_main() - info = model_cost.get(model) - assert info is not None, f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "azure" - assert info["mode"] == "chat" - - input_cost, output_cost, cache_read_cost, _ = AZURE_STANDARD_PRICING[_tier_key(model)] - - assert info["input_cost_per_token"] == pytest.approx(input_cost * 1.1) - assert info["output_cost_per_token"] == pytest.approx(output_cost * 1.1) - assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost * 1.1) - assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2.2) - assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(output_cost * 1.65) - assert info["input_cost_per_token_priority"] == pytest.approx(input_cost * 2.75) - assert info["output_cost_per_token_priority"] == pytest.approx(output_cost * 2.75) - - assert info["max_input_tokens"] == 1050000 - assert info["max_output_tokens"] == 128000 - assert info["supports_reasoning"] is True - - _, provider, _, _ = get_llm_provider(model=model) - assert provider == "azure" - - -def test_gpt_5_6_backup_matches_main(): - """Ensure the bundled model cost map stays in sync with the canonical file.""" - repo_root = Path(__file__).parents[2] - main_path = repo_root / "model_prices_and_context_window.json" - backup_path = repo_root / "litellm" / "model_prices_and_context_window_backup.json" - - with open(main_path) as f: - main_cost = json.load(f) - with open(backup_path) as f: - backup_cost = json.load(f) - - for model in GPT_5_6_MODELS + AZURE_GLOBAL_MODELS + AZURE_REGIONAL_MODELS: - assert backup_cost.get(model) == main_cost.get(model), ( - f"{model} differs between main and backup model cost maps" - ) From 62aebaf035611bf91d6ef1f9f370b62ca9e7bb62 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 30 Jul 2026 14:16:51 -0700 Subject: [PATCH 05/26] fix(pricing): bill gpt-5.6 flex requests above 272k at the flex long-context rate OpenAI publishes a long-context column on the Flex tier, at half the standard long-context rate. We had no field for it, so a >272k flex request fell through to the standard long-context price and billed 2x: Terra $4/$18 instead of $2/$9, Luna $0.40/$1.80 instead of $0.20/$0.90, Sol $10/$45 instead of $5/$22.50. Adding the values to the cost map alone does nothing, because get_model_info builds ModelInfoBase from an explicit kwargs list and silently drops any key not named there. Declare the four *_above_272k_tokens_flex fields and wire them through, then add the values for sol, terra, luna, and the gpt-5.6 alias. That same gap was already swallowing cache_creation_input_token_cost_flex, _priority, and _above_272k_tokens, which were present in the cost map but never reached the calculator; they are wired through here too. Fast mode (ex-Priority) publishes no long-context column, so nothing is added there rather than deriving a rate by analogy. --- ...odel_prices_and_context_window_backup.json | 16 +++++++ litellm/types/utils.py | 16 +++++++ litellm/utils.py | 22 +++++++++ model_prices_and_context_window.json | 16 +++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 48 +++++++++++++++++++ tests/test_litellm/test_utils.py | 8 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 32 +++++++++++++ 7 files changed, 158 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d5f412b9279..86c9eb0b70d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23678,14 +23678,17 @@ "gpt-5.6": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06, "cache_creation_input_token_cost_flex": 3.125e-06, "cache_creation_input_token_cost_priority": 1.25e-05, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 1e-06, "input_cost_per_token": 5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, "input_cost_per_token_batches": 2.5e-06, "input_cost_per_token_flex": 2.5e-06, "input_cost_per_token_priority": 1e-05, @@ -23696,6 +23699,7 @@ "mode": "chat", "output_cost_per_token": 3e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_priority": 6e-05, @@ -23731,14 +23735,17 @@ "gpt-5.6-sol": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06, "cache_creation_input_token_cost_flex": 3.125e-06, "cache_creation_input_token_cost_priority": 1.25e-05, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 1e-06, "input_cost_per_token": 5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, "input_cost_per_token_batches": 2.5e-06, "input_cost_per_token_flex": 2.5e-06, "input_cost_per_token_priority": 1e-05, @@ -23749,6 +23756,7 @@ "mode": "chat", "output_cost_per_token": 3e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_priority": 6e-05, @@ -23784,14 +23792,17 @@ "gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, "cache_creation_input_token_cost_flex": 1.25e-06, "cache_creation_input_token_cost_priority": 5e-06, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 4e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 4e-06, @@ -23802,6 +23813,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, + "output_cost_per_token_above_272k_tokens_flex": 9e-06, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.4e-05, @@ -23837,14 +23849,17 @@ "gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_flex": 1.25e-07, "cache_creation_input_token_cost_priority": 5e-07, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, + "input_cost_per_token_above_272k_tokens_flex": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_token_flex": 1e-07, "input_cost_per_token_priority": 4e-07, @@ -23855,6 +23870,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "output_cost_per_token_above_272k_tokens": 1.8e-06, + "output_cost_per_token_above_272k_tokens_flex": 9e-07, "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, "output_cost_per_token_priority": 2.4e-06, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9df44c6202c..b2d463e2e08 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -200,7 +200,12 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing cache_creation_input_token_cost: Optional[float] cache_creation_input_token_cost_above_200k_tokens: Optional[float] + cache_creation_input_token_cost_above_272k_tokens: Optional[float] + cache_creation_input_token_cost_above_272k_tokens_priority: Optional[float] + cache_creation_input_token_cost_above_272k_tokens_flex: Optional[float] cache_creation_input_token_cost_above_1hr: Optional[float] + cache_creation_input_token_cost_flex: Optional[float] # OpenAI flex service tier pricing + cache_creation_input_token_cost_priority: Optional[float] # OpenAI priority service tier pricing cache_read_input_token_cost: Optional[float] cache_read_input_token_cost_flex: Optional[float] # OpenAI flex service tier pricing cache_read_input_token_cost_priority: Optional[float] # OpenAI priority service tier pricing @@ -208,6 +213,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_read_input_token_cost_above_200k_tokens_priority: Optional[float] cache_read_input_token_cost_above_272k_tokens: Optional[float] cache_read_input_token_cost_above_272k_tokens_priority: Optional[float] + cache_read_input_token_cost_above_272k_tokens_flex: Optional[float] cache_read_input_token_cost_above_512k_tokens: Optional[float] # Smallest prefix this model will actually cache, whatever caching mechanism its provider uses. # Absent means the provider-agnostic default applies; see MINIMUM_PROMPT_CACHE_TOKEN_COUNT. @@ -219,6 +225,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_token_above_200k_tokens_priority: Optional[float] input_cost_per_token_above_272k_tokens: Optional[float] # GPT-5.4/5.4-pro: prompts >272K priced at 2x input input_cost_per_token_above_272k_tokens_priority: Optional[float] + input_cost_per_token_above_272k_tokens_flex: Optional[float] input_cost_per_token_above_512k_tokens: Optional[float] # MiniMax-M3: prompts >512K priced at 2x input input_cost_per_character_above_128k_tokens: Optional[float] # only for vertex ai models input_cost_per_query: Optional[float] # only for rerank models @@ -246,6 +253,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token_above_200k_tokens_priority: Optional[float] output_cost_per_token_above_272k_tokens: Optional[float] # GPT-5.4/5.4-pro: prompts >272K priced at 1.5x output output_cost_per_token_above_272k_tokens_priority: Optional[float] + output_cost_per_token_above_272k_tokens_flex: Optional[float] output_cost_per_token_above_512k_tokens: Optional[float] # MiniMax-M3: prompts >512K priced at 2x output output_cost_per_character_above_128k_tokens: Optional[float] # only for vertex ai models output_cost_per_image: Optional[float] @@ -3158,6 +3166,11 @@ class CustomPricingLiteLLMParams(BaseModel): cache_creation_input_token_cost: Optional[float] = None cache_creation_input_token_cost_above_1hr: Optional[float] = None cache_creation_input_token_cost_above_200k_tokens: Optional[float] = None + cache_creation_input_token_cost_above_272k_tokens: Optional[float] = None + cache_creation_input_token_cost_above_272k_tokens_priority: Optional[float] = None + cache_creation_input_token_cost_above_272k_tokens_flex: Optional[float] = None + cache_creation_input_token_cost_flex: Optional[float] = None + cache_creation_input_token_cost_priority: Optional[float] = None cache_creation_input_audio_token_cost: Optional[float] = None cache_read_input_token_cost: Optional[float] = None cache_read_input_token_cost_flex: Optional[float] = None @@ -3165,6 +3178,7 @@ class CustomPricingLiteLLMParams(BaseModel): cache_read_input_token_cost_above_200k_tokens: Optional[float] = None cache_read_input_token_cost_above_200k_tokens_priority: Optional[float] = None cache_read_input_token_cost_above_272k_tokens_priority: Optional[float] = None + cache_read_input_token_cost_above_272k_tokens_flex: Optional[float] = None cache_read_input_audio_token_cost: Optional[float] = None input_cost_per_character: Optional[float] = None input_cost_per_character_above_128k_tokens: Optional[float] = None @@ -3174,6 +3188,7 @@ class CustomPricingLiteLLMParams(BaseModel): input_cost_per_token_above_200k_tokens: Optional[float] = None input_cost_per_token_above_200k_tokens_priority: Optional[float] = None input_cost_per_token_above_272k_tokens_priority: Optional[float] = None + input_cost_per_token_above_272k_tokens_flex: Optional[float] = None input_cost_per_query: Optional[float] = None input_cost_per_image: Optional[float] = None input_cost_per_image_above_128k_tokens: Optional[float] = None @@ -3193,6 +3208,7 @@ class CustomPricingLiteLLMParams(BaseModel): output_cost_per_token_above_200k_tokens: Optional[float] = None output_cost_per_token_above_200k_tokens_priority: Optional[float] = None output_cost_per_token_above_272k_tokens_priority: Optional[float] = None + output_cost_per_token_above_272k_tokens_flex: Optional[float] = None output_cost_per_character_above_128k_tokens: Optional[float] = None output_cost_per_image: Optional[float] = None output_cost_per_image_token: Optional[float] = None diff --git a/litellm/utils.py b/litellm/utils.py index 944bb61d5e7..4c66955e366 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5410,6 +5410,19 @@ def _get_model_info_helper( cache_creation_input_token_cost_above_200k_tokens=_model_info.get( "cache_creation_input_token_cost_above_200k_tokens", None ), + cache_creation_input_token_cost_above_272k_tokens=_model_info.get( + "cache_creation_input_token_cost_above_272k_tokens", None + ), + cache_creation_input_token_cost_above_272k_tokens_priority=_model_info.get( + "cache_creation_input_token_cost_above_272k_tokens_priority", None + ), + cache_creation_input_token_cost_above_272k_tokens_flex=_model_info.get( + "cache_creation_input_token_cost_above_272k_tokens_flex", None + ), + cache_creation_input_token_cost_flex=_model_info.get("cache_creation_input_token_cost_flex", None), + cache_creation_input_token_cost_priority=_model_info.get( + "cache_creation_input_token_cost_priority", None + ), cache_read_input_token_cost=_model_info.get("cache_read_input_token_cost", None), prompt_cache_min_tokens=_model_info.get("prompt_cache_min_tokens", None), cache_read_input_token_cost_above_200k_tokens=_model_info.get( @@ -5424,6 +5437,9 @@ def _get_model_info_helper( cache_read_input_token_cost_above_272k_tokens_priority=_model_info.get( "cache_read_input_token_cost_above_272k_tokens_priority", None ), + cache_read_input_token_cost_above_272k_tokens_flex=_model_info.get( + "cache_read_input_token_cost_above_272k_tokens_flex", None + ), cache_read_input_token_cost_above_512k_tokens=_model_info.get( "cache_read_input_token_cost_above_512k_tokens", None ), @@ -5442,6 +5458,9 @@ def _get_model_info_helper( input_cost_per_token_above_272k_tokens_priority=_model_info.get( "input_cost_per_token_above_272k_tokens_priority", None ), + input_cost_per_token_above_272k_tokens_flex=_model_info.get( + "input_cost_per_token_above_272k_tokens_flex", None + ), input_cost_per_token_above_512k_tokens=_model_info.get("input_cost_per_token_above_512k_tokens", None), input_cost_per_query=_model_info.get("input_cost_per_query", None), input_cost_per_second=_model_info.get("input_cost_per_second", None), @@ -5483,6 +5502,9 @@ def _get_model_info_helper( output_cost_per_token_above_272k_tokens_priority=_model_info.get( "output_cost_per_token_above_272k_tokens_priority", None ), + output_cost_per_token_above_272k_tokens_flex=_model_info.get( + "output_cost_per_token_above_272k_tokens_flex", None + ), output_cost_per_token_above_512k_tokens=_model_info.get( "output_cost_per_token_above_512k_tokens", None ), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e0917ed87b4..90129dd5ddc 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23753,14 +23753,17 @@ "gpt-5.6": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06, "cache_creation_input_token_cost_flex": 3.125e-06, "cache_creation_input_token_cost_priority": 1.25e-05, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 1e-06, "input_cost_per_token": 5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, "input_cost_per_token_batches": 2.5e-06, "input_cost_per_token_flex": 2.5e-06, "input_cost_per_token_priority": 1e-05, @@ -23771,6 +23774,7 @@ "mode": "chat", "output_cost_per_token": 3e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_priority": 6e-05, @@ -23806,14 +23810,17 @@ "gpt-5.6-sol": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06, "cache_creation_input_token_cost_flex": 3.125e-06, "cache_creation_input_token_cost_priority": 1.25e-05, "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 1e-06, "input_cost_per_token": 5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, "input_cost_per_token_batches": 2.5e-06, "input_cost_per_token_flex": 2.5e-06, "input_cost_per_token_priority": 1e-05, @@ -23824,6 +23831,7 @@ "mode": "chat", "output_cost_per_token": 3e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_flex": 1.5e-05, "output_cost_per_token_priority": 6e-05, @@ -23859,14 +23867,17 @@ "gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, "cache_creation_input_token_cost_flex": 1.25e-06, "cache_creation_input_token_cost_priority": 5e-06, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 4e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 4e-06, @@ -23877,6 +23888,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, + "output_cost_per_token_above_272k_tokens_flex": 9e-06, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.4e-05, @@ -23912,14 +23924,17 @@ "gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_flex": 1.25e-07, "cache_creation_input_token_cost_priority": 5e-07, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, + "input_cost_per_token_above_272k_tokens_flex": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_token_flex": 1e-07, "input_cost_per_token_priority": 4e-07, @@ -23930,6 +23945,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "output_cost_per_token_above_272k_tokens": 1.8e-06, + "output_cost_per_token_above_272k_tokens_flex": 9e-07, "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, "output_cost_per_token_priority": 2.4e-06, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 1ad271900f6..3454f160cfa 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -661,6 +661,54 @@ def test_generic_cost_per_token_gpt56( assert round(completion_cost, 10) == round(output_cost * completion_tokens, 10) +@pytest.mark.parametrize( + "model,flex_long_input_cost,flex_long_output_cost", + [ + ("gpt-5.6", 5e-6, 2.25e-5), + ("gpt-5.6-sol", 5e-6, 2.25e-5), + ("gpt-5.6-terra", 2e-6, 9e-6), + ("gpt-5.6-luna", 2e-7, 9e-7), + ], +) +def test_generic_cost_per_token_gpt56_flex_above_272k( + model, flex_long_input_cost, flex_long_output_cost +): + """A >272K flex request bills the flex long-context rate, not the standard one. + + Flex long-context is half the standard long-context rate. Without the + ``*_above_272k_tokens_flex`` keys these requests silently fell back to the + standard long-context price, billing 2x what OpenAI charges. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + prompt_tokens = 300000 + completion_tokens = 1000 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="openai", + service_tier="flex", + ) + + assert prompt_cost == pytest.approx(flex_long_input_cost * prompt_tokens) + assert completion_cost == pytest.approx(flex_long_output_cost * completion_tokens) + + standard_long_prompt_cost, standard_long_completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="openai", + service_tier=None, + ) + assert prompt_cost == pytest.approx(standard_long_prompt_cost / 2) + assert completion_cost == pytest.approx(standard_long_completion_cost / 2) + + @pytest.mark.parametrize( "model,input_cost,output_cost,cache_read_cost", [ diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b22e69f0942..83cb75636f2 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -768,11 +768,17 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_creation_input_token_cost_above_1hr": {"type": "number"}, "cache_creation_input_token_cost_above_200k_tokens": {"type": "number"}, "cache_creation_input_token_cost_above_272k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_flex": { + "type": "number" + }, "cache_creation_input_token_cost_flex": {"type": "number"}, "cache_creation_input_token_cost_priority": {"type": "number"}, "cache_read_input_token_cost": {"type": "number"}, "cache_read_input_token_cost_above_200k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_272k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_flex": { + "type": "number" + }, "cache_read_input_token_cost_above_512k_tokens": {"type": "number"}, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": { "type": "number" @@ -806,11 +812,13 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_token_priority": {"type": "number"}, "input_cost_per_token_above_200k_tokens_priority": {"type": "number"}, "input_cost_per_token_above_272k_tokens_priority": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_flex": {"type": "number"}, "input_cost_per_audio_token_priority": {"type": "number"}, "output_cost_per_token_flex": {"type": "number"}, "output_cost_per_token_priority": {"type": "number"}, "output_cost_per_token_above_200k_tokens_priority": {"type": "number"}, "output_cost_per_token_above_272k_tokens_priority": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_flex": {"type": "number"}, "regional_processing_uplift_multiplier_eu": {"type": "number"}, "regional_processing_uplift_multiplier_us": {"type": "number"}, "input_cost_per_pixel": {"type": "number"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index bc3110d4d3d..8f51778ded6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25848,6 +25848,16 @@ export interface components { cache_creation_input_token_cost_above_1hr?: number | null; /** Cache Creation Input Token Cost Above 200K Tokens */ cache_creation_input_token_cost_above_200k_tokens?: number | null; + /** Cache Creation Input Token Cost Above 272K Tokens */ + cache_creation_input_token_cost_above_272k_tokens?: number | null; + /** Cache Creation Input Token Cost Above 272K Tokens Flex */ + cache_creation_input_token_cost_above_272k_tokens_flex?: number | null; + /** Cache Creation Input Token Cost Above 272K Tokens Priority */ + cache_creation_input_token_cost_above_272k_tokens_priority?: number | null; + /** Cache Creation Input Token Cost Flex */ + cache_creation_input_token_cost_flex?: number | null; + /** Cache Creation Input Token Cost Priority */ + cache_creation_input_token_cost_priority?: number | null; /** Cache Read Input Audio Token Cost */ cache_read_input_audio_token_cost?: number | null; /** Cache Read Input Token Cost */ @@ -25858,6 +25868,8 @@ export interface components { cache_read_input_token_cost_above_200k_tokens_priority?: number | null; /** Cache Read Input Token Cost Above 272K Tokens */ cache_read_input_token_cost_above_272k_tokens?: number | null; + /** Cache Read Input Token Cost Above 272K Tokens Flex */ + cache_read_input_token_cost_above_272k_tokens_flex?: number | null; /** Cache Read Input Token Cost Above 272K Tokens Priority */ cache_read_input_token_cost_above_272k_tokens_priority?: number | null; /** Cache Read Input Token Cost Above 512K Tokens */ @@ -25916,6 +25928,8 @@ export interface components { input_cost_per_token_above_200k_tokens_priority?: number | null; /** Input Cost Per Token Above 272K Tokens */ input_cost_per_token_above_272k_tokens?: number | null; + /** Input Cost Per Token Above 272K Tokens Flex */ + input_cost_per_token_above_272k_tokens_flex?: number | null; /** Input Cost Per Token Above 272K Tokens Priority */ input_cost_per_token_above_272k_tokens_priority?: number | null; /** Input Cost Per Token Above 512K Tokens */ @@ -26007,6 +26021,8 @@ export interface components { output_cost_per_token_above_200k_tokens_priority?: number | null; /** Output Cost Per Token Above 272K Tokens */ output_cost_per_token_above_272k_tokens?: number | null; + /** Output Cost Per Token Above 272K Tokens Flex */ + output_cost_per_token_above_272k_tokens_flex?: number | null; /** Output Cost Per Token Above 272K Tokens Priority */ output_cost_per_token_above_272k_tokens_priority?: number | null; /** Output Cost Per Token Above 512K Tokens */ @@ -33939,6 +33955,16 @@ export interface components { cache_creation_input_token_cost_above_1hr?: number | null; /** Cache Creation Input Token Cost Above 200K Tokens */ cache_creation_input_token_cost_above_200k_tokens?: number | null; + /** Cache Creation Input Token Cost Above 272K Tokens */ + cache_creation_input_token_cost_above_272k_tokens?: number | null; + /** Cache Creation Input Token Cost Above 272K Tokens Flex */ + cache_creation_input_token_cost_above_272k_tokens_flex?: number | null; + /** Cache Creation Input Token Cost Above 272K Tokens Priority */ + cache_creation_input_token_cost_above_272k_tokens_priority?: number | null; + /** Cache Creation Input Token Cost Flex */ + cache_creation_input_token_cost_flex?: number | null; + /** Cache Creation Input Token Cost Priority */ + cache_creation_input_token_cost_priority?: number | null; /** Cache Read Input Audio Token Cost */ cache_read_input_audio_token_cost?: number | null; /** Cache Read Input Token Cost */ @@ -33949,6 +33975,8 @@ export interface components { cache_read_input_token_cost_above_200k_tokens_priority?: number | null; /** Cache Read Input Token Cost Above 272K Tokens */ cache_read_input_token_cost_above_272k_tokens?: number | null; + /** Cache Read Input Token Cost Above 272K Tokens Flex */ + cache_read_input_token_cost_above_272k_tokens_flex?: number | null; /** Cache Read Input Token Cost Above 272K Tokens Priority */ cache_read_input_token_cost_above_272k_tokens_priority?: number | null; /** Cache Read Input Token Cost Above 512K Tokens */ @@ -34007,6 +34035,8 @@ export interface components { input_cost_per_token_above_200k_tokens_priority?: number | null; /** Input Cost Per Token Above 272K Tokens */ input_cost_per_token_above_272k_tokens?: number | null; + /** Input Cost Per Token Above 272K Tokens Flex */ + input_cost_per_token_above_272k_tokens_flex?: number | null; /** Input Cost Per Token Above 272K Tokens Priority */ input_cost_per_token_above_272k_tokens_priority?: number | null; /** Input Cost Per Token Above 512K Tokens */ @@ -34098,6 +34128,8 @@ export interface components { output_cost_per_token_above_200k_tokens_priority?: number | null; /** Output Cost Per Token Above 272K Tokens */ output_cost_per_token_above_272k_tokens?: number | null; + /** Output Cost Per Token Above 272K Tokens Flex */ + output_cost_per_token_above_272k_tokens_flex?: number | null; /** Output Cost Per Token Above 272K Tokens Priority */ output_cost_per_token_above_272k_tokens_priority?: number | null; /** Output Cost Per Token Above 512K Tokens */ From 91290c60206e1482bfc8e945ea456161eb935b31 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 15:33:52 -0700 Subject: [PATCH 06/26] fix(policy_engine): only hide config policies behind production DB versions in policies list --- .../proxy/policy_engine/policy_endpoints.py | 9 +++++-- .../policy_engine/test_attachment_registry.py | 11 ++++++++ .../test_policy_engine_endpoints.py | 26 ++++++++++++++++++ .../policy_engine/test_policy_versioning.py | 27 +++++++++++++++++++ 4 files changed, 71 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index 787f7069996..718223da5d8 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -80,7 +80,10 @@ async def list_policies(version_status: Optional[str] = None): List all policies from the database and config.yaml. Optionally filter by version_status. Config-defined policies are returned with definition_location "config" and are treated - as production versions. On a name conflict with a DB policy, only the DB policy is returned. + as production versions. On a name conflict with a production DB policy, only the DB policy + is returned, mirroring runtime resolution where only production DB versions override config. + A draft or published DB version does not hide the config policy, since the config version + is still the one being enforced. Query params: - version_status: Optional. One of "draft", "published", "production". @@ -125,7 +128,9 @@ async def list_policies(version_status: Optional[str] = None): if prisma_client is not None else [] ) - db_policy_names = {db_policy.policy_name for db_policy in db_policies} + db_policy_names = { + db_policy.policy_name for db_policy in db_policies if db_policy.version_status == "production" + } include_config = version_status in (None, "production") config_policies = ( [ diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index cf470c66000..cc231e383a3 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -453,3 +453,14 @@ class TestConfigAttachmentsPreservedAcrossDbSync: await registry.sync_attachments_from_db(_prisma_with_attachment_rows([])) assert len(registry.get_all_attachments()) == 1 + + @pytest.mark.asyncio + async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self): + registry = AttachmentRegistry() + registry.load_attachments([{"policy": "config-policy", "scope": "*"}]) + + registry.clear() + await registry.sync_attachments_from_db(_prisma_with_attachment_rows([])) + + assert registry.get_all_attachments() == [] + assert registry.get_config_attachments() == () diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py b/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py index 78126508c2d..9c486540b3d 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py @@ -133,6 +133,32 @@ class TestListPoliciesIncludesConfig: assert response.policies[0].definition_location == "db" assert response.policies[0].guardrails_add == ["db-guard"] + @pytest.mark.asyncio + async def test_draft_db_policy_does_not_hide_enforced_config_policy(self, policy_registry, monkeypatch): + """ + Runtime sync only lets production DB versions override a config policy, + so a draft or published DB version sharing the name must not suppress + the config entry: the config version is still the one being enforced, + and hiding it makes the list API disagree with actual enforcement. + """ + row = _make_policy_row( + policy_id="uuid-1", policy_name="shared-name", version_status="draft", guardrails_add=["db-guard"] + ) + prisma = MagicMock() + prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[row]) + _set_prisma(monkeypatch, prisma) + policy_registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}}) + + response = await policy_endpoints.list_policies() + + assert response.total_count == 2 + config_entry = next(p for p in response.policies if p.definition_location == "config") + assert config_entry.policy_name == "shared-name" + assert config_entry.version_status == "production" + assert config_entry.guardrails_add == ["config-guard"] + db_entry = next(p for p in response.policies if p.definition_location == "db") + assert db_entry.version_status == "draft" + @pytest.mark.asyncio async def test_version_status_filter_excludes_config_policies(self, policy_registry, monkeypatch): row = _make_policy_row(policy_id="uuid-1", policy_name="db-policy", version_status="draft") diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py index 5a840979c1b..ae80373ee23 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py @@ -13,8 +13,10 @@ from litellm.proxy.policy_engine.policy_registry import ( get_policy_registry, ) from litellm.types.proxy.policy_engine import ( + Policy, PolicyCreateRequest, PolicyDBResponse, + PolicyGuardrails, PolicyUpdateRequest, ) @@ -529,3 +531,28 @@ class TestConfigPoliciesPreservedAcrossDbSync: context=None, ) assert resolved.guardrails == ["tooling"] + + @pytest.mark.asyncio + async def test_add_policy_with_config_source_survives_sync(self): + registry = PolicyRegistry() + registry.add_policy( + "late-config-policy", + Policy(guardrails=PolicyGuardrails(add=["tooling"])), + source="config", + ) + + await registry.sync_policies_from_db(_prisma_with_policy_rows([])) + + assert registry.has_policy("late-config-policy") + assert registry.get_source("late-config-policy") == "config" + + @pytest.mark.asyncio + async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self): + registry = PolicyRegistry() + registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}}) + + registry.clear() + await registry.sync_policies_from_db(_prisma_with_policy_rows([])) + + assert not registry.has_policy("config-policy") + assert registry.get_source("config-policy") is None From ec016d1bd86664112561bd1c7f5fd5a0009ada2e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:01:57 -0700 Subject: [PATCH 07/26] fix(policy_engine): decide config policy suppression from fresh db query only --- .../proxy/policy_engine/policy_endpoints.py | 2 +- .../test_policy_engine_endpoints.py | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index 718223da5d8..cff1378c676 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -136,7 +136,7 @@ async def list_policies(version_status: Optional[str] = None): [ _config_policy_to_db_response(policy_name, policy) for policy_name, policy in registry.list_config_policies().items() - if policy_name not in db_policy_names and registry.get_source(policy_name) != "db" + if policy_name not in db_policy_names ] if include_config else [] diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py b/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py index 9c486540b3d..1ca830dc1e6 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_engine_endpoints.py @@ -159,6 +159,33 @@ class TestListPoliciesIncludesConfig: db_entry = next(p for p in response.policies if p.definition_location == "db") assert db_entry.version_status == "draft" + @pytest.mark.asyncio + async def test_stale_registry_provenance_does_not_hide_config_policy(self, policy_registry, monkeypatch): + """ + Another proxy instance can delete or demote the production DB override + between registry syncs. The endpoint's fresh DB query is the source of + truth for conflicts; stale in-memory provenance from the last sync must + not suppress the config entry once no production override exists. + """ + policy_registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}}) + production_row = _make_policy_row(policy_id="uuid-1", policy_name="shared-name", guardrails_add=["db-guard"]) + sync_prisma = MagicMock() + sync_prisma.db.litellm_policytable.find_many = AsyncMock(side_effect=[[production_row], []]) + await policy_registry.sync_policies_from_db(sync_prisma) + assert policy_registry.get_source("shared-name") == "db" + + fresh_prisma = MagicMock() + fresh_prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[]) + _set_prisma(monkeypatch, fresh_prisma) + + response = await policy_endpoints.list_policies() + + assert response.total_count == 1 + entry = response.policies[0] + assert entry.policy_name == "shared-name" + assert entry.definition_location == "config" + assert entry.guardrails_add == ["config-guard"] + @pytest.mark.asyncio async def test_version_status_filter_excludes_config_policies(self, policy_registry, monkeypatch): row = _make_policy_row(policy_id="uuid-1", policy_name="db-policy", version_status="draft") From 899ed6786047ffa4dbf9e6e138777debf9032743 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Thu, 30 Jul 2026 16:42:23 -0700 Subject: [PATCH 08/26] fix(logging): bind litellm_metadata by reference in function_setup so guardrail info reaches spend logs --- litellm/utils.py | 2 +- .../test_litellm_logging.py | 52 +++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 944bb61d5e7..db78cc0af7f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1048,7 +1048,7 @@ def function_setup( if "metadata" in kwargs: litellm_params["metadata"] = kwargs["metadata"] if "litellm_metadata" in kwargs and isinstance(kwargs["litellm_metadata"], dict): - litellm_params["litellm_metadata"] = kwargs["litellm_metadata"].copy() + litellm_params["litellm_metadata"] = kwargs["litellm_metadata"] # For endpoints like /v1/messages that use "litellm_metadata" instead # of "metadata" (to avoid conflicting with provider API metadata fields), # populate litellm_params["metadata"] so callbacks (e.g. Langfuse) that diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index edc257f4c3f..eaa4bd3e3fc 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2992,6 +2992,58 @@ def test_function_setup_litellm_metadata_populates_metadata(): ), "litellm_params['metadata'] should be a copy, not the same object" +def test_function_setup_litellm_metadata_guardrail_writes_visible_after_setup(): + """ + Regression test for LIT-4512: guardrail writes into the request's + "litellm_metadata" bucket that happen AFTER function_setup (the proxy + initializes the logging object before pre-call guardrails run) must be + visible to the logging object and survive merge_litellm_metadata, so + /v1/messages spend logs carry guardrail_information and + applied_guardrails just like /v1/chat/completions. + """ + import litellm + from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + kwargs = { + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "hello"}], + "litellm_call_id": "test-call-id-lit4512", + "litellm_metadata": { + "user_api_key_hash": "sk-hashed-lit4512", + "guardrails": ["pam-ethical-request"], + }, + } + + logging_obj, returned_kwargs = litellm.utils.function_setup( + original_function="anthropic_messages", + rules_obj=litellm.utils.Rules(), + start_time=time.time(), + **kwargs, + ) + + guardrail_entry = { + "guardrail_name": "pam-ethical-request", + "guardrail_mode": "pre_call", + "guardrail_status": "success", + } + _, metadata_bucket = get_or_create_metadata_bucket(returned_kwargs) + metadata_bucket["standard_logging_guardrail_information"] = [guardrail_entry] + metadata_bucket["applied_guardrails"] = ["pam-ethical-request"] + + litellm_params = logging_obj.model_call_details.get("litellm_params", {}) + litellm_metadata = litellm_params.get("litellm_metadata") + assert litellm_metadata is not None + assert litellm_metadata.get("standard_logging_guardrail_information") == [ + guardrail_entry + ], "guardrail writes after function_setup must be visible to the logging object" + assert litellm_metadata.get("applied_guardrails") == ["pam-ethical-request"] + + merged = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params) + assert merged.get("standard_logging_guardrail_information") == [guardrail_entry] + assert merged.get("applied_guardrails") == ["pam-ethical-request"] + + def test_function_setup_metadata_takes_precedence_over_litellm_metadata(): """ Test that when BOTH metadata and litellm_metadata are present (e.g., user sets From 43fad507dec1bf729a35270dc433e2aadc235e9f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:17:32 -0700 Subject: [PATCH 09/26] fix(responses): map all documented in-stream error codes to real HTTP statuses --- litellm/responses/streaming_iterator.py | 45 ++++++++---- .../test_streaming_iterator_error_events.py | 69 +++++++++++++++++++ 2 files changed, 102 insertions(+), 12 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 357a7ecefe6..e85a758269f 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -7,7 +7,8 @@ import traceback import uuid from datetime import datetime from functools import lru_cache -from typing import Any, Dict, List, Literal, Optional +from types import MappingProxyType +from typing import Any, Dict, List, Literal, Mapping, Optional import httpx from openai._streaming import SSEDecoder @@ -48,13 +49,32 @@ def _log_background_task_failure(task: "asyncio.Task[Any]", *, task_name: str) - verbose_logger.error("%s failed: %s", task_name, exception) -_CLIENT_ERROR_CODES: frozenset[str] = frozenset( - ( - "invalid_request_error", - "context_length_exceeded", - "content_policy_violation", - "model_not_found", - ) +_ERROR_CODE_HTTP_STATUS: Mapping[str, int] = MappingProxyType( + { + "server_error": 500, + "rate_limit_exceeded": 429, + "insufficient_quota": 429, + "vector_store_timeout": 504, + "invalid_prompt": 400, + "invalid_image": 400, + "invalid_image_format": 400, + "invalid_base64_image": 400, + "invalid_image_url": 400, + "image_too_large": 400, + "image_too_small": 400, + "image_parse_error": 400, + "image_content_policy_violation": 400, + "invalid_image_mode": 400, + "image_file_too_large": 400, + "unsupported_image_media_type": 400, + "empty_image_file": 400, + "failed_to_download_image": 400, + "image_file_not_found": 400, + "invalid_request_error": 400, + "context_length_exceeded": 400, + "content_policy_violation": 400, + "model_not_found": 400, + } ) @@ -78,12 +98,13 @@ def _error_event_fields(error_obj: object) -> tuple[str, Optional[str], Optional def _status_code_for_error_fields(error_type: Optional[str], error_code: Optional[str]) -> int: - fields = tuple(field for field in (error_type, error_code) if field is not None) + fields = tuple(field for field in (error_code, error_type) if field is not None) if any(field.startswith("rate_limit") or field == "insufficient_quota" for field in fields): return 429 - if any(field in _CLIENT_ERROR_CODES for field in fields): - return 400 - return 500 + return next( + (_ERROR_CODE_HTTP_STATUS[field] for field in fields if field in _ERROR_CODE_HTTP_STATUS), + 500, + ) class BaseResponsesAPIStreamingIterator: diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 1a2dcd0fcb7..3b87246ebdb 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -28,9 +28,11 @@ from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.streaming_iterator import ( + _ERROR_CODE_HTTP_STATUS, BaseResponsesAPIStreamingIterator, ResponsesAPIStreamingIterator, SyncResponsesAPIStreamingIterator, + _status_code_for_error_fields, ) from litellm.types.llms.openai import ( ErrorEvent, @@ -355,3 +357,70 @@ def test_sync_iterator_raises_mid_stream_fallback_on_rate_limit_error_event(): pass assert exc_info.value.status_code == 429 assert isinstance(exc_info.value.original_exception, litellm.APIError) + + +def test_every_openai_sdk_response_error_code_has_explicit_status_mapping(): + from typing import get_args + + from openai.types.responses.response_error import ResponseError + + sdk_codes = set(get_args(ResponseError.model_fields["code"].annotation)) + unmapped = sdk_codes - set(_ERROR_CODE_HTTP_STATUS) + assert unmapped == set(), ( + f"OpenAI SDK ResponseError codes missing from _ERROR_CODE_HTTP_STATUS: {sorted(unmapped)}; " + "classify each new code with an explicit HTTP status instead of letting it default to 500" + ) + + +@pytest.mark.parametrize( + "code,expected_status", + [ + ("server_error", 500), + ("rate_limit_exceeded", 429), + ("insufficient_quota", 429), + ("vector_store_timeout", 504), + ("invalid_prompt", 400), + ("invalid_image", 400), + ("invalid_image_format", 400), + ("invalid_base64_image", 400), + ("invalid_image_url", 400), + ("image_too_large", 400), + ("image_too_small", 400), + ("image_parse_error", 400), + ("image_content_policy_violation", 400), + ("invalid_image_mode", 400), + ("image_file_too_large", 400), + ("unsupported_image_media_type", 400), + ("empty_image_file", 400), + ("failed_to_download_image", 400), + ("image_file_not_found", 400), + ("totally_unknown_future_code", 500), + ], +) +def test_status_code_for_documented_response_error_codes(code: str, expected_status: int): + assert _status_code_for_error_fields(None, code) == expected_status + + +def test_specific_error_code_wins_over_generic_error_type(): + assert _status_code_for_error_fields("server_error", "invalid_image") == 400 + + +def test_maybe_raise_for_response_failed_event_maps_image_code_to_400(): + iterator = _make_iterator() + mock_response_obj = Mock() + mock_response_obj.error = {"code": "image_content_policy_violation", "message": "image rejected"} + chunk = Mock() + chunk.type = "response.failed" + chunk.response = mock_response_obj + with pytest.raises(litellm.APIError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert exc_info.value.status_code == 400 + assert not isinstance(exc_info.value, MidStreamFallbackError) + + +def test_maybe_raise_for_error_event_maps_vector_store_timeout_to_retriable_504(): + iterator = _make_iterator() + chunk = _make_error_chunk("server_error", "vector_store_timeout", "vector store timed out") + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert exc_info.value.status_code == 504 From 35f770f43e10b87483ab28d9f3ad0c67cb0021fd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:23:50 -0700 Subject: [PATCH 10/26] fix(policy_engine): restore config policy immediately when its DB override is removed and keep same-named DB drafts reachable in the UI --- .../proxy/policy_engine/policy_registry.py | 32 +++++++--- .../policy_engine/test_policy_versioning.py | 63 +++++++++++++++++++ .../policies/_components/PolicyTable.test.tsx | 30 +++++++++ .../policies/_components/PolicyTable.tsx | 15 +++-- 4 files changed, 125 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index 9456afaa1b9..7d2b439270a 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -340,7 +340,8 @@ class PolicyRegistry: def remove_policy(self, policy_name: str) -> bool: """ - Remove a policy by name. + Remove a policy by name. If a config-defined policy shares the name, + it is restored immediately instead of waiting for the next DB sync. Args: policy_name: Name of the policy to remove @@ -348,12 +349,18 @@ class PolicyRegistry: Returns: True if policy was removed, False if it didn't exist """ - if policy_name in self._policies: - del self._policies[policy_name] - self._sources = {name: source for name, source in self._sources.items() if name != policy_name} - verbose_proxy_logger.debug(f"Removed policy: {policy_name}") + if policy_name not in self._policies: + return False + config_fallback = self._config_policies.get(policy_name) + if config_fallback is not None: + self._policies[policy_name] = config_fallback + self._sources = {**self._sources, policy_name: "config"} + verbose_proxy_logger.debug(f"Removed policy: {policy_name}; restored config-defined version") return True - return False + del self._policies[policy_name] + self._sources = {name: source for name, source in self._sources.items() if name != policy_name} + verbose_proxy_logger.debug(f"Removed policy: {policy_name}") + return True # ───────────────────────────────────────────────────────────────────────── # Database CRUD Methods @@ -527,10 +534,15 @@ class PolicyRegistry: # Remove from in-memory registry only if this was the production version if version_status == "production": self.remove_policy(policy_name) - result["warning"] = ( - "Production version was deleted. No other version was promoted. " - "Promote another version to production if this policy should remain active." - ) + if self.get_source(policy_name) == "config": + result["warning"] = ( + "Production version was deleted. The config-defined policy with the same name is active again." + ) + else: + result["warning"] = ( + "Production version was deleted. No other version was promoted. " + "Promote another version to production if this policy should remain active." + ) return result except Exception as e: diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py index ae80373ee23..41d856e7baf 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py @@ -556,3 +556,66 @@ class TestConfigPoliciesPreservedAcrossDbSync: assert not registry.has_policy("config-policy") assert registry.get_source("config-policy") is None + + +class TestRemovePolicyRestoresConfigFallback: + """Deleting a same-named DB override must re-activate the config policy immediately, not at the next sync.""" + + def test_remove_policy_restores_config_version_immediately(self): + registry = PolicyRegistry() + registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}}) + registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db") + + assert registry.remove_policy("shared-name") is True + + policy = registry.get_policy("shared-name") + assert policy is not None + assert policy.guardrails.add == ["config-guard"] + assert registry.get_source("shared-name") == "config" + + def test_remove_policy_without_config_fallback_removes_entirely(self): + registry = PolicyRegistry() + registry.add_policy("db-only", Policy(guardrails=PolicyGuardrails(add=["db-guard"]))) + + assert registry.remove_policy("db-only") is True + + assert not registry.has_policy("db-only") + assert registry.get_source("db-only") is None + + def test_remove_missing_policy_returns_false(self): + registry = PolicyRegistry() + + assert registry.remove_policy("missing") is False + + @pytest.mark.asyncio + async def test_delete_production_override_reactivates_config_policy_and_says_so(self): + registry = PolicyRegistry() + registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}}) + registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db") + prisma = MagicMock() + prod_row = _make_row(policy_id="prod-1", policy_name="shared-name", version_status="production") + prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row) + prisma.db.litellm_policytable.delete = AsyncMock() + + result = await registry.delete_policy_from_db(policy_id="prod-1", prisma_client=prisma) + + assert "config" in result["warning"] + policy = registry.get_policy("shared-name") + assert policy is not None + assert policy.guardrails.add == ["config-guard"] + assert registry.get_source("shared-name") == "config" + + @pytest.mark.asyncio + async def test_delete_all_versions_reactivates_config_policy(self): + registry = PolicyRegistry() + registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}}) + registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db") + prisma = MagicMock() + prisma.db.litellm_policytable.delete_many = AsyncMock() + + await registry.delete_all_versions(policy_name="shared-name", prisma_client=prisma) + + assert registry.get_source("shared-name") == "config" + policy = registry.get_policy("shared-name") + assert policy is not None + assert policy.guardrails.add == ["config-guard"] diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.test.tsx index 06c939aa151..9be1bd60ec8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.test.tsx @@ -145,4 +145,34 @@ describe("PolicyTable", () => { await user.click(screen.getByRole("button", { name: /grouped/ })); expect(defaultProps.onViewClick).toHaveBeenCalledWith("prod-id"); }); + + const sameNamedDbDraft: Partial = { + policy_name: "config-policy", + policy_id: "db-draft-id", + version_status: "draft", + version_number: 2, + }; + const configTwin: Partial = { + policy_name: "config-policy", + policy_id: "config-policy", + version_status: "production", + definition_location: "config", + }; + + it("should render a config policy and a same-named DB draft as separate rows", () => { + const policies = [makePolicy(sameNamedDbDraft), makePolicy(configTwin)]; + renderWithProviders(); + expect(screen.getAllByText("config-policy")).toHaveLength(2); + expect(screen.getByText("Config")).toBeInTheDocument(); + }); + + it("should keep a same-named DB draft reachable next to a config policy", async () => { + const user = userEvent.setup(); + const policies = [makePolicy(sameNamedDbDraft), makePolicy(configTwin)]; + renderWithProviders(); + await user.click(screen.getByRole("button", { name: "config-policy" })); + expect(defaultProps.onViewClick).toHaveBeenCalledWith("db-draft-id"); + await user.click(screen.getByTestId("policy-actions-db-draft-id")); + expect(await screen.findByTestId("policy-action-edit")).not.toHaveAttribute("data-disabled"); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx index d6e841c2119..3405ac6b6bb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx @@ -9,16 +9,21 @@ import { Policy } from "@/components/policies/types"; import { getPolicyTableColumns, PolicyRow } from "./PolicyTableColumns"; -/** One row per policy name; primaryPolicy is used for display and for Edit (FlowBuilder loads all versions) */ +/** One row per DB policy name plus one row per config policy, so a config policy never hides same-named DB versions; primaryPolicy is used for display and for Edit (FlowBuilder loads all versions) */ function groupPoliciesByName(policies: Policy[]): PolicyRow[] { - const names = Array.from(new Set(policies.map((policy) => policy.policy_name || "(unnamed)"))); - return names.map((policyName) => { - const versions = policies.filter((policy) => (policy.policy_name || "(unnamed)") === policyName); + const dbPolicies = policies.filter((policy) => policy.definition_location !== "config"); + const names = Array.from(new Set(dbPolicies.map((policy) => policy.policy_name || "(unnamed)"))); + const dbRows = names.map((policyName) => { + const versions = dbPolicies.filter((policy) => (policy.policy_name || "(unnamed)") === policyName); const primary = versions.find((version) => version.version_status === "production") ?? [...versions].sort((a, b) => (b.version_number ?? 0) - (a.version_number ?? 0))[0]; return { policy_name: policyName, primaryPolicy: primary, versionCount: versions.length }; }); + const configRows = policies + .filter((policy) => policy.definition_location === "config") + .map((policy) => ({ policy_name: policy.policy_name || "(unnamed)", primaryPolicy: policy, versionCount: 1 })); + return [...dbRows, ...configRows]; } interface PolicyTableProps { @@ -67,7 +72,7 @@ const PolicyTable: React.FC = ({ row.policy_name} + getRowId={(row) => `${row.primaryPolicy.definition_location ?? "db"}:${row.policy_name}`} sortingMode="client" sorting={sorting} onSortingChange={setSorting} From b42ef469cfb3a87096c455e446bed4d4ff5dd43e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 19:48:35 -0700 Subject: [PATCH 11/26] fix(policy_engine): warn that the config-defined policy reactivates when all DB versions are deleted --- litellm/proxy/policy_engine/policy_registry.py | 12 ++++++++++-- .../proxy/policy_engine/test_policy_versioning.py | 14 +++++++++++++- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index 7d2b439270a..01b88836387 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -1031,12 +1031,20 @@ class PolicyRegistry: prisma_client: The Prisma client instance Returns: - Dict with success message + Dict with "message" and optional "warning" if a config-defined policy took over. """ try: await _policy_table(prisma_client).delete_many(where={"policy_name": policy_name}) self.remove_policy(policy_name) - return {"message": f"All versions of policy '{policy_name}' deleted successfully"} + message = f"All versions of policy '{policy_name}' deleted successfully" + if self.get_source(policy_name) == "config": + return { + "message": message, + "warning": ( + "All DB versions were deleted. The config-defined policy with the same name is active again." + ), + } + return {"message": message} except Exception as e: verbose_proxy_logger.exception(f"Error deleting all versions: {e}") raise Exception(f"Error deleting all versions: {str(e)}") diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py index 41d856e7baf..ebebfde5cd3 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py @@ -613,9 +613,21 @@ class TestRemovePolicyRestoresConfigFallback: prisma = MagicMock() prisma.db.litellm_policytable.delete_many = AsyncMock() - await registry.delete_all_versions(policy_name="shared-name", prisma_client=prisma) + result = await registry.delete_all_versions(policy_name="shared-name", prisma_client=prisma) assert registry.get_source("shared-name") == "config" policy = registry.get_policy("shared-name") assert policy is not None assert policy.guardrails.add == ["config-guard"] + assert "config" in result["warning"] + + async def test_delete_all_versions_without_config_twin_has_no_warning(self): + registry = PolicyRegistry() + registry.add_policy("db-only", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db") + prisma = MagicMock() + prisma.db.litellm_policytable.delete_many = AsyncMock() + + result = await registry.delete_all_versions(policy_name="db-only", prisma_client=prisma) + + assert registry.get_policy("db-only") is None + assert "warning" not in result From 47ebc964eb8e3ed266a3d7a5473299050d262a36 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:55:32 -0700 Subject: [PATCH 12/26] test: patch unified guardrail mapping global instead of loader to fix order-dependent flake --- .../test_passthrough_post_call_guardrails.py | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index a2f7476abd1..470179a0429 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -294,19 +294,22 @@ class TestUnifiedGuardrailCallTypeResolution: response_body = {"candidates": [{"content": {"parts": [{"text": "hello"}]}}]} - with patch( - "litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail.load_guardrail_translation_mappings" - ) as mock_load: - mock_handler_instance = AsyncMock() - mock_handler_instance.process_output_response = AsyncMock( - return_value=response_body - ) - mock_handler_class = MagicMock(return_value=mock_handler_instance) + mock_handler_instance = AsyncMock() + mock_handler_instance.process_output_response = AsyncMock( + return_value=response_body + ) + mock_handler_class = MagicMock(return_value=mock_handler_instance) - from litellm.types.utils import CallTypes - - mock_load.return_value = {CallTypes.pass_through: mock_handler_class} + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( + unified_guardrail as unified_guardrail_module, + ) + from litellm.types.utils import CallTypes + with patch.object( + unified_guardrail_module, + "endpoint_guardrail_translation_mappings", + {CallTypes.pass_through: mock_handler_class}, + ): result = await unified.async_post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, From ed21c2e3023df92556004d1a0852d9b087d9e2a2 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 30 Jul 2026 21:18:33 -0700 Subject: [PATCH 13/26] feat(s3): support SSE-KMS encryption params on both S3 logging paths (#35291) * feat(s3): support SSE-KMS encryption params on both S3 logging paths * fix(s3): ignore non-string SSE config values instead of crashing logger init * Update litellm/integrations/s3.py Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3): invalidate only the mistyped SSE field instead of dropping both --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/s3.py | 44 +++ litellm/integrations/s3_v2.py | 30 ++- tests/test_litellm/integrations/test_s3.py | 156 +++++++++++ tests/test_litellm/integrations/test_s3_v2.py | 250 ++++++++++++++++++ 4 files changed, 468 insertions(+), 12 deletions(-) create mode 100644 tests/test_litellm/integrations/test_s3.py diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index e8252d87572..07bd957b5a3 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -24,6 +24,8 @@ class S3Logger: s3_aws_secret_access_key=None, s3_aws_session_token=None, s3_config=None, + s3_server_side_encryption: str | None = None, + s3_sse_kms_key_id: str | None = None, **kwargs, ): import boto3 @@ -50,11 +52,16 @@ class S3Logger: s3_aws_session_token = litellm.s3_callback_params.get("s3_aws_session_token") s3_config = litellm.s3_callback_params.get("s3_config") s3_path = litellm.s3_callback_params.get("s3_path") + s3_server_side_encryption = litellm.s3_callback_params.get("s3_server_side_encryption") + s3_sse_kms_key_id = litellm.s3_callback_params.get("s3_sse_kms_key_id") # done reading litellm.s3_callback_params s3_use_team_prefix = bool(litellm.s3_callback_params.get("s3_use_team_prefix", False)) self.s3_use_team_prefix = s3_use_team_prefix self.bucket_name = s3_bucket_name self.s3_path = s3_path + self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params( + s3_server_side_encryption, s3_sse_kms_key_id + ) verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}") # Create an S3 client with custom endpoint URL self.s3_client = boto3.client( @@ -136,6 +143,15 @@ class S3Logger: print_verbose(f"\ns3 Logger - Logging payload = {payload_str}") + sse_params = { + key: value + for key, value in { + "ServerSideEncryption": self.s3_server_side_encryption, + "SSEKMSKeyId": self.s3_sse_kms_key_id, + }.items() + if value + } + response = self.s3_client.put_object( Bucket=self.bucket_name, Key=s3_object_key, @@ -144,6 +160,7 @@ class S3Logger: ContentLanguage="en", ContentDisposition=f'inline; filename="{s3_object_download_filename}"', CacheControl="private, immutable, max-age=31536000, s-maxage=0", + **sse_params, ) print_verbose(f"Response from s3:{str(response)}") @@ -155,6 +172,33 @@ class S3Logger: pass +def _validated_sse_value(name: str, value: str | None) -> str | None: + if value is None or isinstance(value, str): + return value + verbose_logger.warning( + f"s3 logging: ignoring {name} because it has invalid type {type(value).__name__}; expected a string" + ) + return None + + +def resolve_sse_params( + server_side_encryption: str | None, + sse_kms_key_id: str | None, +) -> tuple[str | None, str | None]: + valid_sse = _validated_sse_value("s3_server_side_encryption", server_side_encryption) + valid_key_id = _validated_sse_value("s3_sse_kms_key_id", sse_kms_key_id) + algorithm = valid_sse or ("aws:kms" if valid_key_id else None) + if algorithm is None: + return None, None + if valid_key_id and not algorithm.startswith("aws:kms"): + verbose_logger.warning( + f"s3 logging: ignoring s3_sse_kms_key_id because s3_server_side_encryption is {algorithm}; " + "set it to aws:kms to encrypt with the KMS key" + ) + return algorithm, None + return algorithm, valid_key_id + + def get_s3_object_key( s3_path: str, prefix: str, diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 5b953035cfd..7fa78f39460 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -8,13 +8,14 @@ NOTE 1: S3 does not provide a BATCH PUT API endpoint, so we create tasks to uplo import asyncio import time +from collections.abc import Mapping from datetime import datetime from typing import List, Optional, cast import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS -from litellm.integrations.s3 import get_s3_object_key +from litellm.integrations.s3 import get_s3_object_key, resolve_sse_params from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -55,6 +56,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_key_prefix: bool = False, s3_use_virtual_hosted_style: bool = False, s3_server_side_encryption: Optional[str] = None, + s3_sse_kms_key_id: str | None = None, s3_callback_params_override: Optional[dict] = None, **kwargs, ): @@ -94,6 +96,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_key_prefix=s3_use_key_prefix, s3_use_virtual_hosted_style=s3_use_virtual_hosted_style, s3_server_side_encryption=s3_server_side_encryption, + s3_sse_kms_key_id=s3_sse_kms_key_id, ) verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}") @@ -148,6 +151,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_key_prefix: bool = False, s3_use_virtual_hosted_style: bool = False, s3_server_side_encryption: Optional[str] = None, + s3_sse_kms_key_id: str | None = None, params_source: Optional[dict] = None, ): """ @@ -197,10 +201,20 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): bool(params.get("s3_use_virtual_hosted_style", False)) or s3_use_virtual_hosted_style ) - self.s3_server_side_encryption = params.get("s3_server_side_encryption") or s3_server_side_encryption + self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params( + params.get("s3_server_side_encryption") or s3_server_side_encryption, + params.get("s3_sse_kms_key_id") or s3_sse_kms_key_id, + ) return + def _sse_headers(self) -> Mapping[str, str]: + candidates = { + "x-amz-server-side-encryption": self.s3_server_side_encryption, + "x-amz-server-side-encryption-aws-kms-key-id": self.s3_sse_kms_key_id, + } + return {key: value for key, value in candidates.items() if value} + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): await self._async_log_event_base( kwargs=kwargs, @@ -335,11 +349,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): "Content-Language": "en", "Content-Disposition": f'inline; filename="{batch_logging_element.s3_object_download_filename}"', "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", - **( - {"x-amz-server-side-encryption": self.s3_server_side_encryption} - if self.s3_server_side_encryption - else {} - ), + **self._sse_headers(), } req = requests.Request("PUT", url, data=json_string, headers=headers) prepped = req.prepare() @@ -510,11 +520,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): "Content-Language": "en", "Content-Disposition": f'inline; filename="{batch_logging_element.s3_object_download_filename}"', "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", - **( - {"x-amz-server-side-encryption": self.s3_server_side_encryption} - if self.s3_server_side_encryption - else {} - ), + **self._sse_headers(), } req = requests.Request("PUT", url, data=json_string, headers=headers) prepped = req.prepare() diff --git a/tests/test_litellm/integrations/test_s3.py b/tests/test_litellm/integrations/test_s3.py new file mode 100644 index 00000000000..7e997870852 --- /dev/null +++ b/tests/test_litellm/integrations/test_s3.py @@ -0,0 +1,156 @@ +from datetime import datetime +from unittest.mock import MagicMock, patch + +import litellm +from litellm.integrations.s3 import S3Logger + +TEST_KMS_KEY_ARN = "arn:aws:kms:us-east-1:111122223333:key/test-key-id" + + +def _standard_logging_payload() -> dict: + return { + "id": "chatcmpl-test-id", + "metadata": {"user_api_key_team_alias": None}, + } + + +def _log_event_kwargs() -> dict: + return { + "litellm_params": {"metadata": {}}, + "standard_logging_object": _standard_logging_payload(), + } + + +def _run_log_event(callback_params: dict) -> MagicMock: + original = litellm.s3_callback_params + litellm.s3_callback_params = callback_params + try: + with patch("boto3.client") as mock_boto3_client: + mock_s3_client = MagicMock() + mock_boto3_client.return_value = mock_s3_client + logger = S3Logger() + logger.log_event( + kwargs=_log_event_kwargs(), + response_obj={}, + start_time=datetime(2026, 7, 30, 12, 0, 0), + end_time=datetime(2026, 7, 30, 12, 0, 1), + print_verbose=lambda *args, **kwargs: None, + ) + return mock_s3_client + finally: + litellm.s3_callback_params = original + + +def test_put_object_includes_sse_kms_params_when_configured(): + """ + When s3_server_side_encryption and s3_sse_kms_key_id are set in + s3_callback_params, put_object must receive ServerSideEncryption and + SSEKMSKeyId so objects land encrypted with the customer-managed key. + """ + mock_s3_client = _run_log_event( + { + "s3_bucket_name": "test-bucket", + "s3_region_name": "us-east-1", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": TEST_KMS_KEY_ARN, + } + ) + + put_object_kwargs = mock_s3_client.put_object.call_args.kwargs + assert put_object_kwargs["ServerSideEncryption"] == "aws:kms" + assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN + + +def test_put_object_supports_sse_s3_without_key_id(): + """SSE-S3 (AES256) needs only ServerSideEncryption, no key id.""" + mock_s3_client = _run_log_event( + { + "s3_bucket_name": "test-bucket", + "s3_region_name": "us-east-1", + "s3_server_side_encryption": "AES256", + } + ) + + put_object_kwargs = mock_s3_client.put_object.call_args.kwargs + assert put_object_kwargs["ServerSideEncryption"] == "AES256" + assert "SSEKMSKeyId" not in put_object_kwargs + + +def test_put_object_omits_sse_params_by_default(): + """Without SSE config, put_object kwargs must stay unchanged.""" + mock_s3_client = _run_log_event( + { + "s3_bucket_name": "test-bucket", + "s3_region_name": "us-east-1", + } + ) + + put_object_kwargs = mock_s3_client.put_object.call_args.kwargs + assert "ServerSideEncryption" not in put_object_kwargs + assert "SSEKMSKeyId" not in put_object_kwargs + + +def test_put_object_infers_aws_kms_when_only_key_id_set(): + """A key id without an algorithm must infer aws:kms instead of sending an invalid request.""" + mock_s3_client = _run_log_event( + { + "s3_bucket_name": "test-bucket", + "s3_region_name": "us-east-1", + "s3_sse_kms_key_id": TEST_KMS_KEY_ARN, + } + ) + + put_object_kwargs = mock_s3_client.put_object.call_args.kwargs + assert put_object_kwargs["ServerSideEncryption"] == "aws:kms" + assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN + + +def test_put_object_drops_key_id_when_algorithm_is_not_kms(): + """AES256 plus a key id is invalid for S3; the key id must be dropped, not sent.""" + mock_s3_client = _run_log_event( + { + "s3_bucket_name": "test-bucket", + "s3_region_name": "us-east-1", + "s3_server_side_encryption": "AES256", + "s3_sse_kms_key_id": TEST_KMS_KEY_ARN, + } + ) + + put_object_kwargs = mock_s3_client.put_object.call_args.kwargs + assert put_object_kwargs["ServerSideEncryption"] == "AES256" + assert "SSEKMSKeyId" not in put_object_kwargs + + +def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(): + """ + A YAML boolean in s3_server_side_encryption must not crash logger init and + must not discard the valid key id; aws:kms is inferred from the key id. + """ + mock_s3_client = _run_log_event( + { + "s3_bucket_name": "test-bucket", + "s3_region_name": "us-east-1", + "s3_server_side_encryption": True, + "s3_sse_kms_key_id": TEST_KMS_KEY_ARN, + } + ) + + put_object_kwargs = mock_s3_client.put_object.call_args.kwargs + assert put_object_kwargs["ServerSideEncryption"] == "aws:kms" + assert put_object_kwargs["SSEKMSKeyId"] == TEST_KMS_KEY_ARN + + +def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(): + """A mistyped key id (unquoted YAML number) must not disable the valid algorithm.""" + mock_s3_client = _run_log_event( + { + "s3_bucket_name": "test-bucket", + "s3_region_name": "us-east-1", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": 12345, + } + ) + + put_object_kwargs = mock_s3_client.put_object.call_args.kwargs + assert put_object_kwargs["ServerSideEncryption"] == "aws:kms" + assert "SSEKMSKeyId" not in put_object_kwargs diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index f0a33f2ebfc..3977daae92f 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1388,3 +1388,253 @@ def test_s3_server_side_encryption_read_from_callback_params(): assert logger.s3_server_side_encryption == "aws:kms" finally: litellm.s3_callback_params = original + + +@pytest.mark.asyncio +async def test_async_upload_sets_sse_kms_key_id_header_when_configured(): + """ + When s3_sse_kms_key_id is set alongside aws:kms, the PUT must carry + x-amz-server-side-encryption-aws-kms-key-id so objects are encrypted + with the customer-managed KMS key instead of the bucket default. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement + + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_server_side_encryption="aws:kms", + s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id", + ) + + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-sse-kms.json", + payload={"test": "sse-kms"}, + s3_object_download_filename="test-sse-kms.json", + ) + + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put.return_value = response + + await logger.async_upload_data_to_s3(test_element) + + headers = logger.async_httpx_client.put.call_args.kwargs["headers"] + assert headers["x-amz-server-side-encryption"] == "aws:kms" + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == ( + "arn:aws:kms:us-east-1:111122223333:key/test-key-id" + ) + + +def test_sync_upload_sets_sse_kms_key_id_header_when_configured(): + """The sync upload path must carry the same SSE-KMS headers.""" + from unittest.mock import MagicMock + + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement + + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_server_side_encryption="aws:kms", + s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id", + ) + + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-sync-sse-kms.json", + payload={"test": "sync-sse-kms"}, + s3_object_download_filename="test-sync-sse-kms.json", + ) + + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + mock_sync_client = MagicMock() + mock_sync_client.put.return_value = response + + with patch( + "litellm.integrations.s3_v2._get_httpx_client", + return_value=mock_sync_client, + ): + logger.upload_data_to_s3(test_element) + + headers = mock_sync_client.put.call_args.kwargs["headers"] + assert headers["x-amz-server-side-encryption"] == "aws:kms" + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == ( + "arn:aws:kms:us-east-1:111122223333:key/test-key-id" + ) + + +@pytest.mark.asyncio +async def test_async_upload_omits_kms_key_id_header_when_not_configured(): + """SSE without a key id must not emit the KMS key id header.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement + + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_server_side_encryption="AES256", + ) + + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-aes256.json", + payload={"test": "aes256"}, + s3_object_download_filename="test-aes256.json", + ) + + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put.return_value = response + + await logger.async_upload_data_to_s3(test_element) + + headers = logger.async_httpx_client.put.call_args.kwargs["headers"] + assert headers["x-amz-server-side-encryption"] == "AES256" + assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers + + +def test_s3_sse_kms_key_id_read_from_callback_params(): + """s3_sse_kms_key_id can be configured via s3_callback_params.""" + import litellm + + original = litellm.s3_callback_params + litellm.s3_callback_params = { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + } + try: + logger = S3Logger() + assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") + finally: + litellm.s3_callback_params = original + + +@pytest.mark.asyncio +async def test_async_upload_infers_aws_kms_when_only_key_id_set(): + """ + Setting only s3_sse_kms_key_id must not produce an invalid request + (S3 rejects a key id without an algorithm); aws:kms is inferred. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement + + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_sse_kms_key_id="arn:aws:kms:us-east-1:111122223333:key/test-key-id", + ) + + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-kms-only.json", + payload={"test": "kms-only"}, + s3_object_download_filename="test-kms-only.json", + ) + + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put.return_value = response + + await logger.async_upload_data_to_s3(test_element) + + headers = logger.async_httpx_client.put.call_args.kwargs["headers"] + assert headers["x-amz-server-side-encryption"] == "aws:kms" + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == ( + "arn:aws:kms:us-east-1:111122223333:key/test-key-id" + ) + + +def test_s3_sse_kms_key_id_read_from_audit_override_params(): + """The audit-log override path must honor s3_sse_kms_key_id too.""" + import litellm + + original = litellm.s3_callback_params + litellm.s3_callback_params = {"s3_bucket_name": "normal-logs-bucket"} + try: + logger = S3Logger( + s3_callback_params_override={ + "s3_bucket_name": "audit-logs-bucket", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/audit-key-id", + } + ) + assert logger.s3_bucket_name == "audit-logs-bucket" + assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/audit-key-id") + finally: + litellm.s3_callback_params = original + + +def test_kms_key_id_dropped_when_algorithm_is_not_kms(): + """ + AES256 plus a KMS key id is an invalid S3 combination; the key id must be + dropped at init so uploads keep working instead of silently 400ing. + """ + import litellm + + original = litellm.s3_callback_params + litellm.s3_callback_params = { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "AES256", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + } + try: + logger = S3Logger() + assert logger.s3_server_side_encryption == "AES256" + assert logger.s3_sse_kms_key_id is None + finally: + litellm.s3_callback_params = original + + +def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(): + """ + A YAML boolean in s3_server_side_encryption must not crash logger init and + must not discard the valid key id; aws:kms is inferred from the key id. + """ + import litellm + + original = litellm.s3_callback_params + litellm.s3_callback_params = { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": True, + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + } + try: + logger = S3Logger() + assert logger.s3_server_side_encryption == "aws:kms" + assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") + finally: + litellm.s3_callback_params = original + + +def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(): + """A mistyped key id (unquoted YAML number) must not disable the valid algorithm.""" + import litellm + + original = litellm.s3_callback_params + litellm.s3_callback_params = { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": 12345, + } + try: + logger = S3Logger() + assert logger.s3_server_side_encryption == "aws:kms" + assert logger.s3_sse_kms_key_id is None + finally: + litellm.s3_callback_params = original From 1c36f529aa54855c0d698701a7d9410cdb4f8de7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:23:59 -0700 Subject: [PATCH 14/26] fix(pricing): regenerate model prices schema for flex long-context fields --- model_prices_and_context_window.schema.json | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 988f56655fc..882f514b199 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -99,6 +99,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "cache_creation_input_token_cost_above_272k_tokens_flex": { + "type": "number", + "minimum": 0, + "description": "Flex service-tier rate for the same-named base field." + }, "cache_creation_input_token_cost_flex": { "type": "number", "minimum": 0, @@ -133,6 +138,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "cache_read_input_token_cost_above_272k_tokens_flex": { + "type": "number", + "minimum": 0, + "description": "Flex service-tier rate for the same-named base field." + }, "cache_read_input_token_cost_above_272k_tokens_priority": { "type": "number", "minimum": 0, @@ -262,6 +272,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "input_cost_per_token_above_272k_tokens_flex": { + "type": "number", + "minimum": 0, + "description": "Flex service-tier rate for the same-named base field." + }, "input_cost_per_token_above_272k_tokens_priority": { "type": "number", "minimum": 0, @@ -434,6 +449,11 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "output_cost_per_token_above_272k_tokens_flex": { + "type": "number", + "minimum": 0, + "description": "Flex service-tier rate for the same-named base field." + }, "output_cost_per_token_above_272k_tokens_priority": { "type": "number", "minimum": 0, From 18b9e90d123d581c3de238012acf44ad421c598d Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 31 Jul 2026 04:04:04 +0000 Subject: [PATCH 15/26] fix(cost): bill the fast service tier at the priority rate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/llm_cost_calc/utils.py | 22 ++++--- litellm/types/utils.py | 1 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 63 +++++++++++++++++++ 3 files changed, 79 insertions(+), 7 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 85ed0665ebf..fbc06b76c72 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -2,7 +2,8 @@ ## Helper utilities for cost_per_token() from dataclasses import dataclass -from typing import Any, Literal, Optional, Tuple, TypedDict, cast +from types import MappingProxyType +from typing import Any, Literal, Mapping, Optional, Tuple, TypedDict, cast import litellm from litellm._logging import verbose_logger @@ -39,6 +40,14 @@ _VALID_DATA_RESIDENCIES = frozenset(r.value for r in DataResidency) # of being rebuilt for every model_info key on every call. _SERVICE_TIER_SUFFIXES: tuple[str, ...] = tuple(f"_{st.value}" for st in ServiceTier) +_SERVICE_TIER_TO_COST_KEY_SUFFIX: Mapping[str, str] = MappingProxyType( + { + ServiceTier.FLEX.value: ServiceTier.FLEX.value, + ServiceTier.PRIORITY.value: ServiceTier.PRIORITY.value, + ServiceTier.FAST.value: ServiceTier.PRIORITY.value, + } +) + def _get_token_detail_value(details: object, key: str) -> Optional[int]: if isinstance(details, dict): @@ -177,7 +186,7 @@ def _get_service_tier_cost_key(base_key: str, service_tier: Optional[str]) -> st Args: base_key: The base cost key (e.g., "input_cost_per_token") - service_tier: The service tier ("flex", "priority", or None for standard) + service_tier: The service tier ("flex", "priority", "fast", or None for standard) Returns: str: The cost key to use (e.g., "input_cost_per_token_flex" or "input_cost_per_token") @@ -185,12 +194,11 @@ def _get_service_tier_cost_key(base_key: str, service_tier: Optional[str]) -> st if service_tier is None: return base_key - # Only use service tier specific keys for "flex" and "priority" - if service_tier.lower() in [ServiceTier.FLEX.value, ServiceTier.PRIORITY.value]: - return f"{base_key}_{service_tier.lower()}" + suffix = _SERVICE_TIER_TO_COST_KEY_SUFFIX.get(service_tier.lower()) + if suffix is None: + return base_key - # For any other service tier, use standard pricing - return base_key + return f"{base_key}_{suffix}" def _parse_above_token_threshold(key: str) -> float: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 056592bbf93..18991f53e6f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3847,6 +3847,7 @@ class ServiceTier(Enum): AUTO = "auto" FLEX = "flex" PRIORITY = "priority" + FAST = "fast" class DataResidency(Enum): diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3454f160cfa..866f8f71484 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2447,3 +2447,66 @@ def test_generic_cost_per_token_gemini_35_flash_lite(): ) assert prompt_cost == pytest.approx(0.0003) assert completion_cost == pytest.approx(0.00125) + + +def test_fast_service_tier_bills_at_the_priority_rate(_local_model_cost_map): + """Regression: OpenAI's Fast mode replaced Priority Processing and costs 2x standard. + + Before the fix "fast" fell through to standard pricing, so a Fast mode request + was billed at half of what it actually costs.""" + from litellm.types.utils import Usage + + usage = Usage( + prompt_tokens=1_000, + completion_tokens=500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=200), + ) + + standard = generic_cost_per_token( + model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier=None + ) + priority = generic_cost_per_token( + model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="priority" + ) + fast = generic_cost_per_token( + model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="fast" + ) + + expected_prompt = 800 * 1e-05 + 200 * 1e-06 + expected_completion = 500 * 6e-05 + + assert fast == priority + assert fast[0] == pytest.approx(expected_prompt, rel=1e-9) + assert fast[1] == pytest.approx(expected_completion, rel=1e-9) + assert fast[0] == pytest.approx(standard[0] * 2, rel=1e-9) + assert fast[1] == pytest.approx(standard[1] * 2, rel=1e-9) + + +def test_fast_service_tier_is_case_insensitive(_local_model_cost_map): + from litellm.types.utils import Usage + + usage = Usage(prompt_tokens=1_000, completion_tokens=500) + + assert generic_cost_per_token( + model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="FAST" + ) == generic_cost_per_token( + model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="fast" + ) + + +def test_fast_service_tier_matches_priority_above_the_context_threshold(_local_model_cost_map): + """The above-threshold branch resolves its own cost keys, so the alias has to hold there too.""" + from litellm.types.utils import Usage + + usage = Usage(prompt_tokens=300_000, completion_tokens=1_000) + + fast = generic_cost_per_token( + model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="fast" + ) + priority = generic_cost_per_token( + model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="priority" + ) + + assert fast == priority + assert fast[0] == pytest.approx(300_000 * 1e-05, rel=1e-9) + assert fast[1] == pytest.approx(1_000 * 4.5e-05, rel=1e-9) From 3e3a35dabf780776dfecfab8ca6d86a53a072484 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:57:20 -0700 Subject: [PATCH 16/26] test(pricing): cover gpt-5.6 cache-cost plumbing and bedrock_mantle responses billing --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 43 +++++++++++++++++++ ...bedrock_mantle_responses_transformation.py | 33 ++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3454f160cfa..ee5b618e9ea 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -709,6 +709,49 @@ def test_generic_cost_per_token_gpt56_flex_above_272k( assert completion_cost == pytest.approx(standard_long_completion_cost / 2) +@pytest.mark.parametrize( + "service_tier,prompt_tokens,input_rate,cache_write_rate,cache_read_rate", + [ + (None, 100000, 2e-6, 2.5e-6, 2e-7), + ("flex", 100000, 1e-6, 1.25e-6, 1e-7), + ("priority", 100000, 4e-6, 5e-6, 4e-7), + (None, 300000, 4e-6, 5e-6, 4e-7), + ("flex", 300000, 2e-6, 2.5e-6, 2e-7), + ], +) +def test_generic_cost_per_token_gpt56_terra_cache_costs_by_tier_and_context( + service_tier, prompt_tokens, input_rate, cache_write_rate, cache_read_rate +): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + cached_tokens = 50000 + cache_write_tokens = 40000 + text_tokens = prompt_tokens - cached_tokens - cache_write_tokens + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=100, + total_tokens=prompt_tokens + 100, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=cached_tokens, cache_write_tokens=cache_write_tokens + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gpt-5.6-terra", + usage=usage, + custom_llm_provider="openai", + service_tier=service_tier, + ) + + expected_prompt_cost = ( + text_tokens * input_rate + + cached_tokens * cache_read_rate + + cache_write_tokens * cache_write_rate + ) + assert prompt_cost == pytest.approx(expected_prompt_cost) + + @pytest.mark.parametrize( "model,input_cost,output_cost,cache_read_cost", [ diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index e808023debe..bea979aec64 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -1534,6 +1534,39 @@ class TestBedrockMantleResponsesPricing: assert info["output_cost_per_token"] == pytest.approx(output_cost) assert info["max_input_tokens"] == 272000 + @pytest.mark.parametrize( + "model, input_cost, output_cost", + [ + ("openai.gpt-5.6-sol", 5.5e-06, 3.3e-05), + ("openai.gpt-5.6-terra", 2.2e-06, 1.32e-05), + ("openai.gpt-5.6-luna", 2.2e-07, 1.32e-06), + ], + ) + def test_gpt_5_6_responses_call_cost(self, local_cost_map, model, input_cost, output_cost): + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + + input_tokens = 100000 + output_tokens = 10000 + response = ResponsesAPIResponse( + id="resp-1", + created_at=1700000000, + model=model, + output=[], + usage=ResponseAPIUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=input_tokens + output_tokens, + ), + ) + + cost = litellm.completion_cost( + completion_response=response, + model=f"bedrock_mantle/{model}", + custom_llm_provider="bedrock_mantle", + ) + + assert cost == pytest.approx(input_tokens * input_cost + output_tokens * output_cost) + def test_models_registered(self, local_cost_map): assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models From 9545b109b18d2e19e8812e514f7ddda486c2ca53 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:58:23 -0700 Subject: [PATCH 17/26] fix(responses): suppress LIT002 on error-code map frozen by MappingProxyType --- litellm/responses/streaming_iterator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index e85a758269f..dab666ff0d9 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -50,7 +50,7 @@ def _log_background_task_failure(task: "asyncio.Task[Any]", *, task_name: str) - _ERROR_CODE_HTTP_STATUS: Mapping[str, int] = MappingProxyType( - { + { # mutable-ok: immediately frozen by MappingProxyType "server_error": 500, "rate_limit_exceeded": 429, "insufficient_quota": 429, From 7be3ddd0ffbd5de7d504fb1016c999228c964a11 Mon Sep 17 00:00:00 2001 From: Tin Date: Sat, 25 Jul 2026 15:09:51 -0700 Subject: [PATCH 18/26] fix(mcp): recover the tool-name prefix boundary from registered prefixes The gateway publishes a tool as `` and has to recover that boundary on the way back in, to compare a called name against a toolset or allow/deny list and to rebuild the native name sent upstream. Several sites recovered it by cutting at the FIRST separator and others reconstructed it by hand from `MCPServer.name` with a literal `-`, so both disagreed with the prefix the server actually publishes `get_server_prefix` publishes short_prefix, then alias, then server_name, then server_id; it never reads `name`. A server with no alias therefore publishes its hyphen-filled UUID `server_id` as the prefix, and cutting at the first separator leaves most of the UUID glued to the tool name. Every comparison against the stored `(server_id, tool_name)` toolset row then misses: an allowlist denies a tool the list endpoint just advertised, and a disallowed entry stops blocking, which fails open Recover the boundary in one place instead. `match_known_server_prefix` matches a name against the server's registered prefixes, longest first so a prefix that itself contains the separator beats a shorter prefix that is merely its leading segment, and returns None when the name carries none of them. `strip_known_server_prefix` and `is_tool_name_prefixed` both delegate to it, and the sites that receive a wire name call the owner rather than re-deriving the boundary. `split_server_prefix_from_name` stays for the routing pair it was written for, with a docstring saying so The server-level permission checks are the other half. They run after the boundary is already resolved, so their input is bare and the correction there is to derive the wire form rather than strip it back out; stripping a stored entry a second time cuts a boundary the caller already consumed, which breaks a native name that itself opens with the server prefix. Deriving from `get_server_prefix` alone is not enough either, because routing resolves an inbound name against every prefix from `iter_known_server_prefixes`, so enforcement keyed to the published spelling answers for fewer names than are reachable. Turning `LITELLM_USE_SHORT_MCP_TOOL_PREFIX` on republishes every tool under the short ID while an entry stored under the alias stays routable and silently stops being enforced, which is a fail-open on a config nobody edited. `iter_known_tool_name_spellings` yields the bare name plus the wire form under each accepted prefix, and the allow list, the deny list, `allowed_params` and the routing map that `_create_prefixed_tools` builds now all key off that one function, so the set of names enforcement honors and the set routing accepts cannot drift apart `_tool_name_matches` takes the server as a required argument, so a future caller cannot silently fall back to guessing, and it matches against that same spelling set, so `tools/list` hides exactly what dispatch refuses. Answering for fewer spellings in the filter than enforcement honors leaves a blocked tool advertised, which is how the alias-form entry above stayed listed even once the call was refused. The OpenAPI registry lookup builds its key the same way registration does, via `add_server_prefix_to_name` and `get_server_prefix`, because registration used exactly one key; a server whose `name` differs from its published prefix stops missing its own tools --- .../mcp_server/mcp_server_manager.py | 64 +-- .../mcp_server/rest_endpoints.py | 2 +- .../proxy/_experimental/mcp_server/server.py | 50 +- .../proxy/_experimental/mcp_server/utils.py | 66 ++- tests/mcp_tests/test_mcp_server.py | 3 + .../mcp_server/test_mcp_server.py | 259 +++++++++- .../mcp_server/test_mcp_server_manager.py | 468 ++++++++++++++++++ .../mcp_server/test_openapi_tool_auth.py | 4 + .../mcp_server/test_short_mcp_tool_prefix.py | 66 +++ 9 files changed, 902 insertions(+), 80 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index ae1095da336..285e7b26104 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -116,12 +116,13 @@ from litellm.proxy._experimental.mcp_server.utils import ( get_server_prefix, interpolate_headers, is_short_mcp_tool_prefix_enabled, - is_tool_name_prefixed, iter_known_server_prefixes, + iter_known_tool_name_spellings, + match_known_server_prefix, merge_mcp_headers, normalize_server_name, parse_admin_env_vars, - split_server_prefix_from_name, + strip_known_server_prefix, validate_mcp_server_name, ) from litellm.proxy._types import ( @@ -4185,10 +4186,8 @@ class MCPServerManager: # Register every known prefix form (alias, server_name, server_id, # short ID) so call_tool can resolve regardless of which form a # caller / cached client is using. - self.tool_name_to_mcp_server_name_mapping[original_name] = prefix - for known_prefix in iter_known_server_prefixes(server): - qualified = add_server_prefix_to_name(original_name, known_prefix) - self.tool_name_to_mcp_server_name_mapping[qualified] = prefix + for spelling in iter_known_tool_name_spellings(original_name, server): + self.tool_name_to_mcp_server_name_mapping[spelling] = prefix verbose_logger.info(f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}") return prefixed_tools @@ -4261,20 +4260,27 @@ class MCPServerManager: def check_allowed_or_banned_tools(self, tool_name: str, server: MCPServer) -> bool: """ - Check if the tool is allowed or banned for the given server + Check if the tool is allowed or banned for the given server. + + ``tool_name`` is bare: every caller resolves the boundary against the + server's registered prefixes before dispatch (``server.py``'s + ``original_tool_name``, the Responses handler's ``sanitized_tool_name``). + Stored entries are matched by deriving the spellings routing accepts + rather than by stripping the entries, which would cut a second boundary + out of a native name that itself opens with the server prefix. """ from litellm.proxy._experimental.mcp_server.utils import ( server_applies_tool_allowlist, ) + spellings = tuple(iter_known_tool_name_spellings(tool_name, server)) + if server_applies_tool_allowlist(server): if not server.allowed_tools: return False - return tool_name in server.allowed_tools or f"{server.name}-{tool_name}" in server.allowed_tools + return any(spelling in server.allowed_tools for spelling in spellings) if server.disallowed_tools: - return ( - tool_name not in server.disallowed_tools and f"{server.name}-{tool_name}" not in server.disallowed_tools - ) + return all(spelling not in server.disallowed_tools for spelling in spellings) return True def validate_allowed_params(self, tool_name: str, arguments: dict[str, Any], server: MCPServer) -> None: @@ -4282,7 +4288,8 @@ class MCPServerManager: Filter arguments to only include allowed parameters for the given tool. Args: - tool_name: Name of the tool (with or without prefix) + tool_name: Bare tool name, already resolved against the server's + registered prefixes by the caller arguments: Dictionary of arguments to filter server: MCPServer configuration @@ -4292,19 +4299,14 @@ class MCPServerManager: Raises: HTTPException: If allowed_params is configured for this tool but arguments contain disallowed params """ - from litellm.proxy._experimental.mcp_server.utils import ( - split_server_prefix_from_name, - ) - # If no allowed_params configured, return all arguments if not server.allowed_params: return - # Get the unprefixed tool name to match against config - unprefixed_tool_name, _ = split_server_prefix_from_name(tool_name) - - # Check both prefixed and unprefixed tool names - allowed_params_list = server.allowed_params.get(tool_name) or server.allowed_params.get(unprefixed_tool_name) + spellings = iter_known_tool_name_spellings(tool_name, server) + allowed_params_list = next( + (server.allowed_params[name] for name in spellings if name in server.allowed_params), None + ) # If this tool doesn't have allowed_params specified, allow all params if allowed_params_list is None: @@ -4390,8 +4392,11 @@ class MCPServerManager: global_mcp_tool_registry, ) - # Get the tool from the registry - tool = global_mcp_tool_registry.get_tool(f"{server.name}-{tool_name}") + # Registration used add_server_prefix_to_name(base, get_server_prefix(server)), + # and tool_name is the bare base name by the time call_tool reaches here, so + # rebuilding the key the same way reproduces it exactly + registry_key = add_server_prefix_to_name(tool_name, get_server_prefix(server)) + tool = global_mcp_tool_registry.get_tool(registry_key) if tool is None: # Tool not found in registry error_msg = f"OpenAPI tool {tool_name} not found in registry" @@ -5251,7 +5256,7 @@ class MCPServerManager: for tool in tools: # The tool.name here is already prefixed from _get_tools_from_server # Extract original name for mapping - original_name, _ = split_server_prefix_from_name(tool.name) + original_name = strip_known_server_prefix(tool.name, server) self.tool_name_to_mcp_server_name_mapping[original_name] = server.name self.tool_name_to_mcp_server_name_mapping[tool.name] = server.name @@ -5288,13 +5293,10 @@ class MCPServerManager: # If not found and tool name is prefixed, extract the prefix and # match against any known form. - if is_tool_name_prefixed(tool_name, known_server_prefixes=set(prefix_to_server.keys())): - ( - original_tool_name, - server_name_from_prefix, - ) = split_server_prefix_from_name(tool_name) - normalised_prefix = normalize_server_name(server_name_from_prefix) - matched_server = prefix_to_server.get(normalised_prefix) + matched = match_known_server_prefix(tool_name, prefix_to_server.keys()) + if matched is not None: + matched_prefix, original_tool_name = matched + matched_server = prefix_to_server.get(matched_prefix) if matched_server is not None and ( original_tool_name in self.tool_name_to_mcp_server_name_mapping or tool_name in self.tool_name_to_mcp_server_name_mapping diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 9b51513f4ac..1736e34d70c 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -542,7 +542,7 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, ) if allowed_tools_for_server is not None: - tools = [tool for tool in tools if _tool_name_matches(tool.name, allowed_tools_for_server)] + tools = [tool for tool in tools if _tool_name_matches(tool.name, allowed_tools_for_server, server)] return _create_tool_response_objects(tools, server) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index ec07d33f24d..5181f3a2676 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -65,6 +65,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( extract_mcp_tool_result_error_message, get_server_prefix, iter_known_server_prefixes, + iter_known_tool_name_spellings, ) from litellm.proxy._types import ( ProxyException, @@ -1416,34 +1417,32 @@ if MCP_AVAILABLE: return allowed_mcp_servers - def _tool_name_matches(tool_name: str, filter_list: list[str]) -> bool: + def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool: """ Check if a tool name matches any name in the filter list. - Checks both the full tool name and unprefixed version (without server prefix). - This allows users to configure simple tool names regardless of prefixing. - Comparison is case-insensitive to handle OpenAPI operationIds that may be in camelCase. + Matches via the same ``iter_known_tool_name_spellings`` the server-level + permission checks use, so discovery hides exactly what dispatch refuses; + covering fewer spellings here leaves a blocked tool advertised in + ``tools/list``. Comparison is case-insensitive to handle OpenAPI + operationIds that may be in camelCase. Args: tool_name: The tool name to check (may be prefixed like "server-tool_name") filter_list: List of tool names to match against + mcp_server: The server the tool belongs to, whose registered prefixes + locate the boundary exactly. Required: guessing the boundary at + the first separator silently mismatches every tool on a server + whose prefix contains the separator. Returns: - True if the tool name (prefixed or unprefixed) is in the filter list + True if any spelling of the tool name is in the filter list """ - from litellm.proxy._experimental.mcp_server.utils import ( - split_server_prefix_from_name, - ) + filter_list_lower = {f.lower() for f in filter_list} + bare_name = strip_known_server_prefix(tool_name, mcp_server) + spellings = (tool_name, *iter_known_tool_name_spellings(bare_name, mcp_server)) - # Normalize filter list to lowercase for case-insensitive comparison - filter_list_lower = [f.lower() for f in filter_list] - - if tool_name.lower() in filter_list_lower: - return True - - # Check if the unprefixed name is in the list (case-insensitive) - unprefixed_name, _ = split_server_prefix_from_name(tool_name) - return unprefixed_name.lower() in filter_list_lower + return any(spelling.lower() in filter_list_lower for spelling in spellings) def filter_tools_by_allowed_tools( tools: list[MCPTool], @@ -1473,12 +1472,16 @@ if MCP_AVAILABLE: if server_applies_tool_allowlist(mcp_server): if not mcp_server.allowed_tools: return [] - tools_to_return = [tool for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools)] + tools_to_return = [ + tool for tool in tools if _tool_name_matches(tool.name, mcp_server.allowed_tools, mcp_server) + ] # Filter by disallowed_tools (blacklist) if mcp_server.disallowed_tools: tools_to_return = [ - tool for tool in tools_to_return if not _tool_name_matches(tool.name, mcp_server.disallowed_tools) + tool + for tool in tools_to_return + if not _tool_name_matches(tool.name, mcp_server.disallowed_tools, mcp_server) ] return tools_to_return @@ -1498,7 +1501,7 @@ if MCP_AVAILABLE: return tools for tool in tools: - unprefixed, _ = split_server_prefix_from_name(tool.name) + unprefixed = strip_known_server_prefix(tool.name, mcp_server) lookup_key = unprefixed or tool.name if lookup_key in display_name_map: tool.name = display_name_map[lookup_key] @@ -2699,6 +2702,7 @@ if MCP_AVAILABLE: break if mcp_server is not None: server_name = mcp_server.name + original_tool_name = strip_known_server_prefix(name, mcp_server) if requested_server is not None: if mcp_server is not None and mcp_server.server_id != requested_server.server_id: @@ -2716,6 +2720,7 @@ if MCP_AVAILABLE: if mcp_server is None: mcp_server = requested_server server_name = requested_server.name + original_tool_name = strip_known_server_prefix(name, requested_server) # Only enforce server-level permissions when we can resolve a server if server_name: @@ -2887,13 +2892,14 @@ if MCP_AVAILABLE: _request_resolved_auth_headers.reset(_resolved_token) response = CallToolResult(content=cast(Any, local_content), isError=False) - # Try managed MCP server tool (pass the full prefixed name) + # Try managed MCP server tool (the name is bare; the prefix boundary was + # already resolved above against this server's registered prefixes) # Primary and recommended way to use external MCP servers ######################################################### elif mcp_server: response = await _handle_managed_mcp_tool( server_name=server_name, - name=original_tool_name, # Pass the full name (potentially prefixed) + name=original_tool_name, arguments=arguments, user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index afd396adc4c..e8f8c08188e 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -326,8 +326,34 @@ def iter_known_server_prefixes(server: Any) -> Iterator[str]: yield from _emit(server_id) +def iter_known_tool_name_spellings(tool_name: str, server: Any) -> Iterator[str]: + """Yield every name that denotes the bare ``tool_name`` on ``server``. + + The bare name, then its wire spelling under each prefix from + ``iter_known_server_prefixes``. Routing resolves an inbound name against that + whole set, so anything keyed by tool name (the routing map, the allow/deny + lists, ``allowed_params``) must cover it too or it answers for fewer names + than are reachable, which fails open on ``disallowed_tools``. + ``get_server_prefix`` alone covers only the published spelling, and that moves + with the alias and with ``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``. These are + spellings of one tool on one server, so honoring all of them normalizes the + entry rather than widening a grant. + """ + yield tool_name + for prefix in iter_known_server_prefixes(server): + yield add_server_prefix_to_name(tool_name, prefix) + + def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]: - """Return the unprefixed name plus the server name used as prefix.""" + """Return the unprefixed name plus the server name used as prefix. + + Cuts at the FIRST separator, so the two halves are only trustworthy as a + pair: they reassemble into ``prefixed_name`` exactly, which is what makes + this safe for routing. Reading one half on its own is a guess about where the + boundary fell, and that guess is wrong whenever the prefix itself contains + the separator. Callers that compare a half against configuration must use + :func:`match_known_server_prefix` or :func:`strip_known_server_prefix`. + """ if MCP_TOOL_PREFIX_SEPARATOR in prefixed_name: parts = prefixed_name.split(MCP_TOOL_PREFIX_SEPARATOR, 1) if len(parts) == 2: @@ -335,6 +361,27 @@ def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]: return prefixed_name, "" +def match_known_server_prefix(name: str, known_prefixes: Iterable[str]) -> tuple[str, str] | None: + """Return ``(matched_prefix, bare_name)`` when ``name`` carries a known prefix. + + Candidates are normalized and tried LONGEST first, so a prefix that itself + contains :data:`MCP_TOOL_PREFIX_SEPARATOR` (the UUID ``server_id`` used when + a server has no alias, or a legacy hyphenated alias) wins over a shorter + prefix that is merely its leading segment. Returns ``None`` when no candidate + matches, i.e. ``name`` carries none of these prefixes. + """ + candidates = sorted( + {normalize_server_name(prefix) for prefix in known_prefixes if prefix}, + key=len, + reverse=True, + ) + for prefix in candidates: + separator_suffixed = prefix + MCP_TOOL_PREFIX_SEPARATOR + if name.startswith(separator_suffixed): + return prefix, name[len(separator_suffixed) :] + return None + + def strip_known_server_prefix(name: str, server: Optional[Any]) -> str: """Strip ``server``'s registered prefix from a prefixed tool/resource name. @@ -352,11 +399,8 @@ def strip_known_server_prefix(name: str, server: Optional[Any]) -> str: """ if server is None: return split_server_prefix_from_name(name)[0] - for prefix in iter_known_server_prefixes(server): - candidate = normalize_server_name(prefix) + MCP_TOOL_PREFIX_SEPARATOR - if name.startswith(candidate): - return name[len(candidate) :] - return name + matched = match_known_server_prefix(name, iter_known_server_prefixes(server)) + return name if matched is None else matched[1] def is_tool_name_prefixed( @@ -367,15 +411,16 @@ def is_tool_name_prefixed( Check if tool name has a known MCP server prefix. When ``known_server_prefixes`` is provided the function verifies that the - substring before the first separator is an actual registered server - prefix. Without it the check falls back to the legacy heuristic + name actually starts with one of those prefixes followed by the separator, + matching the longest candidate first so a prefix containing the separator + still resolves. Without it the check falls back to the legacy heuristic (separator present anywhere in the name), which can produce false positives for non-MCP tools whose names contain hyphens (e.g. ``text-to-speech``, ``code-review``). Args: tool_name: Tool name to check. - known_server_prefixes: Optional set of normalised server prefixes + known_server_prefixes: Optional set of normalized server prefixes currently registered in the MCP manager. Pass this whenever the caller has access to the server registry so that the check is accurate. @@ -387,8 +432,7 @@ def is_tool_name_prefixed( return False if known_server_prefixes is not None: - candidate_prefix = tool_name.split(MCP_TOOL_PREFIX_SEPARATOR, 1)[0] - return normalize_server_name(candidate_prefix) in known_server_prefixes + return match_known_server_prefix(tool_name, known_server_prefixes) is not None # Legacy fallback – separator present somewhere in the name. return True diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index f1e6539439a..390f6917573 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -2010,6 +2010,9 @@ async def test_get_tools_for_single_server_applies_disallowed_tools_without_allo mock_server.mcp_info = {"server_name": "zapier"} mock_server.name = "zapier" mock_server.server_id = "zapier" + mock_server.server_name = "zapier" + mock_server.alias = None + mock_server.short_prefix = None mock_server.allowed_tools = None mock_server.disallowed_tools = ["send_email"] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index c0affdf46b3..a3577c85621 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1029,6 +1029,9 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): working_server.server_name = "working_server" working_server.auth_type = None working_server.extra_headers = None + working_server.short_prefix = None + working_server.tool_name_to_display_name = None + working_server.tool_name_to_description = None failing_server = MagicMock() failing_server.name = "failing_server" @@ -1039,6 +1042,9 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): failing_server.server_name = "failing_server" failing_server.auth_type = None failing_server.extra_headers = None + failing_server.short_prefix = None + failing_server.tool_name_to_display_name = None + failing_server.tool_name_to_description = None # Mock global_mcp_server_manager mock_manager = MagicMock() @@ -4507,28 +4513,36 @@ def test_tool_name_matches_case_insensitive(): except ImportError: pytest.skip("MCP server not available") + server = MCPServer( + server_id="srv-per-store", + name="per_store", + server_name="per_store", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + ) + # Test case 1: Unprefixed tool name with camelCase in filter list - assert _tool_name_matches("addpet", ["addPet", "updatePet"]) is True - assert _tool_name_matches("updatepet", ["addPet", "updatePet"]) is True - assert _tool_name_matches("deletepet", ["addPet", "updatePet"]) is False + assert _tool_name_matches("addpet", ["addPet", "updatePet"], server) is True + assert _tool_name_matches("updatepet", ["addPet", "updatePet"], server) is True + assert _tool_name_matches("deletepet", ["addPet", "updatePet"], server) is False # Test case 2: Prefixed tool name with camelCase in filter list - assert _tool_name_matches("per_store-addpet", ["addPet", "updatePet"]) is True - assert _tool_name_matches("per_store-updatepet", ["addPet", "updatePet"]) is True - assert _tool_name_matches("per_store-deletepet", ["addPet", "updatePet"]) is False + assert _tool_name_matches("per_store-addpet", ["addPet", "updatePet"], server) is True + assert _tool_name_matches("per_store-updatepet", ["addPet", "updatePet"], server) is True + assert _tool_name_matches("per_store-deletepet", ["addPet", "updatePet"], server) is False # Test case 3: Mixed case variations - assert _tool_name_matches("findPetsByStatus", ["findpetsbystatus"]) is True - assert _tool_name_matches("findpetsbystatus", ["findPetsByStatus"]) is True - assert _tool_name_matches("FINDPETSBYSTATUS", ["findPetsByStatus"]) is True + assert _tool_name_matches("findPetsByStatus", ["findpetsbystatus"], server) is True + assert _tool_name_matches("findpetsbystatus", ["findPetsByStatus"], server) is True + assert _tool_name_matches("FINDPETSBYSTATUS", ["findPetsByStatus"], server) is True # Test case 4: Full prefixed name in filter list (case-insensitive) - assert _tool_name_matches("server-addPet", ["server-addpet"]) is True - assert _tool_name_matches("server-addpet", ["server-addPet"]) is True + assert _tool_name_matches("server-addPet", ["server-addpet"], server) is True + assert _tool_name_matches("server-addpet", ["server-addPet"], server) is True # Test case 5: Ensure non-matching names still don't match - assert _tool_name_matches("addpet", ["deletePet", "updatePet"]) is False - assert _tool_name_matches("server-addpet", ["deletePet", "updatePet"]) is False + assert _tool_name_matches("addpet", ["deletePet", "updatePet"], server) is False + assert _tool_name_matches("server-addpet", ["deletePet", "updatePet"], server) is False def test_filter_tools_by_allowed_tools_case_insensitive(): @@ -4581,6 +4595,7 @@ def test_filter_tools_by_allowed_tools_case_insensitive(): server = MCPServer( server_id="test-server", name="per_store", + server_name="per_store", transport=MCPTransport.http, allowed_tools=["addPet", "updatePet", "findPetsByStatus"], ) @@ -6121,6 +6136,70 @@ async def test_execute_mcp_tool_rest_server_id_authoritative_for_unprefixed_tool assert captured["name"] == "echo" +@pytest.mark.asyncio +async def test_execute_mcp_tool_strips_a_prefix_that_contains_the_separator(): + """A server with no alias publishes its UUID server_id as the tool prefix. + + Splitting that wire name at the first separator leaves a truncated UUID tail + glued to the tool name, which then travels to the upstream server as the tool + to call, into the spend log, and into the server-level allowed_tools check. + """ + from mcp.types import TextContent + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + server_id = "117c814c-1a2b-4c4d-8e8f-0a1b2c3d4e5f" + alias_less_server = MCPServer( + server_id=server_id, + name=server_id, + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + authentication_token="abc123", + ) + + captured: dict = {} + + async def fake_handle_managed_mcp_tool(**kwargs): + captured.update(kwargs) + return mcp_module.CallToolResult( + content=[TextContent(type="text", text="ok")], + isError=False, + ) + + with ( + patch.object( + mcp_module.global_mcp_server_manager, + "_get_mcp_server_from_tool_name", + return_value=alias_less_server, + ), + patch.object( + mcp_module, + "_handle_managed_mcp_tool", + new=fake_handle_managed_mcp_tool, + ), + patch.object( + mcp_module.MCPRequestHandler, + "is_tool_allowed", + return_value=True, + ), + patch.object( + mcp_module.global_mcp_tool_registry, + "get_tool", + return_value=None, + ), + ): + await mcp_module.execute_mcp_tool( + name=f"{server_id}-read_wiki_contents", + arguments={"repoName": "acme/wiki"}, + allowed_mcp_servers=[alias_less_server], + start_time=datetime.now(), + ) + + assert captured["server_name"] == server_id + assert captured["name"] == "read_wiki_contents" + + @pytest.mark.asyncio async def test_execute_mcp_tool_rest_server_id_injects_requested_server_credentials(): """REST server_id must inject the requested server's auth, not a URL-collision peer's.""" @@ -6426,6 +6505,8 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details(): fake_server.mcp_info = None fake_server.server_id = "srv-1" fake_server.server_name = "openapi-petstore" + fake_server.alias = None + fake_server.short_prefix = None fake_tool = MagicMock() fake_tool.name = "list_pets" @@ -6481,7 +6562,13 @@ async def test_execute_mcp_tool_sets_model_in_model_call_details(): @pytest.mark.asyncio async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requested_server(): - """A prefixed REST name that resolves to no tool must still dispatch to the server_id.""" + """A prefixed REST name that resolves to no tool must still dispatch to the server_id. + + The prefix here belongs to a different server, so it is not a prefix boundary on the + routed server and the name travels upstream whole. Stripping it would invoke the routed + server's similarly named tool instead, which is what the tool_server_mismatch 403 exists + to prevent when the prefix does resolve. + """ from mcp.types import TextContent from litellm.proxy._experimental.mcp_server import server as mcp_module @@ -6553,7 +6640,7 @@ async def test_execute_mcp_tool_rest_unresolved_prefixed_name_routes_to_requeste ) assert captured["server_name"] == "rest_target" - assert captured["name"] == "list_things" + assert captured["name"] == "known_prefix-list_things" routed_server = { requested_server.name: requested_server, @@ -7587,22 +7674,28 @@ async def test_aggregate_listing_reports_per_server_outcomes(): working_server = MagicMock() working_server.name = "working_server" working_server.alias = "working" + working_server.short_prefix = None working_server.allowed_tools = None working_server.disallowed_tools = None working_server.server_id = "working_server" working_server.server_name = "working_server" working_server.auth_type = None working_server.extra_headers = None + working_server.tool_name_to_display_name = None + working_server.tool_name_to_description = None broken_server = MagicMock() broken_server.name = "broken_server" broken_server.alias = "broken" + broken_server.short_prefix = None broken_server.allowed_tools = None broken_server.disallowed_tools = None broken_server.server_id = "broken_server" broken_server.server_name = "broken_server" broken_server.auth_type = None broken_server.extra_headers = None + broken_server.tool_name_to_display_name = None + broken_server.tool_name_to_description = None mock_manager = MagicMock() mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["working_server", "broken_server"]) @@ -7924,3 +8017,139 @@ async def test_post_mcp_call_guardrails_propagate_a_block(): user_api_key_auth=None, request_data={}, ) + + +class TestListFiltersHonorThePrefixBoundary: + """The listing filters compare a published (prefixed) tool name against + configured entries, so they have to locate the boundary with the server's + registered prefixes. An alias-less server publishes its UUID server_id as + the prefix, and that prefix contains the separator, so cutting at the first + separator drops every tool on the server from the listing. + """ + + SERVER_ID = "117c814c-1a2b-4c4d-8e8f-0a1b2c3d4e5f" + + @staticmethod + def _alias_less_server(**overrides): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServer + + return MCPServer( + server_id=TestListFiltersHonorThePrefixBoundary.SERVER_ID, + name=TestListFiltersHonorThePrefixBoundary.SERVER_ID, + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + **overrides, + ) + + def _published_tools(self, *bare_names: str): + from mcp.types import Tool as MCPTool + + return [ + MCPTool(name=f"{self.SERVER_ID}-{bare}", description=bare, inputSchema={"type": "object"}) + for bare in bare_names + ] + + def test_bare_allowlist_entry_keeps_the_published_tool(self): + from litellm.proxy._experimental.mcp_server.server import ( + filter_tools_by_allowed_tools, + ) + + server = self._alias_less_server(allowed_tools=["read_wiki_contents"]) + tools = self._published_tools("read_wiki_contents", "read_wiki_structure") + + kept = filter_tools_by_allowed_tools(tools, server) + + assert [tool.name for tool in kept] == [f"{self.SERVER_ID}-read_wiki_contents"] + + def test_bare_blocklist_entry_excludes_the_published_tool(self): + from litellm.proxy._experimental.mcp_server.server import ( + filter_tools_by_allowed_tools, + ) + + server = self._alias_less_server(disallowed_tools=["read_wiki_structure"]) + tools = self._published_tools("read_wiki_contents", "read_wiki_structure") + + kept = filter_tools_by_allowed_tools(tools, server) + + assert [tool.name for tool in kept] == [f"{self.SERVER_ID}-read_wiki_contents"] + + def test_unrelated_entry_does_not_match(self): + from litellm.proxy._experimental.mcp_server.server import _tool_name_matches + + server = self._alias_less_server() + + assert not _tool_name_matches(f"{self.SERVER_ID}-read_wiki_contents", ["read_wiki_structure"], server) + + def test_match_is_still_case_insensitive(self): + from litellm.proxy._experimental.mcp_server.server import _tool_name_matches + + server = self._alias_less_server() + + assert _tool_name_matches(f"{self.SERVER_ID}-findPetsByStatus", ["findpetsbystatus"], server) + + def test_alias_form_entry_matches_a_tool_published_under_the_short_prefix(self, monkeypatch): + # Routing accepts the alias form, so an entry stored before short + # prefixes were turned on still governs the tool. Matching only the + # published spelling left it advertised while dispatch refused it. + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._experimental.mcp_server.server import _tool_name_matches + + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = MCPServer( + server_id=self.SERVER_ID, + name="deepwiki_cfg", + alias="deepwiki_cfg", + short_prefix="eiG", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + ) + + assert _tool_name_matches("eiG-read_wiki_contents", ["deepwiki_cfg-read_wiki_contents"], server) + + def test_discovery_hides_exactly_what_dispatch_refuses(self, monkeypatch): + """Both halves of one decision, driven through both production paths. + + A spelling the blocklist enforces but the filter misses leaves a blocked + tool advertised; the reverse hides a tool that would have been callable. + """ + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServer, + MCPServerManager, + ) + from litellm.proxy._experimental.mcp_server.server import ( + filter_tools_by_allowed_tools, + ) + + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + + def _server(**overrides): + return MCPServer( + server_id=self.SERVER_ID, + name="deepwiki_prod", + alias="deepwiki", + server_name="deepwiki_prod", + short_prefix="eiG", + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + **overrides, + ) + + manager = MCPServerManager() + manager._create_prefixed_tools( + [MCPTool(name="read_wiki_contents", description="", inputSchema={"type": "object"})], + _server(), + ) + registered = sorted(manager.tool_name_to_mcp_server_name_mapping) + assert len(registered) > 1 + + published = MCPTool(name="eiG-read_wiki_contents", description="", inputSchema={"type": "object"}) + for spelling in registered: + server = _server(disallowed_tools=[spelling]) + + refused = not manager.check_allowed_or_banned_tools("read_wiki_contents", server) + hidden = filter_tools_by_allowed_tools([published], server) == [] + + assert refused, spelling + assert hidden, spelling 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 8a8dea0ba28..2977f13caf1 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 @@ -9189,3 +9189,471 @@ class TestDiscoveryFailureLogging: assert "typo_row" in caplog.text assert "authorization_url, token_url" in caplog.text assert "unresolved" in caplog.text + + +def _unrestricted_auth() -> MagicMock: + """A caller with no object_permission, so only server-level checks apply.""" + user_api_key_auth = MagicMock() + user_api_key_auth.object_permission = None + user_api_key_auth.object_permission_id = None + return user_api_key_auth + + +def _permissive_proxy_logging() -> MagicMock: + proxy_logging_obj = MagicMock() + proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(return_value={}) + proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={}) + proxy_logging_obj.pre_call_hook = AsyncMock(return_value={}) + return proxy_logging_obj + + +ALIAS_LESS_SERVER_ID = "117c814c-1a2b-4c4d-8e8f-0a1b2c3d4e5f" + + +class TestServerToolListsHonorThePrefixBoundary: + """The server-level allowed_tools / disallowed_tools / allowed_params checks + receive a BARE tool name. Every caller resolves the prefix boundary before + dispatch (``server.py``'s ``original_tool_name``, the Responses handler's + ``sanitized_tool_name``), and ``call_tool`` hands that same value to the + upstream client verbatim, which only works because it carries no prefix. + + Stored entries may also carry a prefix, and routing accepts *every* prefix + from ``iter_known_server_prefixes`` (short ID, alias, server_name, + server_id), so enforcement derives that whole set via + ``iter_known_tool_name_spellings``. Rebuilding a single comparand as + ``f"{server.name}-{tool_name}"`` used a field the prefix chain never reads + (``get_server_prefix`` is short_prefix, then alias, then server_name, then + server_id) and hardcoded the separator; deriving only ``get_server_prefix`` + covers just the currently published spelling. Either way the check answers + for fewer spellings than are reachable, and on the blocklist arm that is a + fail-open. + """ + + async def _run_check(self, server: MCPServer, name: str, arguments: dict[str, Any] | None = None) -> None: + await MCPServerManager().pre_call_tool_check( + name=name, + arguments=arguments if arguments is not None else {}, + server_name=server.name, + user_api_key_auth=_unrestricted_auth(), + proxy_logging_obj=_permissive_proxy_logging(), + server=server, + ) + + @staticmethod + def _aliased_server(**overrides: Any) -> MCPServer: + return MCPServer( + server_id="dd7f2b9e-2c4a-4f1b-9e0a-8d3c6b5a4f21", + name="petstore_prod", + alias="petstore", + server_name="petstore_prod", + url="https://petstore.example.com/mcp", + transport=MCPTransport.http, + **overrides, + ) + + @staticmethod + def _alias_less_server(**overrides: Any) -> MCPServer: + # No alias and no server_name, so the published prefix is the UUID + # server_id, which itself contains the prefix separator. + return MCPServer( + server_id=ALIAS_LESS_SERVER_ID, + name=ALIAS_LESS_SERVER_ID, + url="https://wiki.example.com/mcp", + transport=MCPTransport.http, + **overrides, + ) + + @pytest.mark.asyncio + async def test_allowlist_entry_prefixed_with_the_alias_matches_a_bare_call(self): + # The dashboard shows tools under the published prefix, so admins store + # "petstore-getpetbyid"; the display name "petstore_prod" is not it. + server = self._aliased_server(allowed_tools=["petstore-getpetbyid"]) + + await self._run_check(server, "getpetbyid") + + @pytest.mark.asyncio + async def test_bare_allowlist_entry_matches_on_an_alias_less_server(self): + server = self._alias_less_server(allowed_tools=["read_wiki_contents"]) + + await self._run_check(server, "read_wiki_contents") + + @pytest.mark.asyncio + async def test_wire_form_allowlist_entry_matches_on_an_alias_less_server(self): + # The published prefix is the UUID server_id, so it contains the + # separator; the derived wire form has to reproduce it whole. + server = self._alias_less_server(allowed_tools=[f"{ALIAS_LESS_SERVER_ID}-read_wiki_contents"]) + + await self._run_check(server, "read_wiki_contents") + + @pytest.mark.asyncio + async def test_wire_form_entry_matches_a_native_name_that_opens_with_the_prefix(self): + # "petstore-getpetbyid" is a real upstream tool name here, so its wire + # form is "petstore-petstore-getpetbyid". Stripping the stored entry + # instead of deriving the wire form cut a boundary the caller had already + # consumed, leaving asymmetric operands that denied a permitted call. + server = self._aliased_server(allowed_tools=["petstore-petstore-getpetbyid"]) + + await self._run_check(server, "petstore-getpetbyid") + + @pytest.mark.asyncio + async def test_wire_form_blocklist_entry_blocks_a_native_name_that_opens_with_the_prefix(self): + server = self._aliased_server(disallowed_tools=["petstore-petstore-getpetbyid"]) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "petstore-getpetbyid") + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_wire_form_blocklist_entry_blocks_under_the_short_prefix_mode(self, monkeypatch): + # short_prefix wins in get_server_prefix but is never server.name, so the + # hand-built comparand could not match a stored wire-form entry and the + # blocklisted tool stayed callable. + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = self._aliased_server(short_prefix="F3X", disallowed_tools=["F3X-deletepet"]) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "deletepet") + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_alias_form_blocklist_entry_still_blocks_under_the_short_prefix_mode(self, monkeypatch): + # Turning short prefixes on republishes every tool under the short ID, + # but routing still resolves the alias form, so an entry an admin stored + # before the flip stays reachable and has to stay enforced. Deriving only + # the published spelling silently stops honoring it: a fail-open on a + # config nobody edited. + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = self._aliased_server(short_prefix="F3X", disallowed_tools=["petstore-deletepet"]) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "deletepet") + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_alias_form_allowlist_entry_still_matches_under_the_short_prefix_mode(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = self._aliased_server(short_prefix="F3X", allowed_tools=["petstore-getpetbyid"]) + + await self._run_check(server, "getpetbyid") + + @pytest.mark.asyncio + async def test_server_name_form_blocklist_entry_still_blocks_under_the_short_prefix_mode(self, monkeypatch): + # server_name sits third in the prefix chain, so it is published only + # when alias and short_prefix are both absent, yet routing accepts it + # regardless. + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = self._aliased_server(short_prefix="F3X", disallowed_tools=["petstore_prod-deletepet"]) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "deletepet") + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_raw_server_id_form_blocklist_entry_still_blocks_under_the_short_prefix_mode(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = self._alias_less_server(disallowed_tools=[f"{ALIAS_LESS_SERVER_ID}-read_wiki_contents"]) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "read_wiki_contents") + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_allowed_params_keyed_by_the_alias_form_are_enforced_under_the_short_prefix_mode(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = self._aliased_server(short_prefix="F3X", allowed_params={"petstore-getpetbyid": ["petid"]}) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check( + server, + "getpetbyid", + arguments={"petid": "7", "include_internal": "true"}, + ) + + assert exc_info.value.status_code == 403 + assert "include_internal" in exc_info.value.detail["error"] + + @pytest.mark.asyncio + async def test_foreign_prefix_entry_does_not_match_under_the_short_prefix_mode(self, monkeypatch): + # Honoring every known prefix must not become "honor any prefix": the + # widened set is this server's spellings only. + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = self._aliased_server(short_prefix="F3X", allowed_tools=["other_server-getpetbyid"]) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "getpetbyid") + + assert exc_info.value.status_code == 403 + + @pytest.mark.parametrize("short_prefix_mode", [False, True]) + @pytest.mark.asyncio + async def test_every_spelling_routing_registers_is_also_enforced(self, monkeypatch, short_prefix_mode): + """The invariant, driven through production code on both sides. + + ``_create_prefixed_tools`` decides which spellings reach dispatch, so + every key it registers has to be a spelling the blocklist can refuse. + Any key routing accepts but enforcement misses is a callable blocked + tool. + """ + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true" if short_prefix_mode else "false") + shape = self._aliased_server(short_prefix="F3X") + + manager = MCPServerManager() + manager._create_prefixed_tools([MCPTool(name="deletepet", description="", inputSchema={})], shape) + registered = sorted(manager.tool_name_to_mcp_server_name_mapping) + assert len(registered) > 1 + + for spelling in registered: + server = self._aliased_server(short_prefix="F3X", disallowed_tools=[spelling]) + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "deletepet") + assert exc_info.value.status_code == 403, spelling + + @pytest.mark.asyncio + async def test_wire_form_allowlist_entry_follows_a_non_default_separator(self): + from litellm.proxy._experimental.mcp_server import utils as mcp_utils + + server = self._aliased_server(allowed_tools=["petstore__getpetbyid"]) + + with patch.object(mcp_utils, "MCP_TOOL_PREFIX_SEPARATOR", "__"): + await self._run_check(server, "getpetbyid") + + @pytest.mark.asyncio + async def test_tool_outside_the_allowlist_is_still_denied(self): + server = self._aliased_server(allowed_tools=["petstore-getpetbyid"]) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "deletepet") + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_allowlist_entry_prefixed_for_another_server_does_not_match(self): + # Reducing both sides must not widen the allowlist across servers: a + # foreign prefix is not one of this server's known prefixes, so the + # entry keeps it and never collapses onto a bare name. + server = self._aliased_server(allowed_tools=["other_server-getpetbyid"]) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "getpetbyid") + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_prefixed_disallowed_entry_blocks_a_bare_call(self): + # Fail-open regression: the blocklist arm answered "not banned" whenever + # the stored entry carried a prefix it failed to reconstruct. + server = self._aliased_server(disallowed_tools=["petstore-deletepet"]) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "deletepet") + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_tool_outside_the_blocklist_is_still_allowed(self): + server = self._aliased_server(disallowed_tools=["petstore-deletepet"]) + + await self._run_check(server, "getpetbyid") + + @pytest.mark.asyncio + async def test_allowed_params_are_enforced_for_a_bare_key(self): + server = self._alias_less_server(allowed_params={"read_wiki_contents": ["repo"]}) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check( + server, + "read_wiki_contents", + arguments={"repo": "acme/wiki", "internal_only": "true"}, + ) + + assert exc_info.value.status_code == 403 + assert "internal_only" in exc_info.value.detail["error"] + + @pytest.mark.asyncio + async def test_allowed_params_are_enforced_for_a_wire_form_key(self): + # A key stored under the published prefix matched nothing, so the lookup + # returned None and the check silently allowed every parameter instead of + # enforcing the configured list. + server = self._alias_less_server(allowed_params={f"{ALIAS_LESS_SERVER_ID}-read_wiki_contents": ["repo"]}) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check( + server, + "read_wiki_contents", + arguments={"repo": "acme/wiki", "internal_only": "true"}, + ) + + assert exc_info.value.status_code == 403 + assert "internal_only" in exc_info.value.detail["error"] + + @pytest.mark.asyncio + async def test_allowed_params_still_accept_the_configured_parameters(self): + server = self._alias_less_server(allowed_params={f"{ALIAS_LESS_SERVER_ID}-read_wiki_contents": ["repo"]}) + + await self._run_check(server, "read_wiki_contents", arguments={"repo": "acme/wiki"}) + + @pytest.mark.asyncio + async def test_allowed_params_are_enforced_for_a_native_name_that_opens_with_the_prefix(self): + server = self._aliased_server(allowed_params={"petstore-petstore-getpetbyid": ["petid"]}) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check( + server, + "petstore-getpetbyid", + arguments={"petid": "7", "include_internal": "true"}, + ) + + assert exc_info.value.status_code == 403 + assert "include_internal" in exc_info.value.detail["error"] + + @pytest.mark.asyncio + async def test_an_explicitly_empty_allowed_params_list_refuses_every_parameter(self): + server = self._alias_less_server(allowed_params={"read_wiki_contents": []}) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "read_wiki_contents", arguments={"repo": "acme/wiki"}) + + assert exc_info.value.status_code == 403 + assert "repo" in exc_info.value.detail["error"] + + @pytest.mark.asyncio + async def test_an_explicitly_empty_allowed_params_list_still_permits_an_argument_free_call(self): + server = self._alias_less_server(allowed_params={"read_wiki_contents": []}) + + await self._run_check(server, "read_wiki_contents", arguments={}) + + +class TestOpenAPIRegistryKeyMatchesRegistration: + """OpenAPI tools are registered under ``add_server_prefix_to_name(base, get_server_prefix(server))``, + so the dispatch lookup has to build its key the same way from the bare name ``call_tool`` + hands it. Rebuilding it as ``f"{server.name}-{bare_name}"`` used a field the prefix chain + never reads and hardcoded the separator, so every call on a server whose published prefix + differs from its display name failed with "not found in registry" instead of dispatching. + """ + + @staticmethod + def _register(server: MCPServer, base_tool_name: str) -> str: + from litellm.proxy._experimental.mcp_server.utils import ( + add_server_prefix_to_name, + get_server_prefix, + ) + + return add_server_prefix_to_name(base_tool_name, get_server_prefix(server)) + + async def _call(self, server: MCPServer, registered_key: str, bare_tool_name: str) -> CallToolResult: + from litellm.proxy._experimental.mcp_server.tool_registry import ( + global_mcp_tool_registry, + ) + + async def handler(**kwargs: Any) -> str: + return "dispatched" + + tool = MagicMock() + tool.handler = handler + + with patch.dict(global_mcp_tool_registry.tools, {registered_key: tool}, clear=True): + return await MCPServerManager()._call_openapi_tool_handler(server, bare_tool_name, {}) + + @pytest.mark.asyncio + async def test_aliased_server_dispatches_when_name_differs_from_published_prefix(self): + server = MCPServer( + server_id="dd7f2b9e-2c4a-4f1b-9e0a-8d3c6b5a4f21", + name="petstore_prod", + alias="petstore", + server_name="petstore_prod", + url=None, + transport=MCPTransport.http, + spec_path="https://example.com/petstore.yaml", + ) + registered_key = self._register(server, "list_pets") + assert registered_key == "petstore-list_pets" + + result = await self._call(server, registered_key, "list_pets") + + assert result.isError is False + assert result.content[0].text == "dispatched" + + @pytest.mark.asyncio + async def test_alias_less_server_dispatches_when_the_prefix_contains_the_separator(self): + server = MCPServer( + server_id=ALIAS_LESS_SERVER_ID, + name=ALIAS_LESS_SERVER_ID, + url=None, + transport=MCPTransport.http, + spec_path="https://example.com/wiki.yaml", + ) + registered_key = self._register(server, "read_wiki_contents") + assert registered_key == f"{ALIAS_LESS_SERVER_ID}-read_wiki_contents" + + result = await self._call(server, registered_key, "read_wiki_contents") + + assert result.isError is False + assert result.content[0].text == "dispatched" + + @pytest.mark.asyncio + async def test_dispatch_keeps_a_native_name_that_opens_with_the_prefix(self): + # Registration prefixes the upstream name whatever it looks like, so + # "petstore-list_pets" is registered as "petstore-petstore-list_pets". + # Stripping the bare name again before rebuilding the key cut that + # leading segment back off and the lookup missed. + server = MCPServer( + server_id="dd7f2b9e-2c4a-4f1b-9e0a-8d3c6b5a4f21", + name="petstore_prod", + alias="petstore", + server_name="petstore_prod", + url=None, + transport=MCPTransport.http, + spec_path="https://example.com/petstore.yaml", + ) + registered_key = self._register(server, "petstore-list_pets") + assert registered_key == "petstore-petstore-list_pets" + + result = await self._call(server, registered_key, "petstore-list_pets") + + assert result.isError is False + assert result.content[0].text == "dispatched" + + @pytest.mark.asyncio + async def test_dispatch_follows_a_non_default_prefix_separator(self): + from litellm.proxy._experimental.mcp_server import utils as mcp_utils + + server = MCPServer( + server_id="dd7f2b9e-2c4a-4f1b-9e0a-8d3c6b5a4f21", + name="petstore_prod", + alias="petstore", + server_name="petstore_prod", + url=None, + transport=MCPTransport.http, + spec_path="https://example.com/petstore.yaml", + ) + + with patch.object(mcp_utils, "MCP_TOOL_PREFIX_SEPARATOR", "__"): + registered_key = self._register(server, "list_pets") + assert registered_key == "petstore__list_pets" + + result = await self._call(server, registered_key, "list_pets") + + assert result.isError is False + assert result.content[0].text == "dispatched" + + @pytest.mark.asyncio + async def test_unregistered_tool_is_still_reported_missing(self): + server = MCPServer( + server_id="dd7f2b9e-2c4a-4f1b-9e0a-8d3c6b5a4f21", + name="petstore_prod", + alias="petstore", + server_name="petstore_prod", + url=None, + transport=MCPTransport.http, + spec_path="https://example.com/petstore.yaml", + ) + + result = await self._call(server, "petstore-list_pets", "delete_pet") + + assert result.isError is True + assert "not found in registry" in result.content[0].text diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index 1e4349c3143..c4b3c7f5f67 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -32,6 +32,8 @@ async def test_openapi_local_tool_runs_pre_call_tool_check(): fake_server.mcp_info = None fake_server.server_id = "srv-1" fake_server.server_name = "openapi-petstore" + fake_server.alias = None + fake_server.short_prefix = None fake_tool = MagicMock() fake_tool.name = "list_pets" @@ -111,6 +113,8 @@ async def test_openapi_local_tool_blocked_when_pre_call_check_raises(): fake_server.mcp_info = None fake_server.server_id = "srv-1" fake_server.server_name = "openapi-petstore" + fake_server.alias = None + fake_server.short_prefix = None fake_tool = MagicMock() fake_tool.name = "delete_pet" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py index 662ef585c6b..6e3ac014840 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -20,7 +20,9 @@ from litellm.proxy._experimental.mcp_server.utils import ( compute_short_server_prefix, get_server_prefix, is_short_mcp_tool_prefix_enabled, + is_tool_name_prefixed, iter_known_server_prefixes, + match_known_server_prefix, strip_known_server_prefix, ) from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -195,6 +197,70 @@ class TestStripKnownServerPrefix: assert strip_known_server_prefix("svc-tool", None) == "tool" +# --------------------------------------------------------------------------- +# match_known_server_prefix — the shared boundary primitive +# --------------------------------------------------------------------------- + + +class TestMatchKnownServerPrefix: + """Locates the boundary by matching registered prefixes instead of cutting + at the first separator, preferring the longest candidate so a prefix that + itself contains the separator still wins.""" + + def test_returns_matched_prefix_and_bare_name(self): + assert match_known_server_prefix("deepwiki-contents", ["deepwiki"]) == ( + "deepwiki", + "contents", + ) + + def test_returns_none_when_no_candidate_matches(self): + assert match_known_server_prefix("contents", ["deepwiki"]) is None + + def test_separator_must_follow_the_prefix(self): + assert match_known_server_prefix("deepwikicontents", ["deepwiki"]) is None + + def test_uuid_prefix_survives_its_own_separators(self): + server_id = "117c814c-1a2b-3c4d-9e8f" + assert match_known_server_prefix(f"{server_id}-contents", [server_id]) == ( + server_id, + "contents", + ) + + def test_longest_candidate_wins_over_leading_segment(self): + # "svc" is a registered prefix in its own right and also the leading + # segment of "svc-prod". A first-separator split hands the tool to + # "svc" with a bare name of "prod-run", attributing it to the wrong + # server; longest-match keeps it on "svc-prod". + assert match_known_server_prefix("svc-prod-run", ["svc", "svc-prod"]) == ( + "svc-prod", + "run", + ) + + def test_candidates_are_normalised_before_matching(self): + assert match_known_server_prefix("my_server-run", ["my server"]) == ( + "my_server", + "run", + ) + + def test_empty_candidate_never_matches_a_leading_separator(self): + assert match_known_server_prefix("-run", [""]) is None + + +class TestIsToolNamePrefixedBoundary: + """The known-prefix gate decides which branch the call path takes, so it has + to agree with the prefix the list path actually emitted.""" + + def test_uuid_prefix_is_recognised(self): + server_id = "117c814c-1a2b-3c4d-9e8f" + assert is_tool_name_prefixed(f"{server_id}-contents", known_server_prefixes={server_id}) + + def test_unrelated_hyphenated_tool_is_still_not_prefixed(self): + # Negative control: an upstream tool whose own name contains the + # separator must not start reading as prefixed just because the gate + # got more permissive about where the boundary can fall. + assert not is_tool_name_prefixed("text-to-speech", known_server_prefixes={"deepwiki"}) + + # --------------------------------------------------------------------------- # Manager-level behaviour: list + reverse-lookup # --------------------------------------------------------------------------- From 33a92bd48f4df5a21bce0f503f0fe2cc7bff2b0a Mon Sep 17 00:00:00 2001 From: tin Date: Mon, 27 Jul 2026 19:07:33 +0000 Subject: [PATCH 19/26] fix(mcp): keep REST tool listing in step with key/team grant enforcement The REST listing filter matched key/team grants through _tool_name_matches, which after the prefix-boundary change answers for every spelling routing accepts. Key-level entries in mcp_tool_permissions and toolset rows name a tool on one server and dispatch compares them bare, so a wire-form entry advertised a tool that tools/call then refused. REST listing now goes through filter_tools_by_key_team_permissions, the same function the MCP list path uses. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/rest_endpoints.py | 18 +++--- tests/mcp_tests/test_mcp_server.py | 59 +++++++++++++++++++ 2 files changed, 67 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 1736e34d70c..b5d2471336f 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -99,9 +99,9 @@ if MCP_AVAILABLE: MCPServer, _apply_toolset_scope, _fire_mcp_tool_call_logging, - _tool_name_matches, execute_mcp_tool, filter_tools_by_allowed_tools, + filter_tools_by_key_team_permissions, ) ######################################################## @@ -530,19 +530,17 @@ if MCP_AVAILABLE: tools = filter_tools_by_allowed_tools(tools, server) # Filter by the key's effective tool permissions through the same - # primitive the MCP protocol path uses (direct grants, toolset grants, - # and team/agent/org ceilings), so REST listing cannot drift from it + # function the MCP protocol path uses (direct grants, toolset grants, + # and team/agent/org ceilings), so REST listing cannot drift from it. + # Entries here are tool names on one server, written bare by every + # writer, and dispatch compares them bare; matching a wider set of + # spellings would advertise a tool that tools/call then refuses if user_api_key_auth: - from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - MCPRequestHandler, - ) - - allowed_tools_for_server = await MCPRequestHandler.get_allowed_tools_for_server( + tools = await filter_tools_by_key_team_permissions( + tools=tools, server_id=server.server_id, user_api_key_auth=user_api_key_auth, ) - if allowed_tools_for_server is not None: - tools = [tool for tool in tools if _tool_name_matches(tool.name, allowed_tools_for_server, server)] return _create_tool_response_objects(tools, server) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 390f6917573..434a9bc3809 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -2039,6 +2039,65 @@ async def test_get_tools_for_single_server_applies_disallowed_tools_without_allo assert [tool.name for tool in result] == ["read_email"] +@pytest.mark.asyncio +async def test_rest_listing_hides_key_grants_dispatch_would_refuse(): + """REST listing must answer for exactly the key/team grants dispatch honors. + + ``mcp_tool_permissions`` and toolset rows name a tool on one server, so both + the MCP list path and ``tools/call`` compare them bare. A wire-form entry + therefore grants nothing, and REST listing that matched the prefixed + spelling would advertise a tool the very next call refuses. + """ + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + from litellm.proxy._experimental.mcp_server.rest_endpoints import ( + _get_tools_for_single_server, + ) + from litellm.proxy._types import UserAPIKeyAuth + from mcp.types import Tool as MCPTool + + server_id = "3c6f6617-d23c-4f48-bfb0-f205e3b27bab" + mock_server = MagicMock() + mock_server.mcp_info = {"server_name": server_id} + mock_server.name = server_id + mock_server.server_id = server_id + mock_server.server_name = None + mock_server.alias = None + mock_server.short_prefix = None + mock_server.allowed_tools = None + mock_server.disallowed_tools = None + mock_server.tool_name_to_display_name = None + + mock_tools = [ + MCPTool( + name="read_wiki_contents", + description="Read a wiki", + inputSchema={"type": "object"}, + ), + ] + + with patch( + "litellm.proxy._experimental.mcp_server.rest_endpoints.global_mcp_server_manager" + ) as mock_manager, patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager" + ) as mock_server_manager, patch.object( + MCPRequestHandler, + "get_allowed_tools_for_server", + AsyncMock(return_value=[f"{server_id}-read_wiki_contents"]), + ): + mock_manager._get_tools_from_server = AsyncMock(return_value=mock_tools) + mock_server_manager.get_mcp_server_by_id.return_value = mock_server + + result = await _get_tools_for_single_server( + mock_server, + "Bearer test_token", + user_api_key_auth=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert result == [] + + @pytest.mark.asyncio async def test_list_tool_rest_api_with_server_specific_auth(): """Test list_tool_rest_api with server-specific auth headers.""" From d200e4a8ea84c4f588d21e105ce4ccf663ad1109 Mon Sep 17 00:00:00 2001 From: Tin Date: Mon, 27 Jul 2026 13:32:37 -0700 Subject: [PATCH 20/26] refactor(mcp): answer every tool-name permission question through one matcher The allow list, the deny list, allowed_params and the discovery filter all ask the same question, "which configured entry names this tool on this server", and each answered it in its own idiom: any() over a spelling tuple, all() over the same tuple negated, a next() that pulled a value out of a dict, and a lowercased set membership. Two review findings on this PR were symptoms of that duplication. Deriving the operands differently at one site produced the over-strip; needing a value rather than a boolean at another produced a truthiness test that read an explicitly empty allowed_params list as "nothing configured" and allowed every parameter. match_known_tool_name returns the matching entry or None, and all four sites read it, so no site can test a container's values to decide membership and the empty-list fail-open is no longer representable. Matching is case-insensitive everywhere, which closes the last divergence between discovery and dispatch: a case-variant disallowed_tools entry used to hide a tool from tools/list while tools/call still executed it. Executable lines over the merge-base drop from +9 to +4, all of it the new owner; mcp_server_manager.py loses 12 lines and the discovery filter loses 17. --- .../mcp_server/mcp_server_manager.py | 36 +++++++------------ .../proxy/_experimental/mcp_server/server.py | 27 ++++---------- .../proxy/_experimental/mcp_server/utils.py | 33 ++++++++++------- .../mcp_server/test_mcp_server.py | 11 +++--- .../mcp_server/test_mcp_server_manager.py | 15 ++++++++ 5 files changed, 60 insertions(+), 62 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 285e7b26104..0462e76c2ad 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -118,6 +118,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( is_short_mcp_tool_prefix_enabled, iter_known_server_prefixes, iter_known_tool_name_spellings, + match_known_tool_name, match_known_server_prefix, merge_mcp_headers, normalize_server_name, @@ -4262,26 +4263,19 @@ class MCPServerManager: """ Check if the tool is allowed or banned for the given server. - ``tool_name`` is bare: every caller resolves the boundary against the - server's registered prefixes before dispatch (``server.py``'s - ``original_tool_name``, the Responses handler's ``sanitized_tool_name``). - Stored entries are matched by deriving the spellings routing accepts - rather than by stripping the entries, which would cut a second boundary - out of a native name that itself opens with the server prefix. + ``tool_name`` is bare: every caller resolves the boundary against the server's + registered prefixes before dispatch (``server.py``'s ``original_tool_name``, the + Responses handler's ``sanitized_tool_name``). Configured entries are matched by + deriving the spellings routing accepts, never by stripping the entry, which would + cut a second boundary out of a native name that opens with the server prefix. """ from litellm.proxy._experimental.mcp_server.utils import ( server_applies_tool_allowlist, ) - spellings = tuple(iter_known_tool_name_spellings(tool_name, server)) - if server_applies_tool_allowlist(server): - if not server.allowed_tools: - return False - return any(spelling in server.allowed_tools for spelling in spellings) - if server.disallowed_tools: - return all(spelling not in server.disallowed_tools for spelling in spellings) - return True + return match_known_tool_name(tool_name, server, server.allowed_tools or ()) is not None + return match_known_tool_name(tool_name, server, server.disallowed_tools or ()) is None def validate_allowed_params(self, tool_name: str, arguments: dict[str, Any], server: MCPServer) -> None: """ @@ -4299,18 +4293,12 @@ class MCPServerManager: Raises: HTTPException: If allowed_params is configured for this tool but arguments contain disallowed params """ - # If no allowed_params configured, return all arguments - if not server.allowed_params: + allowed_params = server.allowed_params or {} + matched = match_known_tool_name(tool_name, server, allowed_params) + if matched is None: return - spellings = iter_known_tool_name_spellings(tool_name, server) - allowed_params_list = next( - (server.allowed_params[name] for name in spellings if name in server.allowed_params), None - ) - - # If this tool doesn't have allowed_params specified, allow all params - if allowed_params_list is None: - return None + allowed_params_list = allowed_params[matched] # Filter arguments to only include allowed parameters disallowed_params = [param for param in arguments.keys() if param not in allowed_params_list] diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5181f3a2676..2c774224098 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -65,7 +65,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( extract_mcp_tool_result_error_message, get_server_prefix, iter_known_server_prefixes, - iter_known_tool_name_spellings, + match_known_tool_name, ) from litellm.proxy._types import ( ProxyException, @@ -1421,28 +1421,13 @@ if MCP_AVAILABLE: """ Check if a tool name matches any name in the filter list. - Matches via the same ``iter_known_tool_name_spellings`` the server-level - permission checks use, so discovery hides exactly what dispatch refuses; - covering fewer spellings here leaves a blocked tool advertised in - ``tools/list``. Comparison is case-insensitive to handle OpenAPI - operationIds that may be in camelCase. - - Args: - tool_name: The tool name to check (may be prefixed like "server-tool_name") - filter_list: List of tool names to match against - mcp_server: The server the tool belongs to, whose registered prefixes - locate the boundary exactly. Required: guessing the boundary at - the first separator silently mismatches every tool on a server - whose prefix contains the separator. - - Returns: - True if any spelling of the tool name is in the filter list + Reads the same owner the server-level permission checks use, so discovery hides + exactly what dispatch refuses. ``mcp_server`` is required: guessing the boundary + at the first separator mismatches every tool on a server whose prefix contains + the separator. """ - filter_list_lower = {f.lower() for f in filter_list} bare_name = strip_known_server_prefix(tool_name, mcp_server) - spellings = (tool_name, *iter_known_tool_name_spellings(bare_name, mcp_server)) - - return any(spelling.lower() in filter_list_lower for spelling in spellings) + return match_known_tool_name(bare_name, mcp_server, filter_list) is not None def filter_tools_by_allowed_tools( tools: list[MCPTool], diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index e8f8c08188e..900d6259a5a 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -23,6 +23,8 @@ import importlib import os from urllib.parse import quote +from litellm.types.mcp_server.mcp_server_manager import MCPServer + # Constants # # NOTE: The environment-backed values below are read once, when this module is @@ -326,24 +328,31 @@ def iter_known_server_prefixes(server: Any) -> Iterator[str]: yield from _emit(server_id) -def iter_known_tool_name_spellings(tool_name: str, server: Any) -> Iterator[str]: - """Yield every name that denotes the bare ``tool_name`` on ``server``. - - The bare name, then its wire spelling under each prefix from - ``iter_known_server_prefixes``. Routing resolves an inbound name against that - whole set, so anything keyed by tool name (the routing map, the allow/deny - lists, ``allowed_params``) must cover it too or it answers for fewer names - than are reachable, which fails open on ``disallowed_tools``. - ``get_server_prefix`` alone covers only the published spelling, and that moves - with the alias and with ``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``. These are - spellings of one tool on one server, so honoring all of them normalizes the - entry rather than widening a grant. +def iter_known_tool_name_spellings(tool_name: str, server: MCPServer) -> Iterator[str]: + """Yield every name that denotes the bare ``tool_name`` on ``server``: the bare name, + then its wire spelling under each prefix ``iter_known_server_prefixes`` accepts. + ``get_server_prefix`` covers only the currently published one, and that moves with the + alias and with ``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``. """ yield tool_name for prefix in iter_known_server_prefixes(server): yield add_server_prefix_to_name(tool_name, prefix) +def match_known_tool_name(tool_name: str, server: MCPServer, names: Iterable[str]) -> str | None: + """Return the entry of ``names`` that denotes ``tool_name`` on ``server``, else ``None``. + + The single question every tool-name-keyed site asks: the allow list, the deny list, + ``allowed_params`` and the discovery filter. Matching spans every spelling routing + accepts and ignores case, so discovery hides exactly what dispatch refuses. Callers + read the returned entry rather than testing a container's values, which is what stops + an explicitly empty ``allowed_params`` list from reading as "nothing configured". + """ + entries = {name.casefold(): name for name in names} + spellings = map(str.casefold, iter_known_tool_name_spellings(tool_name, server)) + return next((entries[spelling] for spelling in spellings if spelling in entries), None) + + def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]: """Return the unprefixed name plus the server name used as prefix. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index a3577c85621..5115fde2687 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -8146,10 +8146,11 @@ class TestListFiltersHonorThePrefixBoundary: published = MCPTool(name="eiG-read_wiki_contents", description="", inputSchema={"type": "object"}) for spelling in registered: - server = _server(disallowed_tools=[spelling]) + for entry in (spelling, spelling.upper()): + server = _server(disallowed_tools=[entry]) - refused = not manager.check_allowed_or_banned_tools("read_wiki_contents", server) - hidden = filter_tools_by_allowed_tools([published], server) == [] + refused = not manager.check_allowed_or_banned_tools("read_wiki_contents", server) + hidden = filter_tools_by_allowed_tools([published], server) == [] - assert refused, spelling - assert hidden, spelling + assert refused, entry + assert hidden, entry 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 2977f13caf1..134a5d89225 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 @@ -9511,6 +9511,21 @@ class TestServerToolListsHonorThePrefixBoundary: assert exc_info.value.status_code == 403 assert "include_internal" in exc_info.value.detail["error"] + @pytest.mark.asyncio + async def test_a_case_variant_blocklist_entry_still_blocks(self): + server = self._aliased_server(disallowed_tools=["PetStore-DeletePet"]) + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "deletepet") + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_a_case_variant_allowlist_entry_grants_the_tool(self): + server = self._aliased_server(allowed_tools=["PetStore-GetPetById"]) + + await self._run_check(server, "getpetbyid") + @pytest.mark.asyncio async def test_an_explicitly_empty_allowed_params_list_refuses_every_parameter(self): server = self._alias_less_server(allowed_params={"read_wiki_contents": []}) From b2d4dde46468dfa07b819aad61116e62b27d469c Mon Sep 17 00:00:00 2001 From: Tin Date: Mon, 27 Jul 2026 16:29:40 -0700 Subject: [PATCH 21/26] fix(mcp): give the key/team grant question one predicate Bugbot flagged REST listing advertising key/team grants that tools/call then refuses. The listing side was fixed by routing through filter_tools_by_key_team_permissions, but the two paths still answered the question with separate implementations that only happened to agree: listing stripped the known prefix and compared bare, dispatch compared whatever name it was handed, and each carried its own reading of None and of an empty list. Changing either side silently diverges from the other, which is how this defect appeared in the first place. MCPRequestHandler.tool_is_granted owns the whole decision, and both is_tool_allowed_for_server and filter_tools_by_key_team_permissions read it. None still means no tool-level restriction and an empty list still grants nothing, now stated once. Grants are stored bare by every writer, so matching stays exact against the bare name, deliberately unlike the server-level lists, which honor every spelling routing accepts. test_key_team_listing_and_dispatch_agree drives both production paths over one matrix. It asserts the expected verdict as well as the agreement, because two paths reading one predicate makes equality alone tautological: a wrong predicate keeps them consistent and the agreement assertion alone survived two mutants that the verdict assertion kills. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 26 +++++---- .../proxy/_experimental/mcp_server/server.py | 8 ++- .../mcp_server/test_mcp_server.py | 58 +++++++++++++++++++ 3 files changed, 77 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 5d8ac8d678f..dc796837277 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1963,6 +1963,18 @@ class MCPRequestHandler: return allowed_tools + @staticmethod + def tool_is_granted(bare_tool_name: str, allowed_tool_names: list[str] | None) -> bool: + """Whether key/team tool permissions reach ``bare_tool_name`` on one server. + + ``None`` means no tool-level restriction; an empty list grants nothing. Entries + name a tool on a single server and every writer stores them bare, so the + comparison is exact against the bare name rather than against the spellings + routing accepts. Both the listing path and the call path answer through here, so + discovery cannot advertise a tool that ``tools/call`` then refuses. + """ + return allowed_tool_names is None or bare_tool_name in allowed_tool_names + @staticmethod async def is_tool_allowed_for_server( tool_name: str, @@ -1973,7 +1985,7 @@ class MCPRequestHandler: Check if a specific tool is allowed for a server based on key/team permissions. Args: - tool_name: Name of the tool to check + tool_name: Bare tool name, already resolved against the server's prefixes server_id: Server ID user_api_key_auth: User auth @@ -1984,17 +1996,7 @@ class MCPRequestHandler: server_id=server_id, user_api_key_auth=user_api_key_auth, ) - - # None means no restrictions (allow all) - if allowed_tools is None: - return True - - # Empty list means no tools allowed - if not allowed_tools: - return False - - # Check if tool is in allowed list - return tool_name in allowed_tools + return MCPRequestHandler.tool_is_granted(tool_name, allowed_tools) @staticmethod def is_tool_allowed( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 2c774224098..48effdb0f6e 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2305,14 +2305,16 @@ if MCP_AVAILABLE: server_id=server_id, user_api_key_auth=user_api_key_auth, ) - if allowed_tool_names is None: - return tools # Tools arrive prefixed with the server's own prefix; strip exactly that # prefix (resolved from the server) rather than the first separator, so a # prefix containing the separator still reduces to the stored bare name. server = global_mcp_server_manager.get_mcp_server_by_id(server_id) - return [t for t in tools if strip_known_server_prefix(t.name, server) in allowed_tool_names] + return [ + t + for t in tools + if MCPRequestHandler.tool_is_granted(strip_known_server_prefix(t.name, server), allowed_tool_names) + ] async def _list_mcp_tools( user_api_key_auth: UserAPIKeyAuth | None = None, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 5115fde2687..b2cdcd51294 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -8154,3 +8154,61 @@ class TestListFiltersHonorThePrefixBoundary: assert refused, entry assert hidden, entry + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "grants,expected", + [ + (None, True), + ([], False), + (["read_wiki_contents"], True), + (["read_wiki_structure"], False), + ([f"{SERVER_ID}-read_wiki_contents"], False), + (["READ_WIKI_CONTENTS"], False), + ], + ) + async def test_key_team_listing_and_dispatch_agree(self, grants, expected): + """The key/team grant question, driven through both production paths. + + Listing and dispatch read one predicate, so a row where the tool is advertised + and then refused (or hidden while callable) cannot exist. Asserting the expected + verdict as well as the agreement matters: both paths reading one predicate makes + equality alone tautological, so a wrong predicate would keep them consistent. + Grants are stored bare, so the wire-form and case-variant rows deny; that is + deliberately unlike the server-level lists, which honor every spelling. + """ + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + from litellm.proxy._experimental.mcp_server.server import ( + filter_tools_by_key_team_permissions, + ) + from litellm.proxy._types import UserAPIKeyAuth + + server = MCPServer( + server_id=self.SERVER_ID, + name=self.SERVER_ID, + url="http://127.0.0.1:5115/mcp", + transport=MCPTransport.http, + ) + published = MCPTool( + name=f"{self.SERVER_ID}-read_wiki_contents", description="", inputSchema={"type": "object"} + ) + auth = UserAPIKeyAuth(api_key="sk-test") + + with patch.object( + MCPRequestHandler, "get_allowed_tools_for_server", AsyncMock(return_value=grants) + ), patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager" + ) as mock_manager: + mock_manager.get_mcp_server_by_id.return_value = server + + listed = await filter_tools_by_key_team_permissions([published], self.SERVER_ID, auth) != [] + callable_ = await MCPRequestHandler.is_tool_allowed_for_server( + tool_name="read_wiki_contents", server_id=self.SERVER_ID, user_api_key_auth=auth + ) + + assert listed == callable_, f"grants={grants!r} listed={listed} callable={callable_}" + assert listed is expected, f"grants={grants!r} expected={expected} got={listed}" From 7962407be02ada648db0e1c6682d179fb00b3312 Mon Sep 17 00:00:00 2001 From: Tin Date: Mon, 27 Jul 2026 17:03:17 -0700 Subject: [PATCH 22/26] fix(mcp): keep tool identity exact, fold case only where registration does Greptile flagged that match_known_tool_name case-folded both the configured entry and the derived spellings. Routing keeps two tools whose names differ only in case as two tools, so folding merged identities the dispatcher separates: on a server exposing getPet and getpet, an allowlist naming getPet also granted getpet, and a blocklist naming getPet also denied getpet. That is unauthorized execution on one arm and the wrong tool denied on the other. Matching is now exact, which is what identity means here. The case leniency it replaces was never typo tolerance; _register_openapi_tools rewrites every operationId through sanitize_openapi_tool_name, so an allowed_tools entry holding the spec's own spelling never equals the registered name. That link is recovered by replaying the same rewrite, and only on servers that carry a spec_path, which is how the rest of the manager already recognizes an OpenAPI server. Every name that rewrite produces is lowercased, so no two tools on such a server can differ only in case and the fold cannot merge anything. Native servers get no folding at all. test_case_folding_applies_to_openapi_ servers_and_not_to_native_ones pins both halves, and two tests pin that a policy naming one tool leaves its case-variant sibling alone. Dropping the spec_path guard, dropping the fold, and forcing the fold path are all killed. The two pre-existing case-insensitivity tests describe OpenAPI servers in their own docstrings but built fixtures without a spec_path, a shape production never produces for one; they now set it. --- .../proxy/_experimental/mcp_server/utils.py | 30 ++++++++++++++----- .../mcp_server/test_mcp_server.py | 25 ++++++++++++---- .../mcp_server/test_mcp_server_manager.py | 18 +++++++---- 3 files changed, 54 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 900d6259a5a..4c0690082d5 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -343,14 +343,30 @@ def match_known_tool_name(tool_name: str, server: MCPServer, names: Iterable[str """Return the entry of ``names`` that denotes ``tool_name`` on ``server``, else ``None``. The single question every tool-name-keyed site asks: the allow list, the deny list, - ``allowed_params`` and the discovery filter. Matching spans every spelling routing - accepts and ignores case, so discovery hides exactly what dispatch refuses. Callers - read the returned entry rather than testing a container's values, which is what stops - an explicitly empty ``allowed_params`` list from reading as "nothing configured". + ``allowed_params`` and the discovery filter, so discovery hides exactly what dispatch + refuses. It spans every spelling routing accepts and no more. A tool's identity is its + exact name, because routing dispatches two names differing only in case as two tools, + and folding case here would let one policy decide both. + + OpenAPI servers are the exception, and not a fuzzy one. ``_register_openapi_tools`` + rewrites every operationId through ``sanitize_openapi_tool_name``, so configuration + holding the spec's own spelling never equals the registered name; replaying that exact + rewrite on the entry recovers the link. It cannot merge identities, because every name + it produces is lowercased, so no two tools on such a server differ only in case. + + Callers read the returned entry rather than testing a container's values, which is what + stops an explicitly empty ``allowed_params`` list from reading as "nothing configured". """ - entries = {name.casefold(): name for name in names} - spellings = map(str.casefold, iter_known_tool_name_spellings(tool_name, server)) - return next((entries[spelling] for spelling in spellings if spelling in entries), None) + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + sanitize_openapi_tool_name, + ) + + spellings = set(iter_known_tool_name_spellings(tool_name, server)) + exact = next((name for name in names if name in spellings), None) + if exact is not None or not getattr(server, "spec_path", None): + return exact + sanitized = {sanitize_openapi_tool_name(spelling) for spelling in spellings} + return next((name for name in names if sanitize_openapi_tool_name(name) in sanitized), None) def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index b2cdcd51294..3a84428add1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -4519,6 +4519,7 @@ def test_tool_name_matches_case_insensitive(): server_name="per_store", url="http://127.0.0.1:5115/mcp", transport=MCPTransport.http, + spec_path="/specs/petstore.yaml", ) # Test case 1: Unprefixed tool name with camelCase in filter list @@ -4597,6 +4598,7 @@ def test_filter_tools_by_allowed_tools_case_insensitive(): name="per_store", server_name="per_store", transport=MCPTransport.http, + spec_path="/specs/petstore.yaml", allowed_tools=["addPet", "updatePet", "findPetsByStatus"], ) @@ -8080,12 +8082,19 @@ class TestListFiltersHonorThePrefixBoundary: assert not _tool_name_matches(f"{self.SERVER_ID}-read_wiki_contents", ["read_wiki_structure"], server) - def test_match_is_still_case_insensitive(self): + def test_case_folding_applies_to_openapi_servers_and_not_to_native_ones(self): + # Registration rewrites operationIds through sanitize_openapi_tool_name, so + # folding recovers a spec-spelled entry on an OpenAPI server. A native server + # gets none of it: routing dispatches two names differing only in case as two + # tools, so one policy must not decide both. from litellm.proxy._experimental.mcp_server.server import _tool_name_matches - server = self._alias_less_server() + native = self._alias_less_server() + openapi = self._alias_less_server(spec_path="/specs/petstore.yaml") - assert _tool_name_matches(f"{self.SERVER_ID}-findPetsByStatus", ["findpetsbystatus"], server) + assert _tool_name_matches(f"{self.SERVER_ID}-findPetsByStatus", ["findpetsbystatus"], openapi) + assert not _tool_name_matches(f"{self.SERVER_ID}-findPetsByStatus", ["findpetsbystatus"], native) + assert _tool_name_matches(f"{self.SERVER_ID}-findpetsbystatus", ["findpetsbystatus"], native) def test_alias_form_entry_matches_a_tool_published_under_the_short_prefix(self, monkeypatch): # Routing accepts the alias form, so an entry stored before short @@ -8111,6 +8120,10 @@ class TestListFiltersHonorThePrefixBoundary: A spelling the blocklist enforces but the filter misses leaves a blocked tool advertised; the reverse hides a tool that would have been callable. + Every spelling routing registers bans, and its upper-cased form bans nothing, + because a tool's identity is its exact name; asserting the verdict and not only + the agreement is what keeps this from passing on a matcher that answers wrongly + but consistently. """ from mcp.types import Tool as MCPTool @@ -8146,14 +8159,14 @@ class TestListFiltersHonorThePrefixBoundary: published = MCPTool(name="eiG-read_wiki_contents", description="", inputSchema={"type": "object"}) for spelling in registered: - for entry in (spelling, spelling.upper()): + for entry, expected in ((spelling, True), (spelling.upper(), False)): server = _server(disallowed_tools=[entry]) refused = not manager.check_allowed_or_banned_tools("read_wiki_contents", server) hidden = filter_tools_by_allowed_tools([published], server) == [] - assert refused, entry - assert hidden, entry + assert refused == hidden, entry + assert refused is expected, entry @pytest.mark.asyncio @pytest.mark.parametrize( 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 134a5d89225..cdef50ad6b5 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 @@ -9512,19 +9512,25 @@ class TestServerToolListsHonorThePrefixBoundary: assert "include_internal" in exc_info.value.detail["error"] @pytest.mark.asyncio - async def test_a_case_variant_blocklist_entry_still_blocks(self): - server = self._aliased_server(disallowed_tools=["PetStore-DeletePet"]) + async def test_a_blocklist_entry_does_not_reach_a_case_variant_sibling_tool(self): + server = self._aliased_server(disallowed_tools=["petstore-getPet"]) with pytest.raises(HTTPException) as exc_info: - await self._run_check(server, "deletepet") + await self._run_check(server, "getPet") assert exc_info.value.status_code == 403 + await self._run_check(server, "getpet") @pytest.mark.asyncio - async def test_a_case_variant_allowlist_entry_grants_the_tool(self): - server = self._aliased_server(allowed_tools=["PetStore-GetPetById"]) + async def test_an_allowlist_entry_does_not_grant_a_case_variant_sibling_tool(self): + server = self._aliased_server(allowed_tools=["petstore-getPet"]) - await self._run_check(server, "getpetbyid") + await self._run_check(server, "getPet") + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "getpet") + + assert exc_info.value.status_code == 403 @pytest.mark.asyncio async def test_an_explicitly_empty_allowed_params_list_refuses_every_parameter(self): From 8018bc3996c533ac642b5fc0e0511cef891315b2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 30 Jul 2026 22:19:30 -0700 Subject: [PATCH 23/26] fix(e2e): exclude skipped tests from coverage-registry numerator The collector read @pytest.mark.covers off every collected item, and collection does not evaluate skips, so a test carrying both a skip and a covers marker reported its cell as covered while asserting nothing. 17 files under tests/e2e do exactly that, which inflated the headline from 290/434 to 311/434. A cell now counts as covered only when at least one test pytest would actually run declares it; a cell claimed by both a live and a skipped test stays covered. Skip state comes from pytest's own evaluator, so skip, skipif (bool and string conditions), and module-level pytestmark resolve exactly as they do in the e2e run. Cells left uncovered this way are listed under the headline and exported as skipped_markers (JSON) and litellm_e2e_coverage_skipped_markers (Prometheus) so the gap surfaces instead of disappearing; the Loki line contract is unchanged. A marker on a skipped test that points outside the registry is still an orphan, so --strict keeps its reach. Because skipif resolves against the environment the collector runs in, the number now depends on that environment; run it where the e2e suite runs. A pytest.skip() call inside a test body remains invisible to a static pass, which the module docstring and README both state. --- tests/e2e/CLAUDE.md | 2 + tests/e2e/coverage_registry/README.md | 10 ++ tests/e2e/coverage_registry/collector.py | 94 +++++++++++++++--- tests/e2e/coverage_registry/test_collector.py | 97 +++++++++++++++++++ 4 files changed, 189 insertions(+), 14 deletions(-) diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 17aee22560c..180639b53e3 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -85,6 +85,8 @@ The metric is coverage: the share of registry rows that have a passing covering Tests do not declare a dashboard module directly. They only declare the registry cell id with `@pytest.mark.covers("...")`; the registry row decides the module, tier, endpoint, and dashboard rollup. Run `python -m coverage_registry.collector --strict` when you want CI to reject unknown marker ids. Add `--fail-on-collection-errors` when the job should also fail on pytest collection errors. +Skipping a test gives its cell back to the gap list: the collector counts a cell as covered only when a test pytest would actually run declares it, and prints the cells left claimed only by skipped tests. So a `@pytest.mark.skip` on a red cell is honest bookkeeping, not a way to keep the number up. + ### Naming grammar per module LLMs - endpoint features (subject = the route), seeded from the Claude Code compat matrix. `chat_completions`, `messages`, and `responses` roll up to `Core LLMs`. Other LLM endpoints, including `batches` and `realtime`, roll up to `Non-Core LLMs`. diff --git a/tests/e2e/coverage_registry/README.md b/tests/e2e/coverage_registry/README.md index aef4c16c89a..5627c88dee4 100644 --- a/tests/e2e/coverage_registry/README.md +++ b/tests/e2e/coverage_registry/README.md @@ -37,6 +37,16 @@ def test_openai_streaming_tool_calls(self) -> None: It is static: a collect-only pass reads the markers, so it runs no test and needs no live proxy. Whether a covered cell currently passes or fails is a separate, live concern. +A skipped test asserts nothing, so its markers do not count. A cell is covered only when +at least one test pytest would actually run declares it; a cell claimed by both a live +test and a skipped one stays covered. Skip state comes from pytest's own evaluator, so +`skip` and `skipif` resolve exactly as they do in the e2e run, which also means a +`skipif` on an absent credential makes that cell uncovered in the environments where the +test cannot run. Cells left uncovered this way are listed under the headline (and counted +by `litellm_e2e_coverage_skipped_markers`) so an unskipped-pending gap is visible rather +than inflating the number. The one skip the collector cannot see is `pytest.skip()` +called from inside a test body, since it does not exist until the test runs. + ``` cd tests/e2e && PYTHONPATH=. python -m coverage_registry.collector ``` diff --git a/tests/e2e/coverage_registry/collector.py b/tests/e2e/coverage_registry/collector.py index 50ef23bcb1a..e20f7884f55 100644 --- a/tests/e2e/coverage_registry/collector.py +++ b/tests/e2e/coverage_registry/collector.py @@ -5,6 +5,13 @@ Coverage here is static: it reads the markers via a collect-only pass, so it run no test and needs no live proxy. Whether a covered cell currently passes or fails (covered_pass vs covered_fail) is a separate, live concern layered on top later. +A skipped test asserts nothing, so its markers do not count: a cell is covered +only when at least one test that pytest would actually run declares it. Skip +state is read with pytest's own evaluator, so `skip` and `skipif` are resolved +exactly as the e2e run resolves them in this environment. The one skip the +collector cannot see is `pytest.skip()` called from inside a test body, which +does not exist until the test runs. + cd tests/e2e && PYTHONPATH=. python -m coverage_registry.collector """ @@ -20,6 +27,7 @@ from pathlib import Path from typing import Literal import pytest +from _pytest.skipping import evaluate_skip_marks from pydantic import BaseModel from .registry import load_registry @@ -28,22 +36,57 @@ from .schema import MODULE_ORDER, Cell, Tier, dashboard_module, loki_module_labe E2E_DIR = Path(__file__).resolve().parent.parent +@dataclass(frozen=True, slots=True) +class CollectedMarkers: + """What a collect-only pass saw: cell ids declared by tests that would run, + cell ids only ever declared by skipped tests, and nodes that failed to import.""" + + covered: frozenset[str] + skipped_only: frozenset[str] + collection_errors: tuple[str, ...] + + +def _is_skipped(item: pytest.Item) -> bool: + """True when pytest would skip this test instead of running it. + + A marker pytest cannot evaluate (for example a bare boolean `skipif` with no + reason) turns into a setup failure at run time, so the test asserts nothing + either way and is treated the same as a skip. + """ + try: + return evaluate_skip_marks(item) is not None + except (pytest.fail.Exception, TypeError): + return True + + class _CoversSink: """Pytest plugin: after collection, capture every cell id declared via - @pytest.mark.covers(...), plus any nodes that failed to import.""" + @pytest.mark.covers(...) split by whether its test would run, plus any nodes + that failed to import.""" def __init__(self) -> None: self.covered_ids: frozenset[str] = frozenset() + self.skipped_only_ids: frozenset[str] = frozenset() self.collection_errors: tuple[str, ...] = () def pytest_collection_finish(self, session: pytest.Session) -> None: - marker_args: tuple[tuple[object, ...], ...] = tuple( - marker.args - for item in session.items + marker_args: tuple[tuple[bool, tuple[object, ...]], ...] = tuple( + (skipped, marker.args) + for item, skipped in ((i, _is_skipped(i)) for i in session.items) for marker in item.iter_markers(name="covers") ) + declared = tuple( + (skipped, arg) + for skipped, args in marker_args + for arg in args + if isinstance(arg, str) + ) self.covered_ids = frozenset( - arg for args in marker_args for arg in args if isinstance(arg, str) + cell_id for skipped, cell_id in declared if not skipped + ) + self.skipped_only_ids = ( + frozenset(cell_id for skipped, cell_id in declared if skipped) + - self.covered_ids ) def pytest_collectreport(self, report: pytest.CollectReport) -> None: @@ -51,10 +94,8 @@ class _CoversSink: self.collection_errors = (*self.collection_errors, report.nodeid) -def collect_covered_ids( - e2e_dir: Path = E2E_DIR, -) -> tuple[frozenset[str], tuple[str, ...]]: - """Return (covered cell ids, nodeids that failed to import).""" +def collect_markers(e2e_dir: Path = E2E_DIR) -> CollectedMarkers: + """Read every @pytest.mark.covers marker in `e2e_dir` via a collect-only pass.""" sink = _CoversSink() with contextlib.redirect_stdout(io.StringIO()): pytest.main( @@ -68,7 +109,11 @@ def collect_covered_ids( ], plugins=[sink], ) - return sink.covered_ids, sink.collection_errors + return CollectedMarkers( + covered=sink.covered_ids, + skipped_only=sink.skipped_only_ids, + collection_errors=sink.collection_errors, + ) @dataclass(frozen=True, slots=True) @@ -93,6 +138,7 @@ class CoverageReport: p0_covered: int p0_gaps: tuple[str, ...] orphan_markers: tuple[str, ...] + skipped_markers: tuple[str, ...] collection_errors: tuple[str, ...] @property @@ -122,6 +168,7 @@ def compute_coverage( cells: tuple[Cell, ...], covered: frozenset[str], collection_errors: tuple[str, ...] = (), + skipped_only: frozenset[str] = frozenset(), ) -> CoverageReport: p0_cells = tuple(c for c in cells if c.tier is Tier.P0) registry_ids = frozenset(c.id for c in cells) @@ -132,7 +179,8 @@ def compute_coverage( p0_total=len(p0_cells), p0_covered=sum(1 for c in p0_cells if c.id in covered), p0_gaps=tuple(sorted(c.id for c in p0_cells if c.id not in covered)), - orphan_markers=tuple(sorted(covered - registry_ids)), + orphan_markers=tuple(sorted((covered | skipped_only) - registry_ids)), + skipped_markers=tuple(sorted(skipped_only & registry_ids)), collection_errors=collection_errors, ) @@ -161,6 +209,15 @@ def render(report: CoverageReport) -> str: if report.orphan_markers else () ) + skipped = ( + ( + f"\n{len(report.skipped_markers)} cell(s) are claimed only by skipped tests, " + f"so they count as uncovered (unskip the test or drop the marker):\n " + + "\n ".join(report.skipped_markers), + ) + if report.skipped_markers + else () + ) warning = ( ( f"\nWARNING: {len(report.collection_errors)} node(s) failed to import during " @@ -170,7 +227,7 @@ def render(report: CoverageReport) -> str: if report.collection_errors else () ) - return "\n".join((*lines, *orphans, *warning)) + return "\n".join((*lines, *orphans, *skipped, *warning)) def _report_dict(report: CoverageReport) -> dict[str, object]: @@ -190,6 +247,7 @@ def _report_dict(report: CoverageReport) -> dict[str, object]: for m in report.modules ], "orphan_markers": list(report.orphan_markers), + "skipped_markers": list(report.skipped_markers), "collection_errors": list(report.collection_errors), } @@ -234,6 +292,9 @@ def render_prometheus(report: CoverageReport) -> str: "# HELP litellm_e2e_coverage_orphan_markers Coverage markers not found in the registry.", "# TYPE litellm_e2e_coverage_orphan_markers gauge", f"litellm_e2e_coverage_orphan_markers {len(report.orphan_markers)}", + "# HELP litellm_e2e_coverage_skipped_markers Registry cells claimed only by skipped tests.", + "# TYPE litellm_e2e_coverage_skipped_markers gauge", + f"litellm_e2e_coverage_skipped_markers {len(report.skipped_markers)}", "# HELP litellm_e2e_coverage_collection_errors Pytest nodes that failed during collection.", "# TYPE litellm_e2e_coverage_collection_errors gauge", f"litellm_e2e_coverage_collection_errors {len(report.collection_errors)}", @@ -286,8 +347,13 @@ def main() -> int: ) args = _CliArgs.model_validate(vars(parser.parse_args())) cells = load_registry() - covered, errors = collect_covered_ids() - report = compute_coverage(cells, covered, errors) + markers = collect_markers() + report = compute_coverage( + cells, + markers.covered, + markers.collection_errors, + markers.skipped_only, + ) output = { "text": render, "json": render_json, diff --git a/tests/e2e/coverage_registry/test_collector.py b/tests/e2e/coverage_registry/test_collector.py index 079ee215866..a85190cc3ba 100644 --- a/tests/e2e/coverage_registry/test_collector.py +++ b/tests/e2e/coverage_registry/test_collector.py @@ -12,6 +12,7 @@ from pathlib import Path import pytest from coverage_registry.collector import ( + collect_markers, compute_coverage, render, render_json, @@ -61,6 +62,27 @@ def test_orphan_marker_is_reported_not_counted() -> None: assert report.orphan_markers == ("llm.ghost",) +def test_cell_claimed_only_by_a_skipped_test_is_uncovered() -> None: + cells = (_llm("llm.a", Tier.P0), _llm("llm.b", Tier.P0)) + report = compute_coverage( + cells, frozenset({"llm.a"}), skipped_only=frozenset({"llm.b"}) + ) + assert (report.covered, report.p0_covered) == (1, 1) + assert report.p0_gaps == ("llm.b",) + assert report.skipped_markers == ("llm.b",) + assert "only by skipped tests" in render(report) + assert '"skipped_markers": [\n "llm.b"\n ]' in render_json(report) + assert "litellm_e2e_coverage_skipped_markers 1" in render_prometheus(report) + + +def test_skipped_marker_outside_the_registry_is_still_an_orphan() -> None: + report = compute_coverage( + (_llm("llm.a", Tier.P0),), frozenset(), skipped_only=frozenset({"llm.ghost"}) + ) + assert report.orphan_markers == ("llm.ghost",) + assert report.skipped_markers == () + + def test_logging_and_guardrail_roll_up_into_one_module() -> None: cells = ( LoggingCell( @@ -175,6 +197,81 @@ def test_loki_render_exposes_exact_stdout_lines_for_loki() -> None: ) +_MARKED_TESTS = ''' +import pytest + + +@pytest.mark.covers("llm.runs") +def test_runs() -> None: + pass + + +@pytest.mark.skip(reason="stage red: product gap") +@pytest.mark.covers("llm.skipped") +def test_skipped() -> None: + pass + + +@pytest.mark.skipif(True, reason="credentials absent in this environment") +@pytest.mark.covers("llm.skipif_true") +def test_skipif_true() -> None: + pass + + +@pytest.mark.skipif(False, reason="credentials present in this environment") +@pytest.mark.covers("llm.skipif_false") +def test_skipif_false() -> None: + pass + + +@pytest.mark.skipif("True") +@pytest.mark.covers("llm.skipif_string") +def test_skipif_string_condition() -> None: + pass + + +@pytest.mark.covers("llm.shared") +def test_shared_cell_runs() -> None: + pass + + +@pytest.mark.skip(reason="stage red: product gap") +@pytest.mark.covers("llm.shared") +def test_shared_cell_skipped() -> None: + pass +''' + +_MODULE_LEVEL_SKIP = ''' +import pytest + +pytestmark = pytest.mark.skipif(True, reason="whole module needs a session fixture") + + +@pytest.mark.covers("llm.module_skipped") +def test_module_level_skip() -> None: + pass +''' + + +def test_collection_counts_only_markers_on_tests_that_would_run( + tmp_path: Path, +) -> None: + """The collect-only pass is the numerator, so a test pytest would skip must not + contribute its cell. A cell stays covered as long as one runnable test claims it.""" + (tmp_path / "test_marked.py").write_text(_MARKED_TESTS) + (tmp_path / "test_module_skip.py").write_text(_MODULE_LEVEL_SKIP) + + markers = collect_markers(tmp_path) + + assert markers.covered == frozenset( + {"llm.runs", "llm.skipif_false", "llm.shared"} + ) + assert markers.skipped_only == frozenset( + {"llm.skipped", "llm.skipif_true", "llm.skipif_string", "llm.module_skipped"} + ) + assert markers.collection_errors == () + + def test_real_registry_loads_and_ids_are_unique() -> None: cells = load_registry() ids = [c.id for c in cells] From 179ebdb86bfefea9cde3d9e9028b7e62927af1eb Mon Sep 17 00:00:00 2001 From: Tin Date: Mon, 27 Jul 2026 17:34:20 -0700 Subject: [PATCH 24/26] fix(mcp): make the operationId to tool-name map a single owner Greptile found that the OpenAPI fallback added a commit ago collapsed operation IDs that registration keeps apart: foo/bar and foo.bar register as two tools but sanitize_openapi_tool_name rewrites both to foo_bar, so a policy naming either also decided the other. The cause was two owners for one map, and picking the wrong one. Registration names an operationId inline at _register_openapi_tools with operation_id.replace(" ", "_").lower(), which keeps / and . ; the separate sanitize_openapi_tool_name replaces every character outside [a-zA-Z0-9_-] and belongs to register_tools_from_openapi, which has no production caller. Nothing made the matcher use the one that actually registers, so it used the lookalike. That inline expression is now openapi_tool_name in utils, and both registration and the matcher call it. Replaying the registering function is the whole safety argument, and it is structural rather than a claim: two operationIds that register as two tools normalize to two names here by construction, because this is the map that registered them. A coarser lookalike cannot be substituted without a test failing. The matcher also loses its exact-then-fallback split. The transform is identity on native servers and idempotent on already-registered names, so normalizing both sides is exact matching where no OpenAPI spec is involved. Executable lines drop by three this round; the branch is +5 over the merge-base for four shared owners that removed duplication at six call sites. --- .../mcp_server/mcp_server_manager.py | 5 ++- .../proxy/_experimental/mcp_server/utils.py | 45 ++++++++++--------- .../mcp_server/test_mcp_server_manager.py | 10 +++++ 3 files changed, 36 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0462e76c2ad..89c83459524 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -118,10 +118,11 @@ from litellm.proxy._experimental.mcp_server.utils import ( is_short_mcp_tool_prefix_enabled, iter_known_server_prefixes, iter_known_tool_name_spellings, - match_known_tool_name, match_known_server_prefix, + match_known_tool_name, merge_mcp_headers, normalize_server_name, + openapi_tool_name, parse_admin_env_vars, strip_known_server_prefix, validate_mcp_server_name, @@ -1785,7 +1786,7 @@ class MCPServerManager: # Generate tool name (without prefix initially) operation_id = operation.get("operationId", f"{method}_{path.replace('/', '_')}") - base_tool_name = operation_id.replace(" ", "_").lower() + base_tool_name = openapi_tool_name(operation_id) # Add server prefix to tool name prefixed_tool_name = add_server_prefix_to_name(base_tool_name, server_prefix) diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 4c0690082d5..698df147247 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -2,7 +2,10 @@ MCP Server Utilities """ +import hashlib +import importlib import json +import os import re from collections.abc import MutableMapping, MutableSequence from typing import ( @@ -17,10 +20,6 @@ from typing import ( Tuple, Union, ) - -import hashlib -import importlib -import os from urllib.parse import quote from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -339,34 +338,36 @@ def iter_known_tool_name_spellings(tool_name: str, server: MCPServer) -> Iterato yield add_server_prefix_to_name(tool_name, prefix) +def openapi_tool_name(operation_id: str) -> str: + """Return the tool name ``_register_openapi_tools`` registers ``operation_id`` under. + + The single transform between a spec's operationId and the name the gateway serves. + Policy recovers the link by replaying this exact function, which is what keeps it from + deciding for a tool it does not name: two operationIds that register as two tools + necessarily normalize to two names here, because this is the map that registered them. + """ + return operation_id.replace(" ", "_").lower() + + def match_known_tool_name(tool_name: str, server: MCPServer, names: Iterable[str]) -> str | None: """Return the entry of ``names`` that denotes ``tool_name`` on ``server``, else ``None``. The single question every tool-name-keyed site asks: the allow list, the deny list, ``allowed_params`` and the discovery filter, so discovery hides exactly what dispatch - refuses. It spans every spelling routing accepts and no more. A tool's identity is its - exact name, because routing dispatches two names differing only in case as two tools, - and folding case here would let one policy decide both. + refuses. It spans every spelling routing accepts and no more, because a tool's identity + is the exact name routing dispatches; anything looser lets one policy decide two tools. - OpenAPI servers are the exception, and not a fuzzy one. ``_register_openapi_tools`` - rewrites every operationId through ``sanitize_openapi_tool_name``, so configuration - holding the spec's own spelling never equals the registered name; replaying that exact - rewrite on the entry recovers the link. It cannot merge identities, because every name - it produces is lowercased, so no two tools on such a server differ only in case. + On an OpenAPI server the configured entry holds the spec's operationId while routing + holds :func:`openapi_tool_name` of it, so both sides go through that map first. Doing it + with the registering function rather than a lookalike is the whole safety argument: a + coarser one collapses operationIds that registration keeps apart. Callers read the returned entry rather than testing a container's values, which is what stops an explicitly empty ``allowed_params`` list from reading as "nothing configured". """ - from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( - sanitize_openapi_tool_name, - ) - - spellings = set(iter_known_tool_name_spellings(tool_name, server)) - exact = next((name for name in names if name in spellings), None) - if exact is not None or not getattr(server, "spec_path", None): - return exact - sanitized = {sanitize_openapi_tool_name(spelling) for spelling in spellings} - return next((name for name in names if sanitize_openapi_tool_name(name) in sanitized), None) + normalize = openapi_tool_name if getattr(server, "spec_path", None) else str + spellings = {normalize(spelling) for spelling in iter_known_tool_name_spellings(tool_name, server)} + return next((name for name in names if normalize(name) in spellings), None) def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]: 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 cdef50ad6b5..db5b64ef131 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 @@ -9511,6 +9511,16 @@ class TestServerToolListsHonorThePrefixBoundary: assert exc_info.value.status_code == 403 assert "include_internal" in exc_info.value.detail["error"] + @pytest.mark.asyncio + async def test_an_entry_does_not_decide_an_operation_id_registration_keeps_separate(self): + server = self._aliased_server(disallowed_tools=["foo/bar"], spec_path="/specs/petstore.yaml") + + with pytest.raises(HTTPException) as exc_info: + await self._run_check(server, "foo/bar") + + assert exc_info.value.status_code == 403 + await self._run_check(server, "foo.bar") + @pytest.mark.asyncio async def test_a_blocklist_entry_does_not_reach_a_case_variant_sibling_tool(self): server = self._aliased_server(disallowed_tools=["petstore-getPet"]) From 4d2b7224fd637c0304789df55384bcd64e14c364 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Thu, 30 Jul 2026 22:38:54 -0700 Subject: [PATCH 25/26] fix(mcp): annotate connected-app reachability on the gateway connect page (#34867) * fix(mcp): annotate connected-app reachability on the gateway connect page The MCP connect page resolved its server grid through the dashboard identity (admin shortcut or view_all returns the whole registry) while the gateway DCR session it sets up resolves servers as an admitted subject through grant sources only, so the page showed servers and tool counts the session is never served. GET /v1/mcp/server now accepts connected_app_view=true and stamps each returned server with connected_app_reachable, computed by the same _reload_admitted_user + get_allowed_mcp_servers pair the live session uses. The connect page requests the flag in connect mode and renders unreachable servers dimmed with a label, excluded from the Connected count and tool-count fetches. Failure to build the admitted set marks everything unreachable, which matches what such a session would actually be served. Default behavior without the param is unchanged for every existing consumer. * fix(mcp): block connecting unavailable servers from the connect-mode detail view A server the connect page marks unavailable could still be added through its detail view Connect action, so the selection could contain servers the connected-app session is never served. The unavailability decision now lives in one predicate, connectUnavailabilityLabel, consumed by the card indicator, the detail view action area, the toggle-on path, the oauth auto-select effect, and the Connected count, so no interaction path can disagree with the label. This also closes the same pre-existing hole for servers marked not supported on this connection, whose detail view likewise offered Connect, and removes a grandfathered nested ternary, ratcheting the eslint suppressions baseline down * fix(mcp): hide unreachable servers on the connect page instead of dimming them Product decision: the connect page should only show what a connected-app session will actually be served, so annotated-unreachable servers are now filtered out of the connect-mode list at fetch time rather than rendered dimmed. Unsupported auth types keep their existing dimmed label since they are a property of the server, not the caller. A user with zero reachable servers gets an explanatory empty state pointing at grants. The list filter is the single source: counts, tabs, auto-select, detail view, and tool-count fetches all derive from the already-filtered state * fix(mcp): guarantee the connect view lists every session-reachable server The connect view's membership came from the dashboard resolver with the admitted-subject answer only annotated on top, so a server reachable by the session but missing from the dashboard list would be invisible on the page; an under-report, the mirror of the bug this PR fixes. The connect view now unions in any session-reachable server the dashboard resolver did not list, built from the registry and redacted through the same ladder, so page membership equals the admitted set by construction in both directions * fix(mcp): honor connected_app_view only for the dashboard UI session credential The reachability view resolves through the owning user's admitted identity, so a caller-passed virtual key could use the param to enumerate servers beyond its own scope (ids, names, descriptions of the owner's wider grants). The view is now gated on is_ui_session_credential, a predicate factored out of resolve_ui_session_team_ids so the two user-identity widening sites share one trust boundary: the SSO-minted dashboard session token acting as its user. Any other credential gets the param as a no-op and the admitted resolver is never consulted for it * fix(mcp): resolve UI sessions with the admitted-user context everywhere, not per endpoint The list endpoint unioned in session-reachable servers itself while tool counts, Connect actions, and credential endpoints still authorized through build_effective_auth_contexts, whose contexts carry team grants but never the user row's own object permission; a user-granted server could render on the connect page while every interaction on it failed. The admitted-user context (the same auth a gateway session resolves with) is now appended inside build_effective_auth_contexts for UI session credentials, so the page list and every per-server action endpoint answer identically, and the list endpoint's one-off union is deleted. Caller-passed keys are still never widened (is_ui_session_credential gate inside the context builder) and a reload failure falls back to team contexts only * fix(mcp): resolve non-admin dashboard sessions as the admitted subject on tool routes Server reachability on the REST tool routes came from the widened context union while tool permission checks ran on the bare session key, which carries no object permission, so a dashboard user could invoke tools their user-level grant excludes. Rather than bookkeeping which context granted which server, the routes now choose one principal at the boundary: acting_user_auth swaps a non-admin UI session for the admitted-subject auth, the same identity a gateway session resolves with, so reachability, per-source fail-closed tool ceilings, rate limits, and billing attribution all bind through the admitted arms that already exist downstream. Admin sessions keep their operator view and caller-passed credentials are never widened. One swap point per route, no per-server principal picking, no parallel permission logic * fix(mcp): derive the connect page's detail view from the reachable server list The detail view held its own copy of the server object, so it outlived the list it came from. When a refetch dropped that server as unreachable, the open detail view kept rendering it and its Connect action still ran: the guard looked the server back up by id or name in the current list, found nothing, and fell through, because a missing target read as "nothing to block" rather than "no longer connectable" Store the selected server's id and derive the row from the list instead. A server the list no longer carries cannot be the detail view's subject, so the stale render, the stale tools query and the guard bypass stop being reachable states rather than being blocked one at a time. handleToggle now takes the server it is toggling, which deletes the lookup that could miss at all * refactor(mcp): one owner for the identity a dashboard session acts as Three call sites reloaded the admitted subject independently, and the management endpoint carried its own copy of the reload, the HTTPException swallow and the logging. admitted_user_context is now the only place that answers "what user identity does this dashboard session act as", and the connected-app reachability helper reads it, which also drops its dead empty-user_id branch That owner now carries the request's tracing span onto the admitted principal. _reload_admitted_user builds a fresh auth from the user row and has no span of its own, so swapping it in on the REST tool routes silently detached every downstream lookup and the tool-call logging from the request's trace Toolset scoping and the acting-as-user swap are mutually exclusive, so they now share one owner on the tools list route. The admitted subject resolves per grant source and a team source deliberately carries none of the caller's object_permission, so a toolset narrowing layered on top would evaporate on every team-granted server: the request would be admitted through the toolset grant and then served tools from servers the toolset never named. A request carrying a toolset name stays on the caller's own credential, exactly as it did before the swap * fix(mcp): commit every async connect-page write against the list as it stands Three continuations in the panel decided against state captured before their await and committed after it, so a reachability refetch landing in between could not be seen handleToggle validated the server at click time and then, once listMCPTools resolved, wrote its name into the selection whatever the list had since become; a server the refresh had dropped was selected anyway. It now re-asks connectableNow at the commit, and that predicate resolves the id against the current list, so absence fails closed instead of reading as nothing to block The load pipeline was worse, because its cancel flag was shared across runs: the successor's effect body reset it to false before the predecessor's fetch resolved, so a superseded load could still run setServers and put the dropped server back on the page outright. The flag is now a per-effect local that only that run's cleanup can clear, which is also what makes unmount stop the chunked tool-count loop again. The load passes its own liveness check down to the tool-count and oauth-status writes rather than having them consult a flag they share with every other run * fix(mcp): write the connect-page server list to its ref as it is committed connectableNow resolves a server id against serversRef, but that ref was a mirror kept in step by a passive effect, so it lagged the state it mirrored by however long React took to render and flush. A continuation resolving inside that window read the previous list: the commit-time reachability check would find a server the refetch had already dropped, call it connectable, and select it, which is the mismatch the check exists to prevent The lag was the whole defect, so the mirror is gone. commitServers writes the ref and the state together, at the one point the list is ever replaced, and the ref is now never older than the last committed list. Readers that want the newest answer (connectableNow, the oauth auto-select effect) get it; rendering still derives from state, so what is on screen is unchanged Pinned by a test that resolves the refetch and the in-flight Connect in the same tick, with no render flushed between them, which is the interleaving the earlier regression could not reach. The two prop mirrors are deliberately untouched: their staleness is inherent to appending to a parent-owned list from an async callback rather than caused by the mirror, and no reachability decision reads them --- litellm/models/mcp_server.py | 1 + .../mcp_server/rest_endpoints.py | 17 +- .../mcp_server/ui_session_utils.py | 72 ++++- .../mcp_management_endpoints.py | 23 ++ .../mcp_server/test_rest_endpoints.py | 171 ++++++++++++ .../mcp_server/test_ui_session_utils.py | 138 ++++++++++ .../test_mcp_management_endpoints.py | 252 ++++++++++++++++++ ui/litellm-dashboard/eslint-suppressions.json | 2 +- .../src/components/chat/MCPAppsPanel.test.tsx | 209 ++++++++++++++- .../src/components/chat/MCPAppsPanel.tsx | 234 +++++++++------- .../src/components/mcp_tools/types.tsx | 1 + .../src/components/networking.tsx | 7 +- 12 files changed, 1015 insertions(+), 112 deletions(-) diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py index 23b26bd8e89..e428d20f99d 100644 --- a/litellm/models/mcp_server.py +++ b/litellm/models/mcp_server.py @@ -102,6 +102,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): byok_description: List[str] = Field(default_factory=list) byok_api_key_help_url: Optional[str] = None has_user_credential: Optional[bool] = None + connected_app_reachable: bool | None = None source_url: Optional[str] = None timeout: Optional[float] = None max_concurrent_requests: Optional[int] = None diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 9b51513f4ac..b8604917825 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -26,6 +26,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( list_fault_http_status, ) from litellm.proxy._experimental.mcp_server.ui_session_utils import ( + acting_user_auth, build_effective_auth_contexts, ) from litellm.proxy._experimental.mcp_server.utils import ( @@ -667,13 +668,19 @@ if MCP_AVAILABLE: """Coerce an Optional[str] Query param to str|None, dropping unresolved FastAPI defaults.""" return value if isinstance(value, str) else None - async def _resolve_toolset_scope( + async def _resolve_acting_auth( toolset_name: str | None, user_api_key_dict: UserAPIKeyAuth, ) -> UserAPIKeyAuth: - """Resolve ``toolset_name`` to its scoped ``UserAPIKeyAuth``, or return unchanged.""" + """The one credential this tools request acts as. + + A toolset name narrows the caller's own credential to that toolset; otherwise a dashboard + session is swapped for its admitted subject. The two are mutually exclusive by construction, + which is why they share an owner: the admitted subject resolves per grant source and a team + source deliberately carries none of the caller's ``object_permission``, so a toolset + narrowing layered on top would evaporate on every team-granted server.""" if not toolset_name: - return user_api_key_dict + return await acting_user_auth(user_api_key_dict) from litellm.proxy.utils import get_prisma_client_or_throw @@ -731,6 +738,7 @@ if MCP_AVAILABLE: try: mcp_server_name = _as_query_str(mcp_server_name) toolset_name = _as_query_str(toolset_name) + user_api_key_dict = await _resolve_acting_auth(toolset_name, user_api_key_dict) # The full catalog (allowlist filter skipped) is admin-only so the # REST endpoint can't be used to enumerate deliberately-disabled tools. @@ -738,8 +746,6 @@ if MCP_AVAILABLE: include_disabled_tools and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN ) - user_api_key_dict = await _resolve_toolset_scope(toolset_name, user_api_key_dict) - if server_id is None: server_id = mcp_server_name @@ -928,6 +934,7 @@ if MCP_AVAILABLE: ) try: + user_api_key_dict = await acting_user_auth(user_api_key_dict) data = await request.json() tool_name = data.get("name") diff --git a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py b/litellm/proxy/_experimental/mcp_server/ui_session_utils.py index 1b37b884987..d1d28574988 100644 --- a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py +++ b/litellm/proxy/_experimental/mcp_server/ui_session_utils.py @@ -1,9 +1,11 @@ -"""Helpers to resolve real team contexts for UI session tokens.""" +"""Helpers to resolve the identity a dashboard UI session token acts as.""" from __future__ import annotations from typing import List +from fastapi import HTTPException + from litellm._logging import verbose_logger from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy._types import UserAPIKeyAuth @@ -23,12 +25,19 @@ def clone_user_api_key_auth_with_team( return cloned_auth +def is_ui_session_credential(user_api_key_auth: UserAPIKeyAuth) -> bool: + """Whether the caller is the dashboard's SSO-minted session token acting as its user, + the only credential shape allowed to widen a request to the owning user's identity.""" + + return user_api_key_auth.team_id == UI_SESSION_TOKEN_TEAM_ID and bool(user_api_key_auth.user_id) + + async def resolve_ui_session_team_ids( user_api_key_auth: UserAPIKeyAuth, ) -> List[str]: """Resolve the real team ids backing a UI session token.""" - if user_api_key_auth.team_id != UI_SESSION_TOKEN_TEAM_ID or not user_api_key_auth.user_id: + if not is_ui_session_credential(user_api_key_auth): return [] from litellm.proxy.auth.auth_checks import get_user_object @@ -68,12 +77,63 @@ async def resolve_ui_session_team_ids( return resolved_team_ids +async def admitted_user_context(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth | None: + """THE owner of "resolve this dashboard session's user identity": the same admitted-subject auth a + gateway OAuth session for this user resolves with, carrying the user row's own object permission, + on this request's tracing span. None for any other credential (a caller-passed key is never + widened) and on reload failure, which every caller reads as "no user-level identity available".""" + + user_id = user_api_key_auth.user_id + if not is_ui_session_credential(user_api_key_auth) or user_id is None: + return None + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + try: + admitted = await MCPRequestHandler._reload_admitted_user(user_id) + except HTTPException as e: + verbose_logger.warning(f"MCP dashboard session: admitted-subject reload failed for {user_id}: {e.detail}") + return None + return admitted.model_copy(update={"parent_otel_span": user_api_key_auth.parent_otel_span}) + + +async def acting_user_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth: + """The principal acting-as-user MCP routes resolve permissions with. A non-admin dashboard + session acts as the admitted subject, the same identity a gateway session resolves with, so + server reachability, per-source tool ceilings, rate limits, and billing bind identically on + both surfaces. An admin session keeps its operator view and any caller-passed credential is + returned unchanged, never widened. + + Do not combine this with a narrowing that rewrites a single credential's ``object_permission`` + (toolset scope): the admitted subject resolves per grant source and a team source deliberately + carries none of the caller's own grants, so the narrowing would silently evaporate on every + team-granted server. A request carrying such a scope keeps the caller's own credential.""" + + if not is_ui_session_credential(user_api_key_auth): + return user_api_key_auth + from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view + + if _user_has_admin_view(user_api_key_auth): + return user_api_key_auth + admitted = await admitted_user_context(user_api_key_auth) + return admitted if admitted is not None else user_api_key_auth + + async def build_effective_auth_contexts( user_api_key_auth: UserAPIKeyAuth, ) -> List[UserAPIKeyAuth]: - """Return auth contexts that reflect the actual teams for UI session tokens.""" + """Every auth context a management or listing surface must resolve a UI session token through: + one per real team backing the session, plus the session user's own admitted identity, so a grant + made directly to the user row is as visible to the dashboard as it is to a gateway session.""" resolved_team_ids = await resolve_ui_session_team_ids(user_api_key_auth) - if resolved_team_ids: - return [clone_user_api_key_auth_with_team(user_api_key_auth, team_id) for team_id in resolved_team_ids] - return [user_api_key_auth] + team_contexts = ( + [clone_user_api_key_auth_with_team(user_api_key_auth, team_id) for team_id in resolved_team_ids] + if resolved_team_ids + else [user_api_key_auth] + ) + admitted_context = await admitted_user_context(user_api_key_auth) + if admitted_context is None: + return team_contexts + return [*team_contexts, admitted_context] diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 64cc13a5543..dcc4dc36d83 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -148,7 +148,9 @@ if MCP_AVAILABLE: global_mcp_server_manager, ) from litellm.proxy._experimental.mcp_server.ui_session_utils import ( + admitted_user_context, build_effective_auth_contexts, + is_ui_session_credential, ) from litellm.proxy._types import ( LiteLLM_MCPServerTable, @@ -939,6 +941,16 @@ if MCP_AVAILABLE: aggregated.setdefault(server.server_id, server) return list(aggregated.values()) + async def _connected_app_reachable_server_ids(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: + """Server ids a connected app authorized by this dashboard user is served on the aggregate + MCP endpoint, resolved through the one owner of the admitted subject so the page and the + session cannot drift. Empty when that identity cannot be built, which is the true answer: + the same user cannot open a gateway session either.""" + admitted = await admitted_user_context(user_api_key_dict) + if admitted is None: + return frozenset() + return frozenset(await global_mcp_server_manager.get_allowed_mcp_servers(admitted)) + @router.get( "/server", description="Returns the mcp server list with associated teams", @@ -953,6 +965,12 @@ if MCP_AVAILABLE: "servers the team has access to plus globally available (allow_all_keys) servers. " "Used by the Create Key UI to show team-scoped MCP servers.", ), + connected_app_view: bool = Query( + False, + description="Annotate each returned server with connected_app_reachable: whether a " + "connected app authorized by the calling user (a gateway OAuth session) is served " + "this server on the aggregate MCP endpoint.", + ), ): """ Get all of the configured mcp servers for the user in the db with their associated teams @@ -1009,6 +1027,11 @@ if MCP_AVAILABLE: servers = await _resolve_accessible_mcp_servers(user_api_key_dict) redacted_mcp_servers = _redact_mcp_credentials_list(servers) + if connected_app_view is True and is_ui_session_credential(user_api_key_dict): + reachable_ids = await _connected_app_reachable_server_ids(user_api_key_dict) + for server in redacted_mcp_servers: + server.connected_app_reachable = server.server_id in reachable_ids + # augment the mcp servers with public status if litellm.public_mcp_servers is not None: for server in redacted_mcp_servers: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index cec62f79e33..329b6d5c45d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -769,8 +769,179 @@ class TestListToolsRestAPI: assert captured["server"] is stub_server assert result["tools"] == ["tool-1"] assert result["error"] is None + + async def test_non_admin_ui_session_resolves_as_admitted_subject(self, monkeypatch): + """LIT-4861: a non-admin dashboard session must act as the admitted subject on this + route, so server reachability AND tool ceilings bind to the user's grants exactly as + they do for a gateway session, never to the bare session key.""" + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + session_auth = UserAPIKeyAuth( + team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user" + ) + admitted_auth = UserAPIKeyAuth(user_id="grant-user", org_id="admitted-org") + + async def fake_reload(user_id): + assert user_id == "grant-user" + return admitted_auth + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + fake_reload, + ) + + seen_server_resolution_auths = [] + + async def fake_get_allowed_mcp_servers(user_api_key_auth=None, **kwargs): + seen_server_resolution_auths.append(user_api_key_auth) + return ["server-1"] + + class StubServer: + alias = "server-1" + server_name = "server-1" + name = "stub" + allowed_tools = None + mcp_info = {"server_name": "stub"} + available_on_public_internet = True + + stub_server = StubServer() + captured = {} + + async def fake_get_tools( + server, + server_auth_header, + raw_headers=None, + user_api_key_auth=None, + extra_headers=None, + apply_tool_filters=True, + ): + captured["user_api_key_auth"] = user_api_key_auth + return ["tool-1"] + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints, + "_get_tools_for_single_server", + fake_get_tools, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + result = await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + user_api_key_dict=session_auth, + ) + + resolved = [*seen_server_resolution_auths, captured["user_api_key_auth"]] + assert seen_server_resolution_auths + assert all(a.org_id == "admitted-org" and a.team_id is None for a in resolved) + assert result["tools"] == ["tool-1"] assert result["message"] == "Successfully retrieved tools" + async def test_toolset_scoped_request_keeps_the_caller_credential(self, monkeypatch): + """LIT-4861: the admitted subject resolves per grant source and a team source deliberately + carries none of the caller's own object_permission, so a toolset narrowing layered on top + would evaporate on every team-granted server. A toolset-scoped request therefore stays on + the caller's own credential, exactly as it did before the acting-as-user swap.""" + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + session_auth = UserAPIKeyAuth( + team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user" + ) + scoped_auth = UserAPIKeyAuth( + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="toolset-scope", + mcp_servers=["toolset-server-1"], + ) + ) + reload_calls: list[str] = [] + scope_inputs: list[UserAPIKeyAuth] = [] + + async def record_reload(user_id): + reload_calls.append(user_id) + return UserAPIKeyAuth(user_id=user_id) + + class StubToolset: + toolset_id = "toolset-1" + + class StubServer: + alias = "toolset-server-1" + server_name = "toolset-server-1" + name = "toolset-server-1" + allowed_tools = None + mcp_info = {"server_name": "toolset-server-1"} + available_on_public_internet = True + + stub_server = StubServer() + + async def fake_get_toolset_by_name_cached(prisma_client, toolset_name): + return StubToolset() + + async def fake_apply_toolset_scope(user_api_key_auth, toolset_id): + scope_inputs.append(user_api_key_auth) + return scoped_auth + + async def fake_get_allowed_mcp_servers(user_api_key_auth=None, **kwargs): + assert user_api_key_auth is scoped_auth + return ["toolset-server-1"] + + async def fake_get_tools(server, server_auth_header, *args, **kwargs): + return ["toolset-tool-1"] + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + record_reload, + ) + monkeypatch.setattr( + "litellm.proxy.utils.get_prisma_client_or_throw", + lambda *args, **kwargs: MagicMock(), + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_toolset_by_name_cached", + fake_get_toolset_by_name_cached, + raising=False, + ) + monkeypatch.setattr(rest_endpoints, "_apply_toolset_scope", fake_apply_toolset_scope, raising=False) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "toolset-server-1" else None, + raising=False, + ) + monkeypatch.setattr(rest_endpoints, "_get_tools_for_single_server", fake_get_tools, raising=False) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + result = await rest_endpoints.list_tool_rest_api( + request, + server_id=None, + toolset_name="research_tools", + user_api_key_dict=session_auth, + ) + + assert result["tools"] == ["toolset-tool-1"] + assert scope_inputs == [session_auth] + assert reload_calls == [] + async def test_include_disabled_tools_is_admin_only(self, monkeypatch): """include_disabled_tools skips the allowlist filter only for PROXY_ADMIN; a non-admin passing it stays filtered so the REST endpoint can't be used diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py index 52120207f76..cd4cba51908 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py @@ -3,6 +3,7 @@ from types import SimpleNamespace from unittest.mock import AsyncMock import pytest +from fastapi import HTTPException from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy._types import UserAPIKeyAuth @@ -120,3 +121,140 @@ async def test_build_effective_auth_contexts_handles_unpicklable_parent_span( assert contexts[0].team_id == "team-span" assert contexts[0].parent_otel_span is parent_span + + +@pytest.mark.asyncio +async def test_build_effective_auth_contexts_appends_admitted_user_context(monkeypatch): + """LIT-4861: the dashboard session must resolve with the user's admitted identity so the + page list and every per-server action endpoint see user-level grants the same way the + gateway session does.""" + user_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-42") + admitted_auth = UserAPIKeyAuth(user_id="user-42") + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.ui_session_utils.resolve_ui_session_team_ids", + AsyncMock(return_value=["team-one"]), + ) + reload_mock = AsyncMock(return_value=admitted_auth) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + reload_mock, + ) + + contexts = await build_effective_auth_contexts(user_auth) + + assert contexts[-1].user_id == "user-42" and contexts[-1].team_id is None + assert [ctx.team_id for ctx in contexts[:-1]] == ["team-one"] + reload_mock.assert_awaited_once_with("user-42") + + +@pytest.mark.asyncio +async def test_build_effective_auth_contexts_never_widens_caller_passed_keys(monkeypatch): + normal_user = UserAPIKeyAuth(team_id="regular-team", user_id="user-1") + reload_mock = AsyncMock() + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + reload_mock, + ) + + contexts = await build_effective_auth_contexts(normal_user) + + assert contexts == [normal_user] + reload_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_build_effective_auth_contexts_survives_admitted_reload_failure(monkeypatch): + user_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-9") + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.ui_session_utils.resolve_ui_session_team_ids", + AsyncMock(return_value=["team-a"]), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + AsyncMock(side_effect=HTTPException(status_code=503, detail="db down")), + ) + + contexts = await build_effective_auth_contexts(user_auth) + + assert [ctx.team_id for ctx in contexts] == ["team-a"] + + +@pytest.mark.asyncio +async def test_acting_user_auth_returns_admitted_subject_for_non_admin_sessions(monkeypatch): + """LIT-4861: acting-as-user MCP routes must resolve a non-admin dashboard session as the + admitted subject so tool ceilings, reachability, and limits bind exactly as on /mcp.""" + from litellm.proxy._experimental.mcp_server.ui_session_utils import acting_user_auth + + user_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-42", user_role="internal_user") + admitted_auth = UserAPIKeyAuth(user_id="user-42") + reload_mock = AsyncMock(return_value=admitted_auth) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + reload_mock, + ) + + result = await acting_user_auth(user_auth) + + assert result.user_id == "user-42" and result.team_id is None + reload_mock.assert_awaited_once_with("user-42") + + +@pytest.mark.asyncio +async def test_acting_user_auth_keeps_admin_sessions_and_passed_keys_unchanged(monkeypatch): + from litellm.proxy._experimental.mcp_server.ui_session_utils import acting_user_auth + + reload_mock = AsyncMock() + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + reload_mock, + ) + + admin_session = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="admin-1", user_role="proxy_admin") + assert await acting_user_auth(admin_session) is admin_session + + passed_key = UserAPIKeyAuth(team_id="regular-team", user_id="user-1", user_role="internal_user") + assert await acting_user_auth(passed_key) is passed_key + + reload_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_acting_user_auth_falls_back_to_session_auth_on_reload_failure(monkeypatch): + from litellm.proxy._experimental.mcp_server.ui_session_utils import acting_user_auth + + user_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-9", user_role="internal_user") + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + AsyncMock(side_effect=HTTPException(status_code=503, detail="db down")), + ) + + assert await acting_user_auth(user_auth) is user_auth + + +@pytest.mark.asyncio +async def test_admitted_user_context_carries_the_request_span(monkeypatch): + """Swapping the principal must not drop the request: the admitted subject is rebuilt from the + user row and carries no span of its own, so every consumer would otherwise lose trace linkage + for the resolution and logging it drives.""" + from litellm.proxy._experimental.mcp_server.ui_session_utils import acting_user_auth + + class DummySpan: + def __init__(self) -> None: + self._lock = threading.RLock() + + parent_span = DummySpan() + user_auth = UserAPIKeyAuth( + team_id=UI_SESSION_TOKEN_TEAM_ID, + user_id="user-42", + user_role="internal_user", + parent_otel_span=parent_span, + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + AsyncMock(return_value=UserAPIKeyAuth(user_id="user-42")), + ) + + assert (await acting_user_auth(user_auth)).parent_otel_span is parent_span + assert (await build_effective_auth_contexts(user_auth))[-1].parent_otel_span is parent_span diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 53dda8f6648..bf119c4fb2f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -6138,3 +6138,255 @@ def test_bundled_openapi_registry_parses_and_entries_are_well_formed(): ) for tool in entry.get("key_tools", []): assert tool.get("name") and tool.get("description"), f"{entry['name']}: malformed key_tool" + + +class TestConnectedAppViewAnnotation: + """LIT-4861: GET /v1/mcp/server?connected_app_view=true must annotate each server with + whether the caller's gateway OAuth sessions (connected apps) are served it on /mcp. + The view is honored only for the dashboard's UI session credential; a caller-passed + virtual key must never be widened to its owning user's identity.""" + + def _ui_session_auth(self, user_role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN) -> UserAPIKeyAuth: + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + return generate_mock_user_api_key_auth(user_role=user_role, team_id=UI_SESSION_TOKEN_TEAM_ID) + + def _mock_manager(self, servers, reachable_ids): + mock_manager = MagicMock() + mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=servers) + mock_manager.get_all_mcp_servers_unfiltered = AsyncMock(return_value=servers) + mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=reachable_ids) + return mock_manager + + def _servers(self): + return [ + generate_mock_mcp_server_db_record(server_id="server-1", alias="Granted"), + generate_mock_mcp_server_db_record(server_id="server-2", alias="Ungranted"), + ] + + @pytest.mark.asyncio + async def test_connected_app_view_annotates_reachability_via_admitted_resolver(self): + caller_auth = self._ui_session_auth() + admitted_auth = UserAPIKeyAuth(user_id="test_user_id") + admitted_auth.mcp_admitted_user_subject = True + mock_manager = self._mock_manager(self._servers(), ["server-1"]) + reload_mock = AsyncMock(return_value=admitted_auth) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[caller_auth]), + ), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + reload_mock, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True) + + flags = {server.server_id: server.connected_app_reachable for server in result} + assert flags == {"server-1": True, "server-2": False} + reload_mock.assert_awaited_once_with("test_user_id") + mock_manager.get_allowed_mcp_servers.assert_awaited_once_with(admitted_auth) + + @pytest.mark.asyncio + async def test_connected_app_view_stamps_view_all_list_and_survives_non_admin_sanitizer(self): + """view_all preempts the manager's admin shortcut with a second whole-registry + shortcut; the annotation must still land, and must survive the non-admin sanitizer.""" + caller_auth = self._ui_session_auth(user_role=LitellmUserRoles.INTERNAL_USER) + mock_manager = self._mock_manager(self._servers(), ["server-2"]) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_user_mcp_management_mode", + return_value="view_all", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + AsyncMock(return_value=UserAPIKeyAuth(user_id="test_user_id")), + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True) + + mock_manager.get_all_mcp_servers_unfiltered.assert_awaited_once() + flags = {server.server_id: server.connected_app_reachable for server in result} + assert flags == {"server-1": False, "server-2": True} + + @pytest.mark.asyncio + async def test_connected_app_view_lists_user_granted_servers_via_admitted_context(self): + """A server granted only through the user's own object permission must be listed and + flagged reachable: the REAL build_effective_auth_contexts appends the admitted-user + context, so the page and every action endpoint resolve it identically.""" + caller_auth = self._ui_session_auth(user_role=LitellmUserRoles.INTERNAL_USER) + admitted_auth = UserAPIKeyAuth(user_id="test_user_id", org_id="admitted-org") + listed_row = generate_mock_mcp_server_db_record(server_id="server-1", alias="TeamGranted") + user_granted_row = generate_mock_mcp_server_db_record(server_id="server-2", alias="UserGranted") + + async def per_context_servers(user_api_key_auth=None): + if user_api_key_auth is not None and user_api_key_auth.org_id == "admitted-org": + return [listed_row, user_granted_row] + return [listed_row] + + mock_manager = MagicMock() + mock_manager.get_all_allowed_mcp_servers = AsyncMock(side_effect=per_context_servers) + mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["server-1", "server-2"]) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy._experimental.mcp_server.ui_session_utils.resolve_ui_session_team_ids", + AsyncMock(return_value=[]), + ), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + AsyncMock(return_value=admitted_auth), + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True) + + flags = {server.server_id: server.connected_app_reachable for server in result} + assert flags == {"server-1": True, "server-2": True} + + @pytest.mark.asyncio + async def test_connected_app_view_fails_closed_when_admitted_reload_fails(self): + caller_auth = self._ui_session_auth() + mock_manager = self._mock_manager(self._servers(), ["server-1"]) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[caller_auth]), + ), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + AsyncMock(side_effect=HTTPException(status_code=401, detail="expired")), + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True) + + assert all(server.connected_app_reachable is False for server in result) + + @pytest.mark.asyncio + async def test_connected_app_view_off_leaves_field_unset(self): + caller_auth = generate_mock_user_api_key_auth() + mock_manager = self._mock_manager(self._servers(), ["server-1"]) + reload_mock = AsyncMock(return_value=UserAPIKeyAuth(user_id="test_user_id")) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[caller_auth]), + ), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + reload_mock, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth) + + assert all(server.connected_app_reachable is None for server in result) + reload_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_connected_app_view_userless_ui_credential_leaves_field_unset(self): + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + caller_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="test_api_key", team_id=UI_SESSION_TOKEN_TEAM_ID + ) + caller_auth.user_id = None + mock_manager = self._mock_manager(self._servers(), ["server-1"]) + reload_mock = AsyncMock(return_value=UserAPIKeyAuth(user_id="test_user_id")) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[caller_auth]), + ), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + reload_mock, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True) + + assert all(server.connected_app_reachable is None for server in result) + reload_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_connected_app_view_ignored_for_caller_passed_virtual_keys(self): + """A virtual key the user passes themselves is never widened to the owning user's + identity: the view param is a no-op and the admitted resolver is never consulted.""" + caller_auth = generate_mock_user_api_key_auth(team_id="some-real-team") + mock_manager = self._mock_manager(self._servers(), ["server-1"]) + reload_mock = AsyncMock(return_value=UserAPIKeyAuth(user_id="test_user_id")) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[caller_auth]), + ), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler._reload_admitted_user", + reload_mock, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_all_mcp_servers, + ) + + result = await fetch_all_mcp_servers(user_api_key_dict=caller_auth, connected_app_view=True) + + assert all(server.connected_app_reachable is None for server in result) + reload_mock.assert_not_awaited() diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 6819b2851f5..4cb3ebe7e16 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2729,7 +2729,7 @@ "count": 1 }, "no-nested-ternary": { - "count": 6 + "count": 4 } }, "src/components/chat/MCPConnectPicker.tsx": { diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx index 656ef157363..5fc3e75195b 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx @@ -1,6 +1,6 @@ import React from "react"; -import { render, screen, fireEvent } from "@testing-library/react"; -import { describe, it, expect, vi, afterEach } from "vitest"; +import { render, screen, fireEvent, waitFor, act } from "@testing-library/react"; +import { describe, it, expect, vi, afterEach, beforeEach } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import MCPAppsPanel from "./MCPAppsPanel"; import { fetchMCPServers, listMCPTools } from "../networking"; @@ -86,3 +86,208 @@ describe("MCPAppsPanel logos", () => { expect(screen.getByAltText("local_logo logo").getAttribute("src")).toBe("/litellm/ui/assets/logos/github.svg"); }); }); + +const connectServers = [ + { + server_id: "s-reach", + server_name: "reachable_srv", + auth_type: "none", + connected_app_reachable: true, + }, + { + server_id: "s-unreach", + server_name: "unreachable_srv", + auth_type: "none", + connected_app_reachable: false, + }, +] as MCPServer[]; + +const renderConnectPanel = (connectMode: boolean, selectedServers: string[] = []) => + render( + + + , + ); + +describe("MCPAppsPanel connected-app reachability (LIT-4861)", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("requests the connected-app view and hides unreachable servers in connect mode", async () => { + vi.mocked(fetchMCPServers).mockResolvedValue(connectServers); + vi.mocked(listMCPTools).mockResolvedValue({ tools: [] }); + + renderConnectPanel(true, ["reachable_srv", "unreachable_srv"]); + + expect(await screen.findByText("reachable_srv")).toBeInTheDocument(); + expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, true); + expect(screen.queryByText("unreachable_srv")).not.toBeInTheDocument(); + expect(screen.getByText("Connected (1)")).toBeInTheDocument(); + const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]); + expect(toolCountFetchedIds).toContain("s-reach"); + expect(toolCountFetchedIds).not.toContain("s-unreach"); + }); + + it("blocks connecting an unsupported server from the detail view in connect mode", async () => { + const detailServers = [ + ...connectServers, + { + server_id: "s-unsup", + server_name: "unsupported_srv", + auth_type: "oauth2_token_exchange", + connected_app_reachable: true, + }, + ] as MCPServer[]; + vi.mocked(fetchMCPServers).mockResolvedValue(detailServers); + vi.mocked(listMCPTools).mockResolvedValue({ tools: [] }); + + renderConnectPanel(true); + + fireEvent.click(await screen.findByText("unsupported_srv")); + expect(await screen.findByRole("heading", { name: "unsupported_srv" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Connect" })).not.toBeInTheDocument(); + expect(screen.getByText("Not supported on this connection")).toBeInTheDocument(); + }); + + it("keeps the detail-view Connect action outside connect mode", async () => { + vi.mocked(fetchMCPServers).mockResolvedValue(connectServers); + vi.mocked(listMCPTools).mockResolvedValue({ tools: [] }); + + renderConnectPanel(false); + + fireEvent.click(await screen.findByText("unreachable_srv")); + expect(await screen.findByRole("heading", { name: "unreachable_srv" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Connect" })).toBeInTheDocument(); + }); + + it("ignores the flag and skips no server outside connect mode", async () => { + vi.mocked(fetchMCPServers).mockResolvedValue(connectServers); + vi.mocked(listMCPTools).mockResolvedValue({ tools: [] }); + + renderConnectPanel(false, ["reachable_srv", "unreachable_srv"]); + + expect(await screen.findByText("unreachable_srv")).toBeInTheDocument(); + expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, false); + expect(screen.queryByText("Not available to connected apps")).not.toBeInTheDocument(); + expect(screen.getByText("Connected (2)")).toBeInTheDocument(); + const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]); + expect(toolCountFetchedIds).toContain("s-unreach"); + }); + + const revocable = (reachable: boolean) => + [ + { server_id: "s-reach", server_name: "reachable_srv", auth_type: "none", connected_app_reachable: true }, + { server_id: "s-drop", server_name: "revoked_srv", auth_type: "none", connected_app_reachable: reachable }, + ] as MCPServer[]; + + const ConnectPanel = ({ + token, + onChange, + client, + }: { + token: string; + onChange: (servers: string[]) => void; + client: QueryClient; + }) => ( + + + + ); + + const newClient = () => new QueryClient({ defaultOptions: { queries: { retry: false } } }); + + it("drops an open detail view when a refetch removes that server from the reachable set", async () => { + vi.mocked(fetchMCPServers).mockResolvedValueOnce(revocable(true)).mockResolvedValueOnce(revocable(false)); + vi.mocked(listMCPTools).mockResolvedValue({ tools: [] }); + + const client = newClient(); + const { rerender } = render(); + + fireEvent.click(await screen.findByText("revoked_srv")); + expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument(); + + rerender(); + + await waitFor(() => expect(screen.queryByRole("heading", { name: "revoked_srv" })).not.toBeInTheDocument()); + expect(screen.queryByRole("button", { name: "Connect" })).not.toBeInTheDocument(); + expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument(); + expect(screen.getByText("reachable_srv")).toBeInTheDocument(); + }); + + it("does not select a server whose Connect finishes after a refetch removed it", async () => { + vi.mocked(fetchMCPServers).mockResolvedValueOnce(revocable(true)).mockResolvedValueOnce(revocable(false)); + vi.mocked(listMCPTools).mockResolvedValue({ tools: [] }); + + const onChange = vi.fn(); + const client = newClient(); + const { rerender } = render(); + + fireEvent.click(await screen.findByText("revoked_srv")); + expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument(); + + let finishConnect: (result: { tools: never[] }) => void = () => {}; + vi.mocked(listMCPTools).mockImplementationOnce(() => new Promise((resolve) => (finishConnect = resolve))); + fireEvent.click(screen.getByRole("button", { name: "Connect" })); + + rerender(); + await waitFor(() => expect(screen.queryByRole("heading", { name: "revoked_srv" })).not.toBeInTheDocument()); + + await act(async () => { + finishConnect({ tools: [] }); + }); + + expect(onChange).not.toHaveBeenCalled(); + expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument(); + expect(screen.getByText("Connected", { exact: false }).textContent).toBe("Connected"); + }); + + it("does not select a server when Connect resolves in the same tick the refetch drops it", async () => { + let finishRefetch: (servers: MCPServer[]) => void = () => {}; + vi.mocked(fetchMCPServers) + .mockResolvedValueOnce(revocable(true)) + .mockImplementationOnce(() => new Promise((resolve) => (finishRefetch = resolve))); + vi.mocked(listMCPTools).mockResolvedValue({ tools: [] }); + + const onChange = vi.fn(); + const client = newClient(); + const { rerender } = render(); + + fireEvent.click(await screen.findByText("revoked_srv")); + expect(await screen.findByRole("heading", { name: "revoked_srv" })).toBeInTheDocument(); + + let finishConnect: (result: { tools: never[] }) => void = () => {}; + vi.mocked(listMCPTools).mockImplementationOnce(() => new Promise((resolve) => (finishConnect = resolve))); + fireEvent.click(screen.getByRole("button", { name: "Connect" })); + + rerender(); + + await act(async () => { + finishRefetch(revocable(false)); + finishConnect({ tools: [] }); + }); + + expect(onChange).not.toHaveBeenCalled(); + expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument(); + }); + + it("does not let a superseded list load overwrite the current reachable set", async () => { + let finishStaleLoad: (servers: MCPServer[]) => void = () => {}; + vi.mocked(fetchMCPServers) + .mockImplementationOnce(() => new Promise((resolve) => (finishStaleLoad = resolve))) + .mockResolvedValueOnce(revocable(false)); + vi.mocked(listMCPTools).mockResolvedValue({ tools: [] }); + + const client = newClient(); + const { rerender } = render(); + rerender(); + + expect(await screen.findByText("reachable_srv")).toBeInTheDocument(); + + await act(async () => { + finishStaleLoad(revocable(true)); + }); + + expect(screen.queryByText("revoked_srv")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index 329e4ec99fa..7dbfc058a77 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -103,16 +103,17 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, const [query, setQuery] = useState(""); const [activeTab, setActiveTab] = useState("all"); const [togglingOn, setTogglingOn] = useState>(new Set()); - const [detailServer, setDetailServer] = useState(null); + const [detailServerId, setDetailServerId] = useState(null); const [toolCounts, setToolCounts] = useState>({}); const [loadingCounts, setLoadingCounts] = useState(false); const [oauthConnected, setOauthConnected] = useState>(new Set()); const [oauthChecking, setOauthChecking] = useState>(new Set()); const serversRef = useRef([]); - useEffect(() => { - serversRef.current = servers; - }, [servers]); + const commitServers = useCallback((next: MCPServer[]) => { + serversRef.current = next; + setServers(next); + }, []); const selectedServersRef = useRef(selectedServers); useEffect(() => { selectedServersRef.current = selectedServers; @@ -124,13 +125,30 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, const nameOf = (s: MCPServer) => s.server_name ?? s.alias ?? s.server_id; - const fetchLoadCancelledRef = useRef(false); + const detailServer = servers.find((s) => s.server_id === detailServerId); + + const connectUnavailabilityLabel = useCallback( + (s: MCPServer): string | null => { + if (!connectMode) return null; + if (isUnsupportedOnGatewayConnect(s.auth_type)) return "Not supported on this connection"; + return null; + }, + [connectMode], + ); + + const connectableNow = useCallback( + (serverId: string): MCPServer | undefined => { + const current = serversRef.current.find((s) => s.server_id === serverId); + return current !== undefined && connectUnavailabilityLabel(current) === null ? current : undefined; + }, + [connectUnavailabilityLabel], + ); const fetchToolCount = useCallback( - async (server: MCPServer) => { + async (server: MCPServer, isCurrentLoad: () => boolean) => { try { const toolsData = await listMCPTools(accessToken, server.server_id); - if (fetchLoadCancelledRef.current) return; + if (!isCurrentLoad()) return; const tools: MCPTool[] = Array.isArray(toolsData?.tools) ? toolsData.tools : []; setToolCounts((prev) => ({ ...prev, [nameOf(server)]: tools.length })); } catch { @@ -141,17 +159,17 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, ); const checkOauthCredential = useCallback( - async (server: MCPServer) => { + async (server: MCPServer, isCurrentLoad: () => boolean) => { try { const status = await getMCPOAuthUserCredentialStatus(accessToken, server.server_id); - if (fetchLoadCancelledRef.current) return; + if (!isCurrentLoad()) return; if (status.has_credential && !status.is_expired) { setOauthConnected((prev) => new Set(prev).add(server.server_id)); } } catch { // ignore } finally { - if (!fetchLoadCancelledRef.current) { + if (isCurrentLoad()) { setOauthChecking((prev) => { const next = new Set(prev); next.delete(server.server_id); @@ -164,70 +182,77 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, ); useEffect(() => { - fetchLoadCancelledRef.current = false; + let current = true; + const isCurrentLoad = () => current; - fetchMCPServers(accessToken) + fetchMCPServers(accessToken, undefined, connectMode) .then(async (serverData) => { - if (fetchLoadCancelledRef.current) return; + if (!isCurrentLoad()) return; const list: MCPServer[] = Array.isArray(serverData) ? serverData : serverData?.data ?? []; - const oauthServers = list.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2); - setServers(list); + const reachable = connectMode ? list.filter((s) => s.connected_app_reachable !== false) : list; + const oauthServers = reachable.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2); + commitServers(reachable); setOauthChecking(new Set(oauthServers.map((s) => s.server_id))); setLoading(false); - oauthServers.forEach((s) => checkOauthCredential(s)); + oauthServers.forEach((s) => checkOauthCredential(s, isCurrentLoad)); setLoadingCounts(true); - const chunks = Array.from({ length: Math.ceil(list.length / TOOLS_FETCH_CONCURRENCY) }, (_, i) => - list.slice(i * TOOLS_FETCH_CONCURRENCY, (i + 1) * TOOLS_FETCH_CONCURRENCY), + const chunks = Array.from({ length: Math.ceil(reachable.length / TOOLS_FETCH_CONCURRENCY) }, (_, i) => + reachable.slice(i * TOOLS_FETCH_CONCURRENCY, (i + 1) * TOOLS_FETCH_CONCURRENCY), ); for (const chunk of chunks) { - if (fetchLoadCancelledRef.current) return; - await Promise.allSettled(chunk.map((s) => fetchToolCount(s))); + if (!isCurrentLoad()) return; + await Promise.allSettled(chunk.map((s) => fetchToolCount(s, isCurrentLoad))); } - if (!fetchLoadCancelledRef.current) setLoadingCounts(false); + if (isCurrentLoad()) setLoadingCounts(false); }) .catch(() => { - if (!fetchLoadCancelledRef.current) { - setServers([]); + if (isCurrentLoad()) { + commitServers([]); setLoading(false); } }); return () => { - fetchLoadCancelledRef.current = true; + current = false; }; - }, [accessToken, fetchToolCount, checkOauthCredential]); + }, [accessToken, connectMode, commitServers, fetchToolCount, checkOauthCredential]); useEffect(() => { if (oauthConnected.size === 0) return; const namesToAdd = serversRef.current - .filter((s) => oauthConnected.has(s.server_id) && !selectedServersRef.current.includes(nameOf(s))) + .filter( + (s) => + oauthConnected.has(s.server_id) && + !selectedServersRef.current.includes(nameOf(s)) && + connectUnavailabilityLabel(s) === null, + ) .map(nameOf); if (namesToAdd.length > 0) { onChangeRef.current([...selectedServersRef.current, ...namesToAdd]); } - }, [oauthConnected]); + }, [oauthConnected, connectUnavailabilityLabel]); - const handleToggle = async (serverName: string, checked: boolean, serverId?: string) => { + const handleToggle = async (server: MCPServer, checked: boolean) => { + const serverName = nameOf(server); if (!checked) { onChange(selectedServers.filter((s) => s !== serverName)); - if (serverId) { - setOauthConnected((prev) => { - const next = new Set(prev); - next.delete(serverId); - return next; - }); - } + setOauthConnected((prev) => { + const next = new Set(prev); + next.delete(server.server_id); + return next; + }); return; } + if (connectableNow(server.server_id) === undefined) return; setTogglingOn((prev) => new Set(prev).add(serverName)); try { - const idToFetch = serverId ?? serverName; - const result = await listMCPTools(accessToken, idToFetch); + const result = await listMCPTools(accessToken, server.server_id); if (result?.error) { MessageManager.warning(`Could not load tools for ${serverName}`); return; } + if (connectableNow(server.server_id) === undefined) return; if (!selectedServersRef.current.includes(serverName)) { onChange([...selectedServersRef.current, serverName]); } @@ -243,11 +268,10 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, }; const renderConnectionIndicator = (server: MCPServer) => { - if (connectMode && isUnsupportedOnGatewayConnect(server.auth_type)) { + const unavailabilityLabel = connectUnavailabilityLabel(server); + if (unavailabilityLabel !== null) { return ( - - Not supported on this connection - + {unavailabilityLabel} ); } if (server.auth_type === AUTH_TYPE.OAUTH2) { @@ -285,11 +309,23 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, !query.trim() || name.toLowerCase().includes(query.toLowerCase()) || (s.description ?? "").toLowerCase().includes(query.toLowerCase()); - const matchesTab = activeTab === "all" || selectedServers.includes(name); + const matchesTab = + activeTab === "all" || (selectedServers.includes(name) && connectUnavailabilityLabel(s) === null); return matchesQuery && matchesTab; }); - const connectedCount = servers.filter((s) => selectedServers.includes(nameOf(s))).length; + const connectedCount = servers.filter( + (s) => selectedServers.includes(nameOf(s)) && connectUnavailabilityLabel(s) === null, + ).length; + + const emptyStateText = () => { + if (servers.length === 0) { + return connectMode + ? "No MCP servers are available to this connection yet. Ask an admin to grant your user or team access." + : "No MCP servers configured. Add servers in Tools -> MCP Servers."; + } + return activeTab === "connected" ? "No servers connected yet." : "No servers match your search."; + }; const totalTools = Object.values(toolCounts).reduce((sum, n) => sum + n, 0); if (detailServer) { @@ -298,12 +334,65 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, const isTogglingOn = togglingOn.has(name); const color = getAvatarColor(name); + const renderDetailAction = () => { + const unavailabilityLabel = connectUnavailabilityLabel(detailServer); + if (unavailabilityLabel !== null) { + return {unavailabilityLabel}; + } + if (detailServer.auth_type !== AUTH_TYPE.OAUTH2) { + return ( + + ); + } + if (oauthConnected.has(detailServer.server_id)) { + return ( + + ); + } + return ( + { + setOauthConnected((prev) => new Set(prev).add(id)); + }} + variant="button" + /> + ); + }; + return (
- {detailServer.auth_type === AUTH_TYPE.OAUTH2 ? ( - oauthConnected.has(detailServer.server_id) ? ( - - ) : ( - { - setOauthConnected((prev) => new Set(prev).add(id)); - }} - variant="button" - /> - ) - ) : ( - - )} + {renderDetailAction()}

Information

@@ -494,13 +542,7 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, ))} ) : filtered.length === 0 ? ( -
- {servers.length === 0 - ? "No MCP servers configured. Add servers in Tools -> MCP Servers." - : activeTab === "connected" - ? "No servers connected yet." - : "No servers match your search."} -
+
{emptyStateText()}
) : (
{filtered.map((server, idx) => { @@ -508,16 +550,16 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, const color = getAvatarColor(name); const isLeftCol = idx % 2 === 0; const count = toolCounts[name]; - const unsupported = !!connectMode && isUnsupportedOnGatewayConnect(server.auth_type); + const unavailable = connectUnavailabilityLabel(server) !== null; return (
setDetailServer(server)} + onClick={() => setDetailServerId(server.server_id)} className={`flex items-center gap-3 p-4 bg-card cursor-pointer transition-colors hover:bg-accent/30 min-w-0 ${ isLeftCol ? "border-r" : "" } ${Math.floor(idx / 2) < Math.floor((filtered.length - 1) / 2) ? "border-b" : ""} ${ - unsupported ? "opacity-50" : "" + unavailable ? "opacity-50" : "" }`} > {server.mcp_info?.logo_url ? ( diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index de497d91afe..44ac25ea955 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -436,6 +436,7 @@ export interface MCPServer { byok_description?: string[] | null; byok_api_key_help_url?: string | null; has_user_credential?: boolean | null; + connected_app_reachable?: boolean | null; /** GitHub / source repository URL */ source_url?: string | null; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 1a0951db967..03cf0e9583c 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -4751,9 +4751,12 @@ export const fetchDiscoverableMCPServers = async (accessToken: string) => { } }; -export const fetchMCPServers = async (accessToken: string, teamId?: string | null) => { +export const fetchMCPServers = async (accessToken: string, teamId?: string | null, connectedAppView?: boolean) => { try { - return await apiClient.get(`/v1/mcp/server`, { accessToken, query: { team_id: teamId || undefined } }); + return await apiClient.get(`/v1/mcp/server`, { + accessToken, + query: { team_id: teamId || undefined, connected_app_view: connectedAppView || undefined }, + }); } catch (error) { console.error("Failed to fetch MCP servers:", error); throw error; From 3c2264cfacc3081492e41c8cb9bff5d7da2a7c4a Mon Sep 17 00:00:00 2001 From: tin-berri Date: Thu, 30 Jul 2026 23:21:23 -0700 Subject: [PATCH 26/26] feat(ui): expose classifier context window fields on Auto-Router screens (LIT-5036) (#35315) PR #35185 added classifier_context_window_size and classifier_context_per_turn_chars to ComplexityRouterConfig; they worked via config.yaml and the API but had no UI control on the Add Model or Edit Auto-Router screens. Wires the two fields into both, shown only when the LLM classifier is selected. --- .../add_model/ClassificationMethodConfig.tsx | 65 +++++++++++++++- .../add_model/ComplexityRouterConfig.test.tsx | 77 ++++++++++++++++++- .../add_model/ComplexityRouterConfig.tsx | 4 + .../add_model/add_auto_router_tab.tsx | 4 + .../build_complexity_router_config.test.ts | 48 ++++++++++++ .../build_complexity_router_config.ts | 14 ++++ ...d_updated_complexity_router_config.test.ts | 65 ++++++++++++++++ .../edit_auto_router_modal.test.tsx | 67 +++++++++++++++- .../edit_auto_router_modal.tsx | 23 +++++- 9 files changed, 359 insertions(+), 8 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 92df8029edc..75acf87a7d8 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -1,7 +1,13 @@ import { InfoCircleOutlined } from "@ant-design/icons"; import { Select as AntdSelect, Card, InputNumber, Radio, Space, Tooltip, Typography } from "antd"; import React from "react"; -import { ClassifierType, ComplexityRouterConfigValue, DEFAULT_CLASSIFIER_TIMEOUT_MS } from "./ComplexityRouterConfig"; +import { + ClassifierType, + ComplexityRouterConfigValue, + DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS, + DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, + DEFAULT_CLASSIFIER_TIMEOUT_MS, +} from "./ComplexityRouterConfig"; const { Text } = Typography; @@ -26,14 +32,23 @@ const ClassificationMethodConfig: React.FC = ({ showValidationErrors && value.classifier_type === "llm" && !value.classifier_llm_config?.model; const handleClassifierTypeChange = (classifierType: ClassifierType) => { - onChange({ + const nextValue: ComplexityRouterConfigValue = { ...value, classifier_type: classifierType, classifier_llm_config: classifierType === "llm" ? value.classifier_llm_config ?? { model: "", timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS } : undefined, - }); + classifier_context_window_size: + classifierType === "llm" + ? value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE + : undefined, + classifier_context_per_turn_chars: + classifierType === "llm" + ? value.classifier_context_per_turn_chars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS + : undefined, + }; + onChange(nextValue); }; const handleClassifierModelChange = (model: string) => { @@ -56,6 +71,20 @@ const ClassificationMethodConfig: React.FC = ({ }); }; + const handleClassifierContextWindowSizeChange = (windowSize: number | null) => { + onChange({ + ...value, + classifier_context_window_size: windowSize ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, + }); + }; + + const handleClassifierContextPerTurnCharsChange = (perTurnChars: number | null) => { + onChange({ + ...value, + classifier_context_per_turn_chars: perTurnChars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS, + }); + }; + return ( <> = ({ response.
+
+ + Context Window Size + + + + Number of prior user turns (tool output and harness reminders excluded) sent to the classifier as context, + so a referring follow-up like "now do the same for the streaming path" is classified against + what it refers to. Set to 0 to send only the current message. + +
+
+ + Context Per-Turn Character Limit + + + + Prior turns longer than this are truncated. + +
)} diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 6b3c1961468..a1e992bd46b 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -99,11 +99,14 @@ describe("ComplexityRouterConfig", () => { fireEvent.click(screen.getByText("Advanced: Classification Method")); fireEvent.click(screen.getByText("LLM Classifier")); - expect(onChange).toHaveBeenCalledWith({ + const expectedValue: ComplexityRouterConfigValue = { ...defaultValue, classifier_type: "llm", classifier_llm_config: { model: "", timeout_ms: 3000 }, - }); + classifier_context_window_size: 3, + classifier_context_per_turn_chars: 200, + }; + expect(onChange).toHaveBeenCalledWith(expectedValue); }); it("should show classifier fields and use the configured values when classifier_type is llm", () => { @@ -111,6 +114,8 @@ describe("ComplexityRouterConfig", () => { ...defaultValue, classifier_type: "llm", classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 750 }, + classifier_context_window_size: 5, + classifier_context_per_turn_chars: 400, }; renderWithProviders(); @@ -119,6 +124,74 @@ describe("ComplexityRouterConfig", () => { expect(screen.getByText("Classifier Model")).toBeInTheDocument(); expect(screen.getByText("Timeout (ms)")).toBeInTheDocument(); expect(screen.getByDisplayValue("750")).toBeInTheDocument(); + expect(screen.getByText("Context Window Size")).toBeInTheDocument(); + expect(screen.getByDisplayValue("5")).toBeInTheDocument(); + expect(screen.getByText("Context Per-Turn Character Limit")).toBeInTheDocument(); + expect(screen.getByDisplayValue("400")).toBeInTheDocument(); + }); + + it("should default classifier context fields to 3 and 200 when llm is selected without explicit values", () => { + const llmValue: ComplexityRouterConfigValue = { + ...defaultValue, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 }, + }; + renderWithProviders(); + + fireEvent.click(screen.getByText("Advanced: Classification Method")); + + const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement; + expect(within(windowSizeSection).getByDisplayValue("3")).toBeInTheDocument(); + + const perTurnCharsSection = screen.getByText("Context Per-Turn Character Limit").closest("div") as HTMLElement; + expect(within(perTurnCharsSection).getByDisplayValue("200")).toBeInTheDocument(); + }); + + it("should hide classifier context fields when classifier_type is heuristic", () => { + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.queryByText("Context Window Size")).not.toBeInTheDocument(); + expect(screen.queryByText("Context Per-Turn Character Limit")).not.toBeInTheDocument(); + }); + + it("should call onChange with the updated classifier_context_window_size when edited", () => { + const onChange = vi.fn(); + const llmValue: ComplexityRouterConfigValue = { + ...defaultValue, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 }, + }; + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + + const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement; + const input = within(windowSizeSection).getByRole("spinbutton"); + fireEvent.change(input, { target: { value: "7" } }); + + expect(onChange).toHaveBeenCalledWith({ + ...llmValue, + classifier_context_window_size: 7, + }); + }); + + it("should call onChange with the updated classifier_context_per_turn_chars when edited", () => { + const onChange = vi.fn(); + const llmValue: ComplexityRouterConfigValue = { + ...defaultValue, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 }, + }; + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + + const perTurnCharsSection = screen.getByText("Context Per-Turn Character Limit").closest("div") as HTMLElement; + const input = within(perTurnCharsSection).getByRole("spinbutton"); + fireEvent.change(input, { target: { value: "500" } }); + + expect(onChange).toHaveBeenCalledWith({ + ...llmValue, + classifier_context_per_turn_chars: 500, + }); }); it("should render the custom technical keywords field", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 1f2edf697a9..de32d15d5a7 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -12,6 +12,8 @@ const { Text } = Typography; export const DEFAULT_CLASSIFIER_TIMEOUT_MS = 3000; export const DEFAULT_TIER_DISTANCE_PENALTY = 0.5; +export const DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE = 3; +export const DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS = 200; export interface ComplexityTiers { SIMPLE: string[]; @@ -40,6 +42,8 @@ export interface ComplexityRouterConfigValue { tiers: ComplexityTiers; classifier_type: ClassifierType; classifier_llm_config?: ClassifierLLMConfig; + classifier_context_window_size?: number; + classifier_context_per_turn_chars?: number; adaptive?: boolean; adaptive_weights?: AdaptiveRouterWeights; tier_distance_penalty?: number; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index cd3294b347f..fb7eabd7110 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -99,6 +99,8 @@ const AddAutoRouterTab: React.FC = ({ tiers, classifier_type: classifierType, classifier_llm_config: classifierLlmConfig, + classifier_context_window_size: classifierContextWindowSize, + classifier_context_per_turn_chars: classifierContextPerTurnChars, adaptive = false, adaptive_weights: adaptiveWeights = DEFAULT_ADAPTIVE_WEIGHTS, tier_distance_penalty: tierDistancePenalty = DEFAULT_TIER_DISTANCE_PENALTY, @@ -142,6 +144,8 @@ const AddAutoRouterTab: React.FC = ({ tiers, classifierType, classifierLlmConfig, + classifierContextWindowSize, + classifierContextPerTurnChars, customTechnicalKeywords, keywordTierRules, semanticMatchingEnabled, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index b5973bf7101..e269a3c9028 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -16,6 +16,8 @@ const baseParams: BuildComplexityRouterConfigParams = { tiers, classifierType: "heuristic", classifierLlmConfig: undefined, + classifierContextWindowSize: undefined, + classifierContextPerTurnChars: undefined, customTechnicalKeywords: [], keywordTierRules: [], semanticMatchingEnabled: false, @@ -75,6 +77,52 @@ describe("buildComplexityRouterConfig", () => { expect(config.classifier_llm_config).toBeUndefined(); }); + it("includes classifier_context_window_size and classifier_context_per_turn_chars only when classifier_type is llm", () => { + const params: BuildComplexityRouterConfigParams = { + ...baseParams, + classifierType: "llm", + classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 }, + classifierContextWindowSize: 5, + classifierContextPerTurnChars: 300, + }; + const config = buildComplexityRouterConfig(params); + expect(config.classifier_context_window_size).toBe(5); + expect(config.classifier_context_per_turn_chars).toBe(300); + }); + + it("omits classifier_context_window_size and classifier_context_per_turn_chars when classifier_type is heuristic even if values linger in state", () => { + const params: BuildComplexityRouterConfigParams = { + ...baseParams, + classifierType: "heuristic", + classifierContextWindowSize: 5, + classifierContextPerTurnChars: 300, + }; + const config = buildComplexityRouterConfig(params); + expect(config.classifier_context_window_size).toBeUndefined(); + expect(config.classifier_context_per_turn_chars).toBeUndefined(); + }); + + it("omits classifier_context_window_size and classifier_context_per_turn_chars when classifier_type is llm but neither was set, leaving the backend default", () => { + const config = buildComplexityRouterConfig({ + ...baseParams, + classifierType: "llm", + classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 }, + }); + expect(config.classifier_context_window_size).toBeUndefined(); + expect(config.classifier_context_per_turn_chars).toBeUndefined(); + }); + + it("allows classifier_context_window_size of 0, distinct from unset, to send no prior-turn context", () => { + const params: BuildComplexityRouterConfigParams = { + ...baseParams, + classifierType: "llm", + classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 }, + classifierContextWindowSize: 0, + }; + const config = buildComplexityRouterConfig(params); + expect(config.classifier_context_window_size).toBe(0); + }); + it("sends keyword_tier_rules with their per-tier targeting preserved (not flattened)", () => { const params: BuildComplexityRouterConfigParams = { ...baseParams, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 02be9b280aa..04a9b0f9bd4 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -12,6 +12,8 @@ export interface BuildComplexityRouterConfigParams { tiers: ComplexityTiers; classifierType: ClassifierType; classifierLlmConfig: ClassifierLLMConfig | undefined; + classifierContextWindowSize: number | undefined; + classifierContextPerTurnChars: number | undefined; customTechnicalKeywords: string[]; keywordTierRules: KeywordTierRule[]; semanticMatchingEnabled: boolean; @@ -29,6 +31,8 @@ export interface ComplexityRouterConfigPayload { tiers: ComplexityTiers; classifier_type: ClassifierType; classifier_llm_config?: ClassifierLLMConfig; + classifier_context_window_size?: number; + classifier_context_per_turn_chars?: number; custom_technical_keywords?: string[]; keyword_tier_rules?: { keywords: string[]; tier: KeywordTierRule["tier"] }[]; semantic_keyword_matching?: boolean; @@ -69,6 +73,8 @@ export const buildComplexityRouterConfig = ({ tiers, classifierType, classifierLlmConfig, + classifierContextWindowSize, + classifierContextPerTurnChars, customTechnicalKeywords, keywordTierRules, semanticMatchingEnabled, @@ -89,6 +95,14 @@ export const buildComplexityRouterConfig = ({ tiers, classifier_type: classifierType, ...(classifierType === "llm" && classifierLlmConfig && { classifier_llm_config: classifierLlmConfig }), + ...(classifierType === "llm" && + classifierContextWindowSize !== undefined && { + classifier_context_window_size: classifierContextWindowSize, + }), + ...(classifierType === "llm" && + classifierContextPerTurnChars !== undefined && { + classifier_context_per_turn_chars: classifierContextPerTurnChars, + }), ...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }), ...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }), escalation_keywords: cleanedEscalationKeywords, diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index e39a7b6e444..eef5e1e4d06 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -86,3 +86,68 @@ describe("buildUpdatedComplexityRouterConfig keyword matching", () => { expect(result.match_threshold).toBe(0.72); }); }); + +const STORED_LLM = { + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] }, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-4o-mini", timeout_ms: 3000 }, + classifier_context_window_size: 5, + classifier_context_per_turn_chars: 300, +}; + +describe("buildUpdatedComplexityRouterConfig classifier context window", () => { + it("round-trips an untouched edit without changing the classifier context values", () => { + const formValue = { + tiers: STORED_LLM.tiers, + classifier_type: "llm" as const, + classifier_llm_config: STORED_LLM.classifier_llm_config, + classifier_context_window_size: 5, + classifier_context_per_turn_chars: 300, + }; + const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue); + + expect(result.classifier_context_window_size).toBe(5); + expect(result.classifier_context_per_turn_chars).toBe(300); + }); + + it("persists an edited classifier context window size and per-turn char limit", () => { + const formValue = { + tiers: STORED_LLM.tiers, + classifier_type: "llm" as const, + classifier_llm_config: STORED_LLM.classifier_llm_config, + classifier_context_window_size: 10, + classifier_context_per_turn_chars: 500, + }; + const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue); + + expect(result.classifier_context_window_size).toBe(10); + expect(result.classifier_context_per_turn_chars).toBe(500); + }); + + it("omits classifier context fields when classifier_type is heuristic even if values linger in state", () => { + const formValue = { + tiers: STORED_LLM.tiers, + classifier_type: "heuristic" as const, + classifier_context_window_size: 5, + classifier_context_per_turn_chars: 300, + }; + const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue); + + expect(result.classifier_context_window_size).toBeUndefined(); + expect(result.classifier_context_per_turn_chars).toBeUndefined(); + }); + + it("does not resurrect a stale stored classifier_context_window_size once the form's own value is unset", () => { + // classifier_context_window_size is a MANAGED key: the form's value must win over whatever + // is still sitting in the stored config, never fall back to it through preservedConfig. + const formValue = { + tiers: STORED_LLM.tiers, + classifier_type: "llm" as const, + classifier_llm_config: STORED_LLM.classifier_llm_config, + }; + const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue); + + expect(result.classifier_context_window_size).toBeUndefined(); + expect(result.classifier_context_per_turn_chars).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index e47bbecf8bd..504234e8977 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -1,7 +1,7 @@ import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; -import { renderWithProviders, screen, waitFor } from "@/../tests/test-utils"; +import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../tests/test-utils"; import NotificationsManager from "@/components/molecules/notifications_manager"; import EditAutoRouterModal from "./edit_auto_router_modal"; @@ -119,3 +119,68 @@ describe("EditAutoRouterModal keyword matching", () => { expect(modelPatchUpdateCall).not.toHaveBeenCalled(); }); }); + +describe("EditAutoRouterModal classifier context window", () => { + beforeEach(() => { + modelPatchUpdateCall.mockClear(); + }); + + const STORED_LLM_CONFIG = { + tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: ["gpt-4o-mini"], COMPLEX: ["gpt-4o-mini"], REASONING: ["gpt-4o-mini"] }, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-4o-mini", timeout_ms: 3000 }, + classifier_context_window_size: 5, + classifier_context_per_turn_chars: 300, + }; + + const renderLlmModal = () => + renderWithProviders( + , + ); + + // Hydration bugs are invisible to the payload-builder unit tests, which only exercise + // buildUpdatedComplexityRouterConfig with a form value the caller already assembled by hand. + // Only driving the real component through open, then save with nothing touched, catches a + // missing initializeForm hydration line. + it("shows the stored classifier context values and preserves them through an untouched open-and-save", async () => { + const user = userEvent.setup(); + renderLlmModal(); + + await user.click(await screen.findByText("Advanced: Classification Method")); + await screen.findByText("Context Window Size"); + expect(screen.getByDisplayValue("5")).toBeInTheDocument(); + expect(screen.getByDisplayValue("300")).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + const config = savedConfig(); + expect(config.classifier_context_window_size).toBe(5); + expect(config.classifier_context_per_turn_chars).toBe(300); + }); + + it("persists an edited classifier context window size", async () => { + const user = userEvent.setup(); + renderLlmModal(); + + await user.click(await screen.findByText("Advanced: Classification Method")); + const windowSizeSection = (await screen.findByText("Context Window Size")).closest("div") as HTMLElement; + const input = within(windowSizeSection).getByRole("spinbutton"); + fireEvent.change(input, { target: { value: "8" } }); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().classifier_context_window_size).toBe(8); + }); +}); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 8fbca822165..2686f99307e 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -33,6 +33,8 @@ const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "tiers", "classifier_type", "classifier_llm_config", + "classifier_context_window_size", + "classifier_context_per_turn_chars", "adaptive", "adaptive_weights", "tier_distance_penalty", @@ -86,6 +88,14 @@ export const buildUpdatedComplexityRouterConfig = ( tiers: value.tiers, classifier_type: value.classifier_type, ...(value.classifier_type === "llm" ? { classifier_llm_config: value.classifier_llm_config } : {}), + ...(value.classifier_type === "llm" && + value.classifier_context_window_size !== undefined && { + classifier_context_window_size: value.classifier_context_window_size, + }), + ...(value.classifier_type === "llm" && + value.classifier_context_per_turn_chars !== undefined && { + classifier_context_per_turn_chars: value.classifier_context_per_turn_chars, + }), ...(customTechnicalKeywords && customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords, @@ -182,7 +192,7 @@ const EditAutoRouterModal: React.FC = ({ parsedConfig = JSON.parse(parsedConfig); } - setComplexityRouterConfig({ + const hydratedComplexityRouterConfig: ComplexityRouterConfigValue = { tiers: { SIMPLE: normalizeTierModels(parsedConfig.tiers?.SIMPLE), MEDIUM: normalizeTierModels(parsedConfig.tiers?.MEDIUM), @@ -191,12 +201,21 @@ const EditAutoRouterModal: React.FC = ({ }, classifier_type: parsedConfig.classifier_type || "heuristic", classifier_llm_config: parsedConfig.classifier_llm_config, + classifier_context_window_size: + typeof parsedConfig.classifier_context_window_size === "number" + ? parsedConfig.classifier_context_window_size + : undefined, + classifier_context_per_turn_chars: + typeof parsedConfig.classifier_context_per_turn_chars === "number" + ? parsedConfig.classifier_context_per_turn_chars + : undefined, adaptive: parsedConfig.adaptive || false, adaptive_weights: parsedConfig.adaptive_weights, tier_distance_penalty: parsedConfig.tier_distance_penalty, adaptive_eligible: parsedConfig.adaptive_eligible || "all", return_raw_model_name: parsedConfig.return_raw_model_name || false, - }); + }; + setComplexityRouterConfig(hydratedComplexityRouterConfig); setCustomTechnicalKeywords( Array.isArray(parsedConfig.custom_technical_keywords) ? parsedConfig.custom_technical_keywords : [], );