From 229e13678320cc18fc496ed4bfe7f6af601f6477 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 11:26:54 -0700 Subject: [PATCH 01/21] fix(mcp): honor admin-entered OAuth URLs on authorize after issuer yield Co-authored-by: Cursor --- .../mcp_server/discoverable_endpoints.py | 55 +++++++++----- .../mcp_server/mcp_server_manager.py | 6 +- .../types/mcp_server/mcp_server_manager.py | 12 ++++ .../mcp_server/test_discoverable_endpoints.py | 72 +++++++++++++++++++ 4 files changed, 126 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index aef4f5dc721..6ede1c553e8 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -663,6 +663,23 @@ def _endpoint_not_configured_detail( ) +async def _server_with_oauth_endpoints(mcp_server: MCPServer) -> MCPServer: + """Join deferred OAuth discovery only when this server still has no authorize URL. + + Admin-entered endpoints live on ``configured_*`` after an anchored issuer empties the + resolved fields. Those already let authorize/token run, so discovery is not awaited + and cannot 503 over a leftover pin. A server with nothing configured still joins the + deferred task; no slot is a no-op and the caller 400s. + """ + if mcp_server.effective_authorization_url is not None: + return mcp_server + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load + global_mcp_server_manager, + ) + + return await global_mcp_server_manager.ensure_oauth_metadata_discovered(mcp_server) + + def _raise_unless_oauth2_discovery_server( mcp_server: MCPServer | None, mcp_server_name: str | None, @@ -697,7 +714,7 @@ def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool: returns directly to the client's redirect URI without transiting the gateway. Gateway-side redirect trust and the ``/callback`` state relay therefore only apply to the short-circuit arm, where the upstream only knows the gateway's own callback.""" - return mcp_server.is_dcr_bridge and bool(mcp_server.registration_url) and not mcp_server.client_id + return mcp_server.is_dcr_bridge and bool(mcp_server.effective_registration_url) and not mcp_server.client_id def _require_s256_pkce( @@ -745,7 +762,7 @@ def _redirect_to_upstream_authorize( **({"scope": scope_value} if scope_value else {}), **({"resource": upstream_resource} if upstream_resource else {}), } - parsed_auth_url: Final = urlparse(mcp_server.authorization_url or "") + parsed_auth_url: Final = urlparse(mcp_server.effective_authorization_url or "") merged_params: Final = {**dict(parse_qsl(parsed_auth_url.query)), **passthrough_params} return RedirectResponse(urlunparse(parsed_auth_url._replace(query=urlencode(merged_params)))) @@ -812,11 +829,12 @@ async def authorize_with_server( ephemeral_dcr_client: "EphemeralDcrClient | None" = None, ): _raise_if_not_oauth2(mcp_server) - if mcp_server.authorization_url is None: + resolved_server: Final = await _server_with_oauth_endpoints(mcp_server) + if resolved_server.effective_authorization_url is None: raise HTTPException( status_code=400, detail=_endpoint_not_configured_detail( - mcp_server, + resolved_server, "authorization url", "set Authorization URL and Token URL manually", "set Issuer to discover them from the identity provider (RFC 8414)", @@ -913,7 +931,7 @@ async def authorize_with_server( if upstream_resource: params["resource"] = upstream_resource - parsed_auth_url: Final = urlparse(mcp_server.authorization_url) + parsed_auth_url: Final = urlparse(resolved_server.effective_authorization_url) existing_params: Final = dict(parse_qsl(parsed_auth_url.query)) existing_params.update(params) final_url: Final = urlunparse(parsed_auth_url._replace(query=urlencode(existing_params))) @@ -946,11 +964,13 @@ async def exchange_token_with_server( if grant_type not in ("authorization_code", "refresh_token"): raise HTTPException(status_code=400, detail="Unsupported grant_type") - if mcp_server.token_url is None: + resolved_server: Final = await _server_with_oauth_endpoints(mcp_server) + token_url: Final = resolved_server.effective_token_url + if token_url is None: raise HTTPException( status_code=400, detail=_endpoint_not_configured_detail( - mcp_server, + resolved_server, "token url", "set Token URL manually", "set Issuer to discover it from the identity provider (RFC 8414)", @@ -1067,7 +1087,7 @@ async def exchange_token_with_server( async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) try: response: Final = await async_client.post( - mcp_server.token_url, + token_url, headers={"Accept": "application/json", **token_request.headers}, data=token_data, ) @@ -1551,7 +1571,8 @@ async def mint_ephemeral_dcr_client(request: Request, mcp_server: MCPServer) -> bounded by the server count even when the request origin varies) so parallel authorize requests cannot each register an upstream client; the cache stamps nothing onto the server record and correctness never depends on it because the sealed state carries the client through the flow.""" - if mcp_server.registration_url is None: + registration_url: Final = mcp_server.effective_registration_url + if registration_url is None: return None request_base_url: Final = get_request_base_url(request) cache_key: Final = f"mcp_ephemeral_dcr_client:{mcp_server.server_id}:{request_base_url}" @@ -1571,7 +1592,7 @@ async def mint_ephemeral_dcr_client(request: Request, mcp_server: MCPServer) -> "token_endpoint_auth_method": "none", } response: Final = await _post_dcr_registration( - registration_url=mcp_server.registration_url, + registration_url=registration_url, register_data=register_data, server_id=mcp_server.server_id, ) @@ -1617,7 +1638,7 @@ async def resolve_ephemeral_dcr_client( usable to generate orphan IdP clients).""" if not (mcp_server.is_true_passthrough or (mcp_server.is_oauth_delegate and not mcp_server.is_dcr_bridge)): return None - if mcp_server.authorization_url is None: + if mcp_server.effective_authorization_url is None: raise HTTPException( status_code=400, detail="MCP server authorization url is not set", @@ -1661,21 +1682,23 @@ async def register_client_with_server( ): return dummy_return - if mcp_server.authorization_url is None: + resolved_server: Final = await _server_with_oauth_endpoints(mcp_server) + if resolved_server.effective_authorization_url is None: raise HTTPException( status_code=400, detail=_endpoint_not_configured_detail( - mcp_server, + resolved_server, "authorization url", "set Authorization URL and Token URL manually", "set Issuer to discover them from the identity provider (RFC 8414)", ), ) - if mcp_server.registration_url is None: + registration_url: Final = resolved_server.effective_registration_url + if registration_url is None: return dummy_return - bridge_relay: Final = _dcr_bridge_relays_client_registration(mcp_server) + bridge_relay: Final = _dcr_bridge_relays_client_registration(resolved_server) if bridge_relay and not client_redirect_uris: raise HTTPException( status_code=400, @@ -1690,7 +1713,7 @@ async def register_client_with_server( "token_endpoint_auth_method": token_endpoint_auth_method or ("none" if bridge_relay else ""), } response: Final = await _post_dcr_registration( - registration_url=mcp_server.registration_url, + registration_url=registration_url, register_data=register_data, server_id=mcp_server.server_id, ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7ab26db0f3e..fe0c73a5efc 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -523,7 +523,7 @@ def _oauth_endpoints_unresolved(server: MCPServer) -> bool: # can come from resource discovery, so a server that resolved its endpoints but no scopes is # still unresolved for its flow. return True - if server.is_dcr_bridge and not server.client_id and server.registration_url is None: + if server.is_dcr_bridge and not server.client_id and server.effective_registration_url is None: # A DCR bridge with no admin-configured client can only register callers through the # upstream's registration endpoint, so a build that resolved the authorize and token # endpoints but not registration_endpoint (partial metadata) is still unresolved for its @@ -535,8 +535,8 @@ def _oauth_endpoints_unresolved(server: MCPServer) -> bool: return _flow_endpoints_missing( server.auth_type, MCPServerManager.effective_oauth2_flow(server), - server.authorization_url, - server.token_url, + server.effective_authorization_url, + server.effective_token_url, server.token_exchange_endpoint, ) diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index d09503cdc4d..401793a79e4 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -183,6 +183,18 @@ class MCPServer(BaseModel): def __str__(self) -> str: return self.__repr__() + @property + def effective_authorization_url(self) -> str | None: + return self.authorization_url or self.configured_authorization_url + + @property + def effective_token_url(self) -> str | None: + return self.token_url or self.configured_token_url + + @property + def effective_registration_url(self) -> str | None: + return self.registration_url or self.configured_registration_url + @property def has_client_credentials(self) -> bool: """True if this server should use the OAuth2 client_credentials (M2M) flow. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index bcac27a4a14..790b5c7ad22 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -8854,6 +8854,78 @@ async def test_authorize_wall_names_the_issuer_for_anchored_servers(): assert "idp.example.com" not in detail_text +@pytest.mark.asyncio +async def test_authorize_uses_admin_entered_github_oauth_urls_after_issuer_yield(): + """GitHub MCP servers store Authorization URL and Token URL on the row. 1.99 can empty + the resolved authorization_url when a leftover issuer is treated as a pin (RFC 8414 + yield). The UI authorize must still redirect to the admin-entered GitHub authorize URL + instead of 400ing that discovery against api.githubcopilot.com failed.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize_with_server, + ) + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="ecac50c4-8eca-438a-af80-9bdebadafc69", + name="github_mcp", + alias="github_mcp", + server_name="github_mcp", + url="https://api.githubcopilot.com/mcp/", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + client_id="github-app-client", + authorization_url=None, + token_url=None, + issuer="https://github.com", + issuer_is_anchored=True, + configured_authorization_url="https://github.com/login/oauth/authorize", + configured_token_url="https://github.com/login/oauth/access_token", + ) + mock_request = MagicMock() + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper") as mock_encrypt: + mock_encrypt.return_value = "mocked_encrypted_state" + response = await authorize_with_server( + request=mock_request, + mcp_server=server, + client_id="github-app-client", + redirect_uri="http://127.0.0.1:60108/callback", + state="state123", + ) + + assert response.status_code == 307 + assert "https://github.com/login/oauth/authorize" in response.headers["location"] + assert "client_id=github-app-client" in response.headers["location"] + + +def test_oauth_endpoints_count_admin_entered_urls_as_resolved(): + """A leftover issuer empties the resolved authorize/token fields but must not keep the + server on the deferred-discovery retry path when the admin already stored those URLs.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _oauth_endpoints_unresolved, + ) + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="github-configured", + name="github_mcp", + url="https://api.githubcopilot.com/mcp/", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + authorization_url=None, + token_url=None, + configured_authorization_url="https://github.com/login/oauth/authorize", + configured_token_url="https://github.com/login/oauth/access_token", + ) + assert _oauth_endpoints_unresolved(server) is False + + def test_passthrough_authorization_code_round_trips_and_rejects_hostile_input(): """The passthrough gateway code seals and recovers the ephemeral DCR client and upstream code, and is total over hostile input: a raw upstream code opens to None, and a tampered or From 91e7eb115d8efb2c342d29058134cf6fe6594f0e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 26 Aug 2026 11:46:54 -0700 Subject: [PATCH 02/21] fix(proxy): sync search tools into the router on management writes Creating a search tool through the UI only wrote the row; the router was updated solely by the add_deployment job, so the tool was unusable for up to PROXY_CONFIG_RELOAD_INTERVAL_SECONDS (30s by default) even on the worker that served the write. Tools declared in config.yaml load straight into the router at startup, which is why they never showed the delay. The create, update and delete endpoints now refresh the router inline, matching what the MCP server endpoints already do. The refresh is best-effort: the row is already committed, so a failure must not surface as a 500 and push the caller into a retry that creates duplicates. Two related gaps go with it. _init_search_tools_in_db skipped the router update whenever the merged list came back empty, so deleting the last search tool left it live in memory forever. And in store_model_in_db-off deployments the add_deployment job is never scheduled, so DB-backed search tools never reached the router at all; that branch now loads them at startup and keeps them fresh on its own interval, the same way MCP servers already do. --- litellm/proxy/proxy_server.py | 28 +++- .../search_tool_management.py | 28 +++- .../test_search_tool_management.py | 156 ++++++++++++++++++ .../proxy/proxy_server/test_proxy_config.py | 55 +++++- 4 files changed, 253 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9b5d1b9bfea..c596042bb71 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7470,11 +7470,9 @@ class ProxyConfig: len(db_search_tools), ) - if llm_router is not None and search_tools: + if llm_router is not None: await SearchAPIRouter.update_router_search_tools(router_instance=llm_router, search_tools=search_tools) verbose_proxy_logger.info("Successfully loaded %s search tool(s) into router", len(search_tools)) - elif llm_router is not None: - verbose_proxy_logger.debug("No search tools found in config or database, skipping router update") else: verbose_proxy_logger.debug( "Router not initialized yet, search tools will be added when router is created" @@ -7485,6 +7483,19 @@ class ProxyConfig: "litellm.proxy.proxy_server.py::ProxyConfig:_init_search_tools_in_db - %s", e ) + async def reload_search_tools_from_db(self) -> None: + """Refresh this worker's router from the search tools table. + + Driven by the management endpoints so the worker that served the write is correct + immediately, and by the periodic job in store_model_in_db-off deployments. Gated the same + way as startup, so an admin who excluded search_tools from supported_db_objects opts out. + """ + if not self._should_load_db_object(object_type="search_tools"): + return + if prisma_client is None: + return + await self._init_search_tools_in_db(prisma_client=prisma_client) + @staticmethod def _merge_config_and_db_search_tools( config_search_tools: list[SearchToolTypedDict], @@ -9104,7 +9115,18 @@ class ProxyStartupEvent: if store_model_in_db is not True: await proxy_config.init_mcp_servers_from_db() + # Without this branch's own refresh, a UI-created search tool never reaches the router: + # the add_deployment job that carries it in store_model_in_db=True mode is not scheduled. + await proxy_config.reload_search_tools_from_db() if prisma_client is not None: + scheduler.add_job( + proxy_config.reload_search_tools_from_db, + "interval", + seconds=config_reload_interval_seconds, + id="reload_search_tools_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) # DB-backed MCP servers are live objects in every mode, so the registry refresh that # store_model_in_db=True deployments get via the add_deployment job must run here # too; without it, a server whose OAuth discovery failed at startup is rebuilt only diff --git a/litellm/proxy/search_endpoints/search_tool_management.py b/litellm/proxy/search_endpoints/search_tool_management.py index 69edf681e4d..81a008cf4c8 100644 --- a/litellm/proxy/search_endpoints/search_tool_management.py +++ b/litellm/proxy/search_endpoints/search_tool_management.py @@ -51,6 +51,20 @@ def _convert_datetime_to_str(value: datetime | str | None) -> str | None: TeamObjectLookup: TypeAlias = Callable[[str, UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable]] +async def _refresh_router_search_tools() -> None: + """Push the search tools table into this worker's router. + + Best-effort: the row is already committed, so a refresh failure must not surface as a 500 and + push the caller into a retry that creates duplicates. + """ + from litellm.proxy.proxy_server import proxy_config + + try: + await proxy_config.reload_search_tools_from_db() + except Exception as e: # noqa: BLE001 # the row is committed; no refresh failure may reach the caller + verbose_proxy_logger.exception("Search tool router refresh failed after a management write: %s", e) + + async def _team_object_from_db(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.proxy_server import ( @@ -305,8 +319,10 @@ async def create_search_tool(request: CreateSearchToolRequest): search_tool=request.search_tool, prisma_client=prisma_client ) + await _refresh_router_search_tools() + verbose_proxy_logger.debug( - "Successfully added search tool '%s' to database. Router will be updated by the cron job.", + "Successfully added search tool '%s' to database.", result.get("search_tool_name"), ) @@ -388,8 +404,10 @@ async def update_search_tool(search_tool_id: str, request: UpdateSearchToolReque prisma_client=prisma_client, ) + await _refresh_router_search_tools() + verbose_proxy_logger.debug( - "Successfully updated search tool '%s' in database. Router will be updated by the cron job.", + "Successfully updated search tool '%s' in database.", result.get("search_tool_name"), ) @@ -445,9 +463,9 @@ async def delete_search_tool(search_tool_id: str): search_tool_id=search_tool_id, prisma_client=prisma_client ) - verbose_proxy_logger.debug( - "Successfully deleted search tool from database. Router will be updated by the cron job." - ) + await _refresh_router_search_tools() + + verbose_proxy_logger.debug("Successfully deleted search tool from database.") return result except HTTPException as e: diff --git a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py index 7b895cd7fdb..70e9a96b316 100644 --- a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py +++ b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py @@ -992,3 +992,159 @@ async def test_list_search_tools_reports_a_missing_real_team_as_404(): assert response.status_code == 404 assert "search_tools" not in response.json() + + +# --------------------------------------------------------------------------- +# Router sync on management writes (LIT-3379) +# +# The proxy resolves prisma_client / proxy_config / llm_router from +# litellm.proxy.proxy_server module globals at call time and reaches its DB layer through a +# module-level registry singleton, so there is no constructor or parameter to inject through. +# Patching those globals is the only seam that exercises the endpoint end to end. +# --------------------------------------------------------------------------- + + +def _search_tool_row(name: str, provider: str = "tavily") -> dict: + return { + "search_tool_id": f"{name}-id", + "search_tool_name": name, + "litellm_params": {"search_provider": provider, "api_key": "sk-test"}, + "search_tool_info": {"description": name}, + } + + +def _fake_registry(db_rows: list) -> MagicMock: + """A registry singleton whose writes land in db_rows, so the refresh reads back real state.""" + + async def _add(search_tool, **_): + row = _search_tool_row( + search_tool["search_tool_name"], + provider=search_tool.get("litellm_params", {}).get("search_provider", "tavily"), + ) + db_rows.append(row) + return row + + async def _update(search_tool_id, search_tool, **_): + row = _search_tool_row( + search_tool["search_tool_name"], + provider=search_tool.get("litellm_params", {}).get("search_provider", "tavily"), + ) + db_rows[:] = [row if existing["search_tool_id"] == search_tool_id else existing for existing in db_rows] + return row + + async def _delete(search_tool_id, **_): + db_rows[:] = [existing for existing in db_rows if existing["search_tool_id"] != search_tool_id] + return {"message": "deleted", "search_tool_name": search_tool_id} + + async def _get_by_id(search_tool_id, **_): + return next((row for row in db_rows if row["search_tool_id"] == search_tool_id), None) + + registry = MagicMock() + registry.add_search_tool_to_db = AsyncMock(side_effect=_add) + registry.update_search_tool_in_db = AsyncMock(side_effect=_update) + registry.delete_search_tool_from_db = AsyncMock(side_effect=_delete) + registry.get_search_tool_by_id_from_db = AsyncMock(side_effect=_get_by_id) + return registry + + +@contextlib.contextmanager +def _live_router_and_db(db_rows: list): + """Drive the endpoints against a real ProxyConfig so the router refresh actually runs.""" + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config.update_config_state({}) + fake_router = MagicMock() + fake_router.search_tools = list(db_rows) + + with contextlib.ExitStack() as stack: + stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client", MagicMock())) # test-quality-ok: proxy globals are the only seam; see the module note above + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_config", proxy_config)) # test-quality-ok: proxy globals are the only seam; see the module note above + stack.enter_context(patch("litellm.proxy.proxy_server.llm_router", fake_router)) # test-quality-ok: proxy globals are the only seam; see the module note above + stack.enter_context( + patch( # test-quality-ok: proxy globals are the only seam; see the module note above + "litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY", + _fake_registry(db_rows), + ) + ) + stack.enter_context( + patch( # test-quality-ok: proxy globals are the only seam; see the module note above + "litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db", + AsyncMock(side_effect=lambda **_: list(db_rows)), + ) + ) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + yield fake_router + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_create_search_tool_reaches_the_router_before_the_response(): + """A UI-created tool must be usable immediately, not only after the next config reload tick.""" + with _live_router_and_db([]) as fake_router: + response = TestClient(app).post( + "/search_tools", + json={ + "search_tool": { + "search_tool_name": "tavily-search", + "litellm_params": {"search_provider": "tavily"}, + } + }, + ) + + assert response.status_code == 200 + assert [tool["search_tool_name"] for tool in fake_router.search_tools] == ["tavily-search"] + + +@pytest.mark.asyncio +async def test_update_search_tool_reaches_the_router_before_the_response(): + with _live_router_and_db([_search_tool_row("tavily-search", provider="tavily")]) as fake_router: + response = TestClient(app).put( + "/search_tools/tavily-search-id", + json={ + "search_tool": { + "search_tool_name": "tavily-search", + "litellm_params": {"search_provider": "exa_ai"}, + } + }, + ) + + assert response.status_code == 200 + assert fake_router.search_tools[0]["litellm_params"]["search_provider"] == "exa_ai" + + +@pytest.mark.asyncio +async def test_delete_search_tool_removes_it_from_the_router(): + """Deleting the last tool must clear the router; the old empty-list guard left it live.""" + with _live_router_and_db([_search_tool_row("tavily-search")]) as fake_router: + response = TestClient(app).delete("/search_tools/tavily-search-id") + + assert response.status_code == 200 + assert fake_router.search_tools == [] + + +@pytest.mark.asyncio +async def test_create_search_tool_survives_a_failing_router_refresh(): + """The row is already committed, so a refresh failure must not turn into a 500.""" + with _live_router_and_db([]): + with patch( # test-quality-ok: forcing the refresh to fail needs the refresh itself replaced + "litellm.proxy.proxy_server.ProxyConfig.reload_search_tools_from_db", + AsyncMock(side_effect=RuntimeError("registry boom")), + ): + response = TestClient(app).post( + "/search_tools", + json={ + "search_tool": { + "search_tool_name": "tavily-search", + "litellm_params": {"search_provider": "tavily"}, + } + }, + ) + + assert response.status_code == 200 + assert response.json()["search_tool_name"] == "tavily-search" diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index ee0de8840f6..7678b0cab9e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -1253,26 +1253,69 @@ async def test_ProxyConfig__init_search_tools_in_db_loads_merged_tools(monkeypat @pytest.mark.asyncio -async def test_ProxyConfig__init_search_tools_in_db_skips_empty_router_update(monkeypatch): +async def test_ProxyConfig__init_search_tools_in_db_clears_router_when_last_tool_is_deleted(monkeypatch): + """Deleting the last search tool must clear the router, not leave the tool live in memory.""" from litellm.proxy import proxy_server - from litellm.router_utils.search_api_router import SearchAPIRouter pc = ProxyConfig() pc.update_config_state({}) + fake_router = MagicMock() + fake_router.search_tools = [{"search_tool_name": "deleted-search", "litellm_params": {}}] mock_get_db_tools = AsyncMock(return_value=[]) - mock_update_router = AsyncMock() - monkeypatch.setattr(proxy_server, "llm_router", MagicMock()) + monkeypatch.setattr(proxy_server, "llm_router", fake_router) monkeypatch.setattr( "litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db", mock_get_db_tools, ) - monkeypatch.setattr(SearchAPIRouter, "update_router_search_tools", mock_update_router) await pc._init_search_tools_in_db(prisma_client=MagicMock()) mock_get_db_tools.assert_awaited_once() - mock_update_router.assert_not_awaited() + assert fake_router.search_tools == [] + + +@pytest.mark.asyncio +async def test_ProxyConfig_reload_search_tools_from_db_refreshes_router(monkeypatch): + from litellm.proxy import proxy_server + + pc = ProxyConfig() + mock_init = AsyncMock() + monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + + await pc.reload_search_tools_from_db() + + mock_init.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_ProxyConfig_reload_search_tools_from_db_honors_supported_db_objects(monkeypatch): + from litellm.proxy import proxy_server + + pc = ProxyConfig() + mock_init = AsyncMock() + monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]}) + + await pc.reload_search_tools_from_db() + + mock_init.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_ProxyConfig_reload_search_tools_from_db_noops_without_prisma(monkeypatch): + from litellm.proxy import proxy_server + + pc = ProxyConfig() + mock_init = AsyncMock() + monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init) + monkeypatch.setattr(proxy_server, "prisma_client", None) + + await pc.reload_search_tools_from_db() + + mock_init.assert_not_awaited() # --------------------------------------------------------------------------- From 8fcbc357e591435176248c995a1ce1cb2e0f4e51 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:45:25 -0700 Subject: [PATCH 03/21] fix(mcp): gate deferred oauth discovery on the endpoint each flow needs and read the resolved server The token exchange no longer joins deferred discovery when the token url is already stored, so it cannot 503 over an unreachable issuer it needs nothing from. After a request joins discovery, authorize and token now read the resolved server for the DCR bridge relay decision and the rest of the flow, so a registration endpoint resolved mid-request routes a front-door client to its own redirect binding. The encrypt seam in the issuer-yield authorize test now uses a real salt key instead of patching an SDK internal. --- .../mcp_server/discoverable_endpoints.py | 95 ++++++----- .../mcp_server/test_discoverable_endpoints.py | 150 ++++++++++++++++-- 2 files changed, 191 insertions(+), 54 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 6ede1c553e8..928373d93d8 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -3,7 +3,7 @@ import html as _html import json import secrets import time -from collections.abc import Mapping +from collections.abc import Callable, Mapping from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Optional from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse @@ -663,15 +663,18 @@ def _endpoint_not_configured_detail( ) -async def _server_with_oauth_endpoints(mcp_server: MCPServer) -> MCPServer: - """Join deferred OAuth discovery only when this server still has no authorize URL. +async def _server_with_oauth_endpoints( + mcp_server: MCPServer, + needed_endpoint: Callable[[MCPServer], str | None], +) -> MCPServer: + """Join deferred OAuth discovery only when the endpoint this caller needs is still missing. Admin-entered endpoints live on ``configured_*`` after an anchored issuer empties the - resolved fields. Those already let authorize/token run, so discovery is not awaited - and cannot 503 over a leftover pin. A server with nothing configured still joins the - deferred task; no slot is a no-op and the caller 400s. + resolved fields. A caller whose needed endpoint already resolves never awaits discovery + and cannot 503 over a leftover pin. A server still missing it joins the deferred task; + no slot is a no-op and the caller 400s. """ - if mcp_server.effective_authorization_url is not None: + if needed_endpoint(mcp_server) is not None: return mcp_server from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import with mcp_server_manager at module load global_mcp_server_manager, @@ -829,7 +832,7 @@ async def authorize_with_server( ephemeral_dcr_client: "EphemeralDcrClient | None" = None, ): _raise_if_not_oauth2(mcp_server) - resolved_server: Final = await _server_with_oauth_endpoints(mcp_server) + resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, lambda s: s.effective_authorization_url) if resolved_server.effective_authorization_url is None: raise HTTPException( status_code=400, @@ -841,7 +844,7 @@ async def authorize_with_server( ), ) - if mcp_server.is_dcr_bridge: + if resolved_server.is_dcr_bridge: # Enforce S256 PKCE on both bridge arms. The relay arm forwards the validated, # now-non-optional pair to the upstream authorize; the short-circuit arm keeps # calling this for its enforcement side effect, then falls through to the gateway @@ -850,9 +853,9 @@ async def authorize_with_server( # A gateway-minted ephemeral client is registered against {base}/callback, so its # flow must run the short-circuit arm; the relay arm is only for clients that # registered themselves through the front door and hold their own redirect binding. - if _dcr_bridge_relays_client_registration(mcp_server) and ephemeral_dcr_client is None: + if _dcr_bridge_relays_client_registration(resolved_server) and ephemeral_dcr_client is None: return _redirect_to_upstream_authorize( - mcp_server=mcp_server, + mcp_server=resolved_server, client_id=client_id, redirect_uri=redirect_uri, state=state, @@ -878,7 +881,7 @@ async def authorize_with_server( # litellm key, so the browser session is the only identity source; without one there is nothing to # bind, so send the user through login first. Every other oauth2 server keeps the identity-less state. litellm_user_id: str | None = None - if mcp_server.is_dcr_bridge and mcp_server.is_oauth_delegate: + if resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate: from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import _user_id_from_session_cookie, ) @@ -888,7 +891,7 @@ async def authorize_with_server( return _redirect_to_litellm_login(request) denial: Final = await _bridge_authorize_access_denial( litellm_user_id=litellm_user_id, - mcp_server=mcp_server, + mcp_server=resolved_server, redirect_uri=redirect_uri, state=state, ) @@ -902,7 +905,7 @@ async def authorize_with_server( code_challenge_method=code_challenge_method, client_redirect_uri=redirect_uri, litellm_user_id=litellm_user_id, - mcp_server_id=mcp_server.server_id if (litellm_user_id or ephemeral_dcr_client) else None, + mcp_server_id=resolved_server.server_id if (litellm_user_id or ephemeral_dcr_client) else None, dcr_client_id=ephemeral_dcr_client.client_id if ephemeral_dcr_client else None, dcr_client_secret=ephemeral_dcr_client.client_secret if ephemeral_dcr_client else None, dcr_token_endpoint_auth_method=ephemeral_dcr_client.token_endpoint_auth_method @@ -912,22 +915,22 @@ async def authorize_with_server( relay_state: Final = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES) params: Final = { - "client_id": mcp_server.client_id if mcp_server.client_id else client_id, + "client_id": resolved_server.client_id if resolved_server.client_id else client_id, "redirect_uri": f"{request_base_url}/callback", "state": relay_state, "response_type": response_type or "code", } if scope: params["scope"] = scope - elif mcp_server.scopes: - params["scope"] = " ".join(mcp_server.scopes) + elif resolved_server.scopes: + params["scope"] = " ".join(resolved_server.scopes) if code_challenge: params["code_challenge"] = code_challenge if code_challenge_method: params["code_challenge_method"] = code_challenge_method - upstream_resource: Final = resolve_upstream_resource(mcp_server) + upstream_resource: Final = resolve_upstream_resource(resolved_server) if upstream_resource: params["resource"] = upstream_resource @@ -964,7 +967,7 @@ async def exchange_token_with_server( if grant_type not in ("authorization_code", "refresh_token"): raise HTTPException(status_code=400, detail="Unsupported grant_type") - resolved_server: Final = await _server_with_oauth_endpoints(mcp_server) + resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, lambda s: s.effective_token_url) token_url: Final = resolved_server.effective_token_url if token_url is None: raise HTTPException( @@ -985,16 +988,16 @@ async def exchange_token_with_server( # recovered from a sealed code) must authenticate the way its own registration was granted, # not the way the server row is configured; callers that carry no method keep the row's method # as before. - resolved_client_id: Final = mcp_server.client_id if mcp_server.client_id else client_id - resolved_client_secret: Final = mcp_server.client_secret if mcp_server.client_id else client_secret + resolved_client_id: Final = resolved_server.client_id if resolved_server.client_id else client_id + resolved_client_secret: Final = resolved_server.client_secret if resolved_server.client_id else client_secret resolved_auth_method: Final = ( - mcp_server.token_endpoint_auth_method - if mcp_server.client_id - else (client_token_endpoint_auth_method or mcp_server.token_endpoint_auth_method) + resolved_server.token_endpoint_auth_method + if resolved_server.client_id + else (client_token_endpoint_auth_method or resolved_server.token_endpoint_auth_method) ) try: token_request: Final = build_upstream_oauth2_token_request( - mcp_server, + resolved_server, auth_method=resolved_auth_method, client_id=resolved_client_id, client_secret=resolved_client_secret, @@ -1007,14 +1010,14 @@ async def exchange_token_with_server( bridge_upstream_refresh: SecretStr | None = None bridge_upstream_scope: str | None = None refresh_request_scope: str | None = None - is_bridge: Final = mcp_server.is_oauth_delegate and mcp_server.is_dcr_bridge + is_bridge: Final = resolved_server.is_oauth_delegate and resolved_server.is_dcr_bridge if grant_type == "refresh_token": # Phase 1 for a bridge refresh: open the client's refresh envelope, re-validate the sealed # identity, and unwrap the real upstream refresh token BEFORE building token_data, so the exchange # sends the upstream token and never the envelope. A failure returns without touching the upstream. if is_bridge: - prepared_refresh: Final = await _prepare_bridge_refresh(mcp_server, refresh_token) + prepared_refresh: Final = await _prepare_bridge_refresh(resolved_server, refresh_token) if not isinstance(prepared_refresh, _BridgeRefreshReady): return _bridge_mint_error_response(prepared_refresh) bridge_mint_ready = prepared_refresh.ready @@ -1051,13 +1054,13 @@ async def exchange_token_with_server( # A raw upstream code (scripted path) opens to None and the code is used as-is. bridge_identity = open_bridge_authorization_code(code) if bridge_identity is not None: - if bridge_identity.mcp_server_id != mcp_server.server_id: + if bridge_identity.mcp_server_id != resolved_server.server_id: raise HTTPException( status_code=400, detail="Authorization code was issued for a different MCP server", ) code = bridge_identity.upstream_code - bridge_token_relay: Final = _dcr_bridge_relays_client_registration(mcp_server) + bridge_token_relay: Final = _dcr_bridge_relays_client_registration(resolved_server) if bridge_token_relay and not redirect_uri: raise HTTPException( status_code=400, @@ -1079,7 +1082,7 @@ async def exchange_token_with_server( # Phase 1 for a bridge authorization_code mint: resolve identity (the SSO user recovered above, or # the presented litellm key) and the envelope keys BEFORE the exchange consumes the single-use code. if is_bridge: - prepared: Final = await _prepare_bridge_mint(request, mcp_server, bridge_identity) + prepared: Final = await _prepare_bridge_mint(request, resolved_server, bridge_identity) if not isinstance(prepared, _BridgeMintReady): return _bridge_mint_error_response(prepared) bridge_mint_ready = prepared @@ -1096,8 +1099,8 @@ async def exchange_token_with_server( except httpx.HTTPStatusError as exc: fault: Final = classify_upstream_token_rejection( exc.response, - credential_source=_token_credential_source(mcp_server), - log_context=mcp_server.server_id, + credential_source=_token_credential_source(resolved_server), + log_context=resolved_server.server_id, ) upstream_rejected_bridge_refresh: Final = ( is_bridge @@ -1110,7 +1113,7 @@ async def exchange_token_with_server( "bridge refresh: the upstream rejected the sealed refresh token for server=%s with " "invalid_grant (revoked or expired at the IdP); returning invalid_grant so the client " "re-runs authorization_code rather than an opaque upstream error", - mcp_server.server_id, + resolved_server.server_id, ) return _bridge_mint_error_response("invalid_refresh") return render_token_fault(fault) @@ -1123,22 +1126,22 @@ async def exchange_token_with_server( # Validate token response against server-configured rules before any storage. # This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc. - if mcp_server.token_validation and isinstance(mcp_server.token_validation, dict): + if resolved_server.token_validation and isinstance(resolved_server.token_validation, dict): _validate_token_response( token_response=token_response, - validation_rules=mcp_server.token_validation, - server_id=mcp_server.server_id, + validation_rules=resolved_server.token_validation, + server_id=resolved_server.server_id, ) # Store server-side when the server is configured for per-user OAuth and # the calling client has provided a valid LiteLLM identity. # Errors are non-fatal: the token is still returned to the client. - if mcp_server.needs_user_oauth_token: + if resolved_server.needs_user_oauth_token: user_id: Final = await _extract_user_id_from_request(request) if user_id: try: await _store_per_user_token_server_side( - server=mcp_server, + server=resolved_server, user_id=user_id, token_response=token_response, ) @@ -1146,7 +1149,7 @@ async def exchange_token_with_server( verbose_logger.warning( "exchange_token_with_server: server-side storage failed for user=%s server=%s: %s", user_id, - mcp_server.server_id, + resolved_server.server_id, exc, ) else: @@ -1156,7 +1159,7 @@ async def exchange_token_with_server( "requires the stored token, so the client will be challenged with 401 on reconnect. " "Ensure the request carries a valid LiteLLM key (x-litellm-api-key or Authorization), " "or store it via POST /mcp/server/{id}/oauth-user-credential.", - mcp_server.server_id, + resolved_server.server_id, ) # A DCR-bridge oauth_delegate server hands the client a gateway-bound envelope (identity plus the @@ -1167,7 +1170,9 @@ async def exchange_token_with_server( token_response = {**token_response, "scope": refresh_request_scope} # Phase 3: seal the upstream grant into the client-held envelope; failures map through the same # OAuth-shaped response as the phase-1 preconditions. - minted: Final = _finish_bridge_mint(bridge_mint_ready, mcp_server, token_response, datetime.now(timezone.utc)) + minted: Final = _finish_bridge_mint( + bridge_mint_ready, resolved_server, token_response, datetime.now(timezone.utc) + ) return minted if isinstance(minted, JSONResponse) else _bridge_mint_error_response(minted) raw_access_token: Final = token_response.get("access_token") if isinstance(token_response, dict) else None @@ -1682,7 +1687,7 @@ async def register_client_with_server( ): return dummy_return - resolved_server: Final = await _server_with_oauth_endpoints(mcp_server) + resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, lambda s: s.effective_authorization_url) if resolved_server.effective_authorization_url is None: raise HTTPException( status_code=400, @@ -1715,13 +1720,15 @@ async def register_client_with_server( response: Final = await _post_dcr_registration( registration_url=registration_url, register_data=register_data, - server_id=mcp_server.server_id, + server_id=resolved_server.server_id, ) token_response = response.json() if persist_credentials and not bridge_relay: - persistence_result = await _persist_dcr_client_registration(mcp_server, token_response, current_redirect_uri) + persistence_result = await _persist_dcr_client_registration( + resolved_server, token_response, current_redirect_uri + ) if persistence_result == "reused": return dummy_return diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 790b5c7ad22..46c9c21e80d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -8855,7 +8855,7 @@ async def test_authorize_wall_names_the_issuer_for_anchored_servers(): @pytest.mark.asyncio -async def test_authorize_uses_admin_entered_github_oauth_urls_after_issuer_yield(): +async def test_authorize_uses_admin_entered_github_oauth_urls_after_issuer_yield(monkeypatch): """GitHub MCP servers store Authorization URL and Token URL on the row. 1.99 can empty the resolved authorization_url when a leftover issuer is treated as a pin (RFC 8414 yield). The UI authorize must still redirect to the admin-entered GitHub authorize URL @@ -8887,15 +8887,14 @@ async def test_authorize_uses_admin_entered_github_oauth_urls_after_issuer_yield mock_request.base_url = "https://litellm.example.com/" mock_request.headers = {} - with patch("litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper") as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state" - response = await authorize_with_server( - request=mock_request, - mcp_server=server, - client_id="github-app-client", - redirect_uri="http://127.0.0.1:60108/callback", - state="state123", - ) + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-lit-6255") + response = await authorize_with_server( + request=mock_request, + mcp_server=server, + client_id="github-app-client", + redirect_uri="http://127.0.0.1:60108/callback", + state="state123", + ) assert response.status_code == 307 assert "https://github.com/login/oauth/authorize" in response.headers["location"] @@ -8926,6 +8925,137 @@ def test_oauth_endpoints_count_admin_entered_urls_as_resolved(): assert _oauth_endpoints_unresolved(server) is False +@pytest.mark.asyncio +async def test_token_exchange_with_configured_token_url_never_joins_discovery(monkeypatch): + """A server can hold an admin-entered Token URL while its Authorization URL is absent. The + token exchange must post to that stored endpoint without awaiting deferred discovery, which + can 503 against an unreachable issuer even though nothing it resolves is needed here.""" + from litellm.proxy._experimental.mcp_server import ( + discoverable_endpoints, + mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="token-url-only", + name="token_url_only", + server_name="token_url_only", + alias="token_url_only", + url="https://mcp.example.com/mcp/", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + client_secret="cs", + authorization_url=None, + token_url=None, + issuer="https://idp.example.com", + issuer_is_anchored=True, + configured_token_url="https://idp.example.com/oauth/token", + ) + + async def fail_discovery(_srv): + raise AssertionError("the exchange joined deferred discovery despite a stored token url") + + monkeypatch.setattr( + mcp_server_manager.global_mcp_server_manager, + "ensure_oauth_metadata_discovered", + fail_discovery, + ) + fake_http_response = MagicMock() + fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"} + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + monkeypatch.setattr( + discoverable_endpoints, + "get_async_httpx_client", + lambda llm_provider: fake_http_client, + ) + mock_request = MagicMock() + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + response = await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code="upstream-code", + redirect_uri="http://127.0.0.1:3000/cb", + client_id="cid", + client_secret=None, + code_verifier=None, + ) + + assert response.status_code == 200 + assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/oauth/token" + + +@pytest.mark.asyncio +async def test_bridge_authorize_relays_with_registration_url_resolved_by_deferred_discovery(monkeypatch): + """When deferred discovery resolves a DCR-bridge server during the authorize request, the + relay-vs-short-circuit call must read the resolved server: a client that registered itself + through the front door keeps its own redirect binding instead of being routed through the + gateway callback the upstream never granted it.""" + from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize_with_server, + ) + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="bridge-deferred", + name="bridge_deferred", + server_name="bridge_deferred", + alias="bridge_deferred", + url="https://mcp.example.com/mcp/", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + dcr_bridge=True, + authorization_url=None, + token_url=None, + registration_url=None, + ) + resolved = server.model_copy( + update={ + "authorization_url": "https://idp.example.com/oauth/authorize", + "token_url": "https://idp.example.com/oauth/token", + "registration_url": "https://idp.example.com/oauth/register", + } + ) + + async def resolve_discovery(_srv): + return resolved + + monkeypatch.setattr( + mcp_server_manager.global_mcp_server_manager, + "ensure_oauth_metadata_discovered", + resolve_discovery, + ) + mock_request = MagicMock() + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + response = await authorize_with_server( + request=mock_request, + mcp_server=server, + client_id="front-door-client", + redirect_uri="http://127.0.0.1:60110/client-callback", + state="state456", + code_challenge="E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM", + code_challenge_method="S256", + ) + + assert response.status_code == 307 + location = response.headers["location"] + assert location.startswith("https://idp.example.com/oauth/authorize") + assert "redirect_uri=http%3A%2F%2F127.0.0.1%3A60110%2Fclient-callback" in location + + def test_passthrough_authorization_code_round_trips_and_rejects_hostile_input(): """The passthrough gateway code seals and recovers the ephemeral DCR client and upstream code, and is total over hostile input: a raw upstream code opens to None, and a tampered or From 19e6f03a3c08c71c9307d5588d7277604acb4c64 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:39:35 -0700 Subject: [PATCH 04/21] fix(mcp): let root oauth routes defer discovery to the endpoint-gated flow join --- .../mcp_server/discoverable_endpoints.py | 28 +--- .../mcp_server/mcp_server_manager.py | 8 -- .../mcp_server/test_discoverable_endpoints.py | 125 ++++++++++++++---- 3 files changed, 108 insertions(+), 53 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 928373d93d8..6274cfcef8b 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1785,17 +1785,10 @@ async def authorize( lookup_name: Final[str | None] = mcp_server_name or client_id client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) mcp_server = ( - await global_mcp_server_manager.get_resolved_mcp_server_by_name(lookup_name, client_ip=client_ip) - if lookup_name - else None + global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) if lookup_name else None ) if mcp_server is None and mcp_server_name is None: - unresolved_server: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) - mcp_server = ( - await global_mcp_server_manager.ensure_oauth_metadata_discovered(unresolved_server) - if unresolved_server is not None - else None - ) + mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if mcp_server is None: raise HTTPException(status_code=404, detail="MCP server not found") _raise_if_not_oauth2(mcp_server) @@ -1876,14 +1869,9 @@ async def token_endpoint( lookup_name: Final = mcp_server_name or client_id client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) - mcp_server = await global_mcp_server_manager.get_resolved_mcp_server_by_name(lookup_name, client_ip=client_ip) + mcp_server = global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) if mcp_server is None and mcp_server_name is None: - unresolved_server: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) - mcp_server = ( - await global_mcp_server_manager.ensure_oauth_metadata_discovered(unresolved_server) - if unresolved_server is not None - else None - ) + mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if mcp_server is None: raise HTTPException(status_code=404, detail="MCP server not found") return await exchange_token_with_server( @@ -2714,10 +2702,9 @@ async def register_client(request: Request, mcp_server_name: str | None = None): return await register_aggregate_client(request=request, request_body=data) resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if resolved: - resolved_server: Final = await global_mcp_server_manager.ensure_oauth_metadata_discovered(resolved) return await register_client_with_server( request=request, - mcp_server=resolved_server, + mcp_server=resolved, client_name=data.get("client_name", ""), grant_types=data.get("grant_types", []), response_types=data.get("response_types", []), @@ -2727,10 +2714,7 @@ async def register_client(request: Request, mcp_server_name: str | None = None): ) return dummy_return - mcp_server: Final = await global_mcp_server_manager.get_resolved_mcp_server_by_name( - mcp_server_name, - client_ip=client_ip, - ) + mcp_server: Final = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) if mcp_server is None: return dummy_return return await register_client_with_server( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index fe0c73a5efc..308813039ca 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -6205,14 +6205,6 @@ class MCPServerManager: return server return None - async def get_resolved_mcp_server_by_name( - self, - server_name: str, - client_ip: str | None = None, - ) -> MCPServer | None: - server: Final = self.get_mcp_server_by_name(server_name, client_ip=client_ip) - return await self.ensure_oauth_metadata_discovered(server) if server is not None else None - def get_filtered_registry(self, client_ip: str | None = None) -> dict[str, MCPServer]: """ Get registry filtered by client IP access control. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 46c9c21e80d..f67e9a67d53 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -79,24 +79,23 @@ def _resolved_oauth_metadata(): @pytest.mark.asyncio -async def test_authorize_resolves_cold_oauth_metadata(): +async def test_authorize_resolves_cold_oauth_metadata(monkeypatch): + """The route hands the registered server to the flow, whose deferred-discovery join resolves + the cold metadata; the redirect must land on the discovered authorization endpoint.""" from litellm.proxy._experimental.mcp_server import discoverable_endpoints from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-lit-6255") server = _unresolved_oauth_server() global_mcp_server_manager.registry[server.server_id] = server global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True) request = _mock_callback_request("https://litellm.example.com/") - expected = MagicMock() - with ( - patch.object( - global_mcp_server_manager, - "_discover_oauth_metadata_for_server", - new=AsyncMock(return_value=_resolved_oauth_metadata()), - ) as discovery, - patch.object(discoverable_endpoints, "authorize_with_server", new=AsyncMock(return_value=expected)) as relay, - ): + with patch.object( + global_mcp_server_manager, + "_discover_oauth_metadata_for_server", + new=AsyncMock(return_value=_resolved_oauth_metadata()), + ) as discovery: response = await discoverable_endpoints.authorize( request=request, client_id="client-id", @@ -105,12 +104,14 @@ async def test_authorize_resolves_cold_oauth_metadata(): ) discovery.assert_awaited_once_with(server) - assert relay.await_args.kwargs["mcp_server"].authorization_url == "https://idp.example.com/authorize" - assert response is expected + assert response.status_code == 307 + assert response.headers["location"].startswith("https://idp.example.com/authorize") @pytest.mark.asyncio async def test_token_resolves_cold_oauth_metadata(): + """The route hands the registered server to the exchange, whose deferred-discovery join + resolves the cold metadata; the exchange must post to the discovered token endpoint.""" from litellm.proxy._experimental.mcp_server import discoverable_endpoints from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager @@ -118,7 +119,11 @@ async def test_token_resolves_cold_oauth_metadata(): global_mcp_server_manager.registry[server.server_id] = server global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True) request = _mock_callback_request("https://litellm.example.com/") - expected = MagicMock() + fake_http_response = MagicMock() + fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"} + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) with ( patch.object( @@ -127,8 +132,10 @@ async def test_token_resolves_cold_oauth_metadata(): new=AsyncMock(return_value=_resolved_oauth_metadata()), ) as discovery, patch.object( - discoverable_endpoints, "exchange_token_with_server", new=AsyncMock(return_value=expected) - ) as relay, + discoverable_endpoints, + "get_async_httpx_client", + new=lambda llm_provider: fake_http_client, + ), ): response = await discoverable_endpoints.token_endpoint( request=request, @@ -139,20 +146,26 @@ async def test_token_resolves_cold_oauth_metadata(): ) discovery.assert_awaited_once_with(server) - assert relay.await_args.kwargs["mcp_server"].token_url == "https://idp.example.com/token" - assert response is expected + assert response.status_code == 200 + assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/token" @pytest.mark.asyncio async def test_register_resolves_cold_oauth_metadata(): + """The route hands the registered server to the registration flow, whose deferred-discovery + join resolves the cold metadata; DCR must post to the discovered registration endpoint.""" from litellm.proxy._experimental.mcp_server import discoverable_endpoints from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager - server = _unresolved_oauth_server() + server = _unresolved_oauth_server().model_copy(update={"client_id": None}) global_mcp_server_manager.registry[server.server_id] = server global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True) request = _mock_callback_request("https://litellm.example.com/") - expected = MagicMock() + fake_http_response = MagicMock() + fake_http_response.json.return_value = {"client_id": "generated-client", "client_secret": "generated-secret"} + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) with ( patch.object( @@ -162,14 +175,16 @@ async def test_register_resolves_cold_oauth_metadata(): ) as discovery, patch.object(discoverable_endpoints, "_read_request_body", new=AsyncMock(return_value={})), patch.object( - discoverable_endpoints, "register_client_with_server", new=AsyncMock(return_value=expected) - ) as relay, + discoverable_endpoints, + "get_async_httpx_client", + new=lambda llm_provider: fake_http_client, + ), ): response = await discoverable_endpoints.register_client(request=request, mcp_server_name=server.server_name) discovery.assert_awaited_once_with(server) - assert relay.await_args.kwargs["mcp_server"].registration_url == "https://idp.example.com/register" - assert response is expected + assert response.status_code == 200 + assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/register" @pytest.fixture @@ -8994,6 +9009,70 @@ async def test_token_exchange_with_configured_token_url_never_joins_discovery(mo assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/oauth/token" +@pytest.mark.asyncio +async def test_root_token_route_with_configured_token_url_never_joins_discovery(monkeypatch): + """A root POST /token that falls back to the sole OAuth2 server must reach the exchange's + endpoint-gated discovery join instead of awaiting full discovery at the route: with the + token url admin-entered, a failing or slow discovery must not turn the exchange into a 503.""" + from litellm.proxy._experimental.mcp_server import ( + discoverable_endpoints, + mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + manager = mcp_server_manager.global_mcp_server_manager + server = MCPServer( + server_id="sole-token-url-only", + name="sole_token_url_only", + server_name="sole_token_url_only", + alias="sole_token_url_only", + url="https://mcp.example.com/mcp/", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + client_secret="cs", + issuer="https://idp.example.com", + issuer_is_anchored=True, + configured_token_url="https://idp.example.com/oauth/token", + ) + saved_registry = dict(manager.registry) + manager.registry.clear() + manager.registry[server.server_id] = server + manager._set_oauth_discovery_deferred(server.server_id, True) + + async def fail_discovery(_srv): + raise AssertionError("the root token route joined deferred discovery despite a stored token url") + + monkeypatch.setattr(manager, "ensure_oauth_metadata_discovered", fail_discovery) + fake_http_response = MagicMock() + fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"} + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + monkeypatch.setattr( + discoverable_endpoints, + "get_async_httpx_client", + lambda llm_provider: fake_http_client, + ) + request = _mock_callback_request("https://litellm.example.com/") + + try: + response = await discoverable_endpoints.token_endpoint( + request=request, + grant_type="authorization_code", + code="upstream-code", + redirect_uri="http://127.0.0.1:3000/cb", + client_id="unregistered-dcr-client", + ) + finally: + manager.registry.clear() + manager.registry.update(saved_registry) + + assert response.status_code == 200 + assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/oauth/token" + + @pytest.mark.asyncio async def test_bridge_authorize_relays_with_registration_url_resolved_by_deferred_discovery(monkeypatch): """When deferred discovery resolves a DCR-bridge server during the authorize request, the From 465ebb1bdd3a143c72179b2babe027d51ebb12ba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 14:08:09 -0700 Subject: [PATCH 05/21] fix(mcp): join discovery for a clientless DCR bridge still missing its registration endpoint --- .../mcp_server/discoverable_endpoints.py | 15 ++++- .../mcp_server/test_discoverable_endpoints.py | 60 +++++++++++++++++++ 2 files changed, 73 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 6274cfcef8b..46feeab4dc3 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -832,7 +832,7 @@ async def authorize_with_server( ephemeral_dcr_client: "EphemeralDcrClient | None" = None, ): _raise_if_not_oauth2(mcp_server) - resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, lambda s: s.effective_authorization_url) + resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _register_flow_needed_endpoint) if resolved_server.effective_authorization_url is None: raise HTTPException( status_code=400, @@ -1653,6 +1653,17 @@ async def resolve_ephemeral_dcr_client( return await mint_ephemeral_dcr_client(request, mcp_server) +def _register_flow_needed_endpoint(mcp_server: MCPServer) -> str | None: + """The register flow's deferred-discovery join gate. A DCR bridge with no admin-configured + client can only register callers through the upstream's registration endpoint + (``_oauth_endpoints_unresolved`` keeps its discovery slot armed for exactly this shape), so + the flow must keep joining discovery while registration is still missing instead of silently + degrading to the dummy short-circuit. Every other shape only needs the authorization url.""" + if mcp_server.is_dcr_bridge and not mcp_server.client_id and mcp_server.effective_registration_url is None: + return None + return mcp_server.effective_authorization_url + + async def register_client_with_server( request: Request, mcp_server: MCPServer, @@ -1687,7 +1698,7 @@ async def register_client_with_server( ): return dummy_return - resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, lambda s: s.effective_authorization_url) + resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _register_flow_needed_endpoint) if resolved_server.effective_authorization_url is None: raise HTTPException( status_code=400, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index f67e9a67d53..b1a99b498e5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -187,6 +187,66 @@ async def test_register_resolves_cold_oauth_metadata(): assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/register" +@pytest.mark.asyncio +async def test_register_route_bridge_missing_registration_url_joins_discovery(): + """A clientless DCR bridge whose authorize and token urls are admin-entered still relays + registration upstream: the flow must join deferred discovery for the missing registration + endpoint instead of short-circuiting to dummy credentials because authorization resolves.""" + import json + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="bridge-partial-metadata", + name="bridge_partial_metadata", + server_name="bridge_partial_metadata", + alias="bridge_partial_metadata", + url="https://mcp.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth_delegate, + dcr_bridge=True, + client_id=None, + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + ) + global_mcp_server_manager.registry[server.server_id] = server + global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True) + request = _mock_callback_request("https://litellm.example.com/") + fake_http_response = MagicMock() + fake_http_response.json.return_value = {"client_id": "generated-client", "client_secret": "generated-secret"} + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + + with ( + patch.object( # test-quality-ok: innermost discovery seam on a module-global manager; the route-to-flow join under test stays real + global_mcp_server_manager, + "_discover_oauth_metadata_for_server", + new=AsyncMock(return_value=_resolved_oauth_metadata()), + ) as discovery, + patch.object( # test-quality-ok: the MagicMock Request carries no body; this seam feeds the RFC 7591 redirect_uris + discoverable_endpoints, + "_read_request_body", + new=AsyncMock(return_value={"redirect_uris": ["https://client.example.com/cb"]}), + ), + patch.object( # test-quality-ok: keeps the DCR POST off the network so its target URL can be asserted + discoverable_endpoints, + "get_async_httpx_client", + new=lambda llm_provider: fake_http_client, + ), + ): + response = await discoverable_endpoints.register_client(request=request, mcp_server_name=server.server_name) + + discovery.assert_awaited_once_with(server) + assert fake_http_client.post.await_args.args[0] == "https://idp.example.com/register" + assert fake_http_client.post.await_args.kwargs["json"]["redirect_uris"] == ["https://client.example.com/cb"] + assert response.status_code == 200 + assert json.loads(response.body.decode("utf-8"))["client_id"] == "generated-client" + + @pytest.fixture def trust_xff(): """Force ``IPAddressUtils.is_request_from_trusted_proxy`` to True. From 2548e960f18d5260fc4fb511f2d67406cb8ce8f7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 14:43:21 -0700 Subject: [PATCH 06/21] fix(cost-map): correct Gemini TTS and native-audio rates Gemini 2.5 Flash Preview TTS, Gemini 2.5 Pro Preview TTS, and the three gemini-2.5-flash-native-audio entries carried rates copied from the text models, so audio output was billed 2x to 6x under Google's published prices. Set the published per-token rates on all ten keys, add output_cost_per_audio_token to the native-audio entries, and drop the long-context tier rates Google does not publish for Pro TTS. --- ...odel_prices_and_context_window_backup.json | 84 +++++----- model_prices_and_context_window.json | 84 +++++----- .../test_gemini_tts_native_audio_pricing.py | 144 ++++++++++++++++++ 3 files changed, 228 insertions(+), 84 deletions(-) create mode 100644 tests/test_litellm/test_gemini_tts_native_audio_pricing.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d9a13ef1b98..467acce9850 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -21091,18 +21091,15 @@ }, "gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", + "output_cost_per_token": 2e-05, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -22212,11 +22209,11 @@ "supports_image_size": false }, "gemini/gemini-2.5-flash-preview-tts": { - "input_cost_per_token": 3e-07, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "mode": "audio_speech", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_token": 1e-05, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ], @@ -23156,19 +23153,16 @@ }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token": 1e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_token": 2e-05, "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -48804,15 +48798,16 @@ } }, "gemini-2.5-flash-native-audio-latest": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -48829,15 +48824,16 @@ "gemini_native_audio": true }, "gemini-2.5-flash-native-audio-preview-09-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -48854,15 +48850,16 @@ "gemini_native_audio": true }, "gemini-2.5-flash-native-audio-preview-12-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -48912,15 +48909,16 @@ "gemini_audio_only_live": true }, "gemini/gemini-2.5-flash-native-audio-latest": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -48939,15 +48937,16 @@ "gemini_native_audio": true }, "gemini/gemini-2.5-flash-native-audio-preview-09-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -48966,15 +48965,16 @@ "gemini_native_audio": true }, "gemini/gemini-2.5-flash-native-audio-preview-12-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -49043,11 +49043,11 @@ "rpm": 10 }, "gemini-2.5-flash-preview-tts": { - "input_cost_per_token": 3e-07, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "mode": "audio_speech", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_token": 1e-05, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d9a13ef1b98..467acce9850 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -21091,18 +21091,15 @@ }, "gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", + "output_cost_per_token": 2e-05, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -22212,11 +22209,11 @@ "supports_image_size": false }, "gemini/gemini-2.5-flash-preview-tts": { - "input_cost_per_token": 3e-07, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "mode": "audio_speech", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_token": 1e-05, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ], @@ -23156,19 +23153,16 @@ }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_audio_token": 7e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "input_cost_per_token": 1e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_token": 2e-05, "rpm": 10000, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-pro-preview", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_modalities": [ "text" ], @@ -48804,15 +48798,16 @@ } }, "gemini-2.5-flash-native-audio-latest": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -48829,15 +48824,16 @@ "gemini_native_audio": true }, "gemini-2.5-flash-native-audio-preview-09-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -48854,15 +48850,16 @@ "gemini_native_audio": true }, "gemini-2.5-flash-native-audio-preview-12-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -48912,15 +48909,16 @@ "gemini_audio_only_live": true }, "gemini/gemini-2.5-flash-native-audio-latest": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -48939,15 +48937,16 @@ "gemini_native_audio": true }, "gemini/gemini-2.5-flash-native-audio-preview-09-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -48966,15 +48965,16 @@ "gemini_native_audio": true }, "gemini/gemini-2.5-flash-native-audio-preview-12-2025": { - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/realtime" ], @@ -49043,11 +49043,11 @@ "rpm": 10 }, "gemini-2.5-flash-preview-tts": { - "input_cost_per_token": 3e-07, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "mode": "audio_speech", - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/pricing", + "output_cost_per_token": 1e-05, + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/audio/speech" ] diff --git a/tests/test_litellm/test_gemini_tts_native_audio_pricing.py b/tests/test_litellm/test_gemini_tts_native_audio_pricing.py new file mode 100644 index 00000000000..88fd14e436d --- /dev/null +++ b/tests/test_litellm/test_gemini_tts_native_audio_pricing.py @@ -0,0 +1,144 @@ +import json +from pathlib import Path +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.types.utils import CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, Usage + +REPO_ROOT: Final = Path(__file__).parents[2] +MAIN_PATH: Final = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH: Final = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +FLASH_TTS_KEYS: Final = ("gemini-2.5-flash-preview-tts", "gemini/gemini-2.5-flash-preview-tts") +PRO_TTS_KEYS: Final = ("gemini-2.5-pro-preview-tts", "gemini/gemini-2.5-pro-preview-tts") +NATIVE_AUDIO_KEYS: Final = tuple( + f"{prefix}gemini-2.5-flash-native-audio-{suffix}" + for prefix in ("", "gemini/") + for suffix in ("latest", "preview-09-2025", "preview-12-2025") +) + +FLASH_TTS_INPUT: Final = 5e-07 +FLASH_TTS_AUDIO_OUTPUT: Final = 1e-05 +PRO_TTS_INPUT: Final = 1e-06 +PRO_TTS_AUDIO_OUTPUT: Final = 2e-05 +NATIVE_AUDIO_TEXT_INPUT: Final = 5e-07 +NATIVE_AUDIO_AUDIO_INPUT: Final = 3e-06 +NATIVE_AUDIO_TEXT_OUTPUT: Final = 2e-06 +NATIVE_AUDIO_AUDIO_OUTPUT: Final = 1.2e-05 + +PUBLISHED_RATES: Final = { + **{ + key: {"input_cost_per_token": FLASH_TTS_INPUT, "output_cost_per_token": FLASH_TTS_AUDIO_OUTPUT} + for key in FLASH_TTS_KEYS + }, + **{ + key: {"input_cost_per_token": PRO_TTS_INPUT, "output_cost_per_token": PRO_TTS_AUDIO_OUTPUT} + for key in PRO_TTS_KEYS + }, + **{ + key: { + "input_cost_per_token": NATIVE_AUDIO_TEXT_INPUT, + "input_cost_per_audio_token": NATIVE_AUDIO_AUDIO_INPUT, + "output_cost_per_token": NATIVE_AUDIO_TEXT_OUTPUT, + "output_cost_per_audio_token": NATIVE_AUDIO_AUDIO_OUTPUT, + } + for key in NATIVE_AUDIO_KEYS + }, +} +ALL_KEYS: Final = tuple(PUBLISHED_RATES) +LONG_CONTEXT_TIER_FIELDS: Final = ( + "input_cost_per_token_above_200k_tokens", + "output_cost_per_token_above_200k_tokens", + "cache_read_input_token_cost_above_200k_tokens", +) + + +def _load(path: Path) -> dict: + with open(path, encoding="utf-8") as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model", ALL_KEYS) +@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) +def test_published_rates_are_registered(model: str, path: Path): + info = _load(path)[model] + for field, value in PUBLISHED_RATES[model].items(): + assert info[field] == value, f"{model} {field} in {path.name}: {info.get(field)} != {value}" + + +@pytest.mark.parametrize("model", PRO_TTS_KEYS) +@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) +def test_pro_tts_has_no_long_context_tier(model: str, path: Path): + info = _load(path)[model] + for field in LONG_CONTEXT_TIER_FIELDS: + assert field not in info, f"{model} has {field} but Google publishes one flat TTS rate" + + +@pytest.mark.parametrize("model", ALL_KEYS) +def test_backup_matches_main(model: str): + assert _load(BACKUP_PATH)[model] == _load(MAIN_PATH)[model] + + +@pytest.mark.parametrize( + ("model", "provider", "input_rate", "audio_output_rate"), + ( + ("gemini-2.5-flash-preview-tts", "gemini", FLASH_TTS_INPUT, FLASH_TTS_AUDIO_OUTPUT), + ("gemini-2.5-pro-preview-tts", "gemini", PRO_TTS_INPUT, PRO_TTS_AUDIO_OUTPUT), + ("gemini-2.5-pro-preview-tts", "vertex_ai", PRO_TTS_INPUT, PRO_TTS_AUDIO_OUTPUT), + ), +) +def test_tts_audio_output_is_billed_at_the_audio_rate( + model: str, provider: str, input_rate: float, audio_output_rate: float, local_model_cost_map +): + usage: Final = Usage( + prompt_tokens=9, + completion_tokens=49, + total_tokens=58, + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=9), + completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=49, text_tokens=0), + ) + prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=provider) + assert prompt_cost == pytest.approx(9 * input_rate) + assert completion_cost == pytest.approx(49 * audio_output_rate) + + +@pytest.mark.parametrize("model", NATIVE_AUDIO_KEYS) +def test_native_audio_output_is_billed_at_the_audio_rate(model: str, local_model_cost_map): + usage: Final = Usage( + prompt_tokens=377, + completion_tokens=84, + total_tokens=461, + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=377), + completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=48, reasoning_tokens=36, text_tokens=0), + ) + prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="gemini") + assert prompt_cost == pytest.approx(377 * NATIVE_AUDIO_TEXT_INPUT) + assert completion_cost == pytest.approx(48 * NATIVE_AUDIO_AUDIO_OUTPUT + 36 * NATIVE_AUDIO_TEXT_OUTPUT) + + +@pytest.mark.parametrize("model", NATIVE_AUDIO_KEYS) +def test_native_audio_input_is_billed_at_the_audio_rate(model: str, local_model_cost_map): + usage: Final = Usage( + prompt_tokens=1000, + completion_tokens=0, + total_tokens=1000, + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100, audio_tokens=900), + ) + prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="gemini") + assert prompt_cost == pytest.approx(100 * NATIVE_AUDIO_TEXT_INPUT + 900 * NATIVE_AUDIO_AUDIO_INPUT) From 3576c773eb3103412e3147a9c49de90de49d8f39 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 26 Aug 2026 15:02:19 -0700 Subject: [PATCH 07/21] fix(proxy): serialize the search tool router refresh reload_search_tools_from_db is a read-modify-write of the shared llm_router global: it reads the whole table, merges the config tools in, and replaces router.search_tools wholesale. Two of those interleaving lets the older snapshot's assignment land last and put back a tool the newer one deleted, so a revoked tool keeps serving on the provider key it carried until the next reload. Take MODEL_RECONCILE_LOCK, which add_deployment already uses to serialize the same shape of work on the same global. It has to go on this entry point rather than in _init_search_tools_in_db, because _init_non_llm_objects_in_db calls that while already holding the lock and asyncio.Lock is not reentrant. A separate search-tools-only lock would not close the race: the periodic reconcile reaches _init_search_tools_in_db under MODEL_RECONCILE_LOCK, so only that same lock orders an endpoint refresh against a cron tick. Ordering across workers is unchanged and still reconciles on the next tick. --- litellm/proxy/proxy_server.py | 9 +++- .../proxy/proxy_server/test_proxy_config.py | 45 +++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7434ac3c494..18cc57cd135 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7491,12 +7491,19 @@ class ProxyConfig: Driven by the management endpoints so the worker that served the write is correct immediately, and by the periodic job in store_model_in_db-off deployments. Gated the same way as startup, so an admin who excluded search_tools from supported_db_objects opts out. + + Serialized by MODEL_RECONCILE_LOCK for the reason add_deployment documents: the body is a + read-modify-write of the shared ``llm_router`` global, so two of them interleaving lets the + older snapshot's wholesale assignment land last and restore a tool the newer one deleted. + The lock belongs here rather than in _init_search_tools_in_db, which _init_non_llm_objects_in_db + already calls while holding it. """ if not self._should_load_db_object(object_type="search_tools"): return if prisma_client is None: return - await self._init_search_tools_in_db(prisma_client=prisma_client) + async with MODEL_RECONCILE_LOCK: + await self._init_search_tools_in_db(prisma_client=prisma_client) @staticmethod def _merge_config_and_db_search_tools( diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 7678b0cab9e..d1dada4d10e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -1304,6 +1304,51 @@ async def test_ProxyConfig_reload_search_tools_from_db_honors_supported_db_objec mock_init.assert_not_awaited() +@pytest.mark.asyncio +async def test_ProxyConfig_reload_search_tools_from_db_serializes_overlapping_refreshes(monkeypatch): + """An older snapshot must not land last and restore a tool a newer refresh deleted.""" + import asyncio + + from litellm.proxy import proxy_server + + pc = ProxyConfig() + pc.update_config_state({}) + fake_router = MagicMock() + fake_router.search_tools = [] + + stale_read_started = asyncio.Event() + fresh_write_committed = asyncio.Event() + snapshots = iter( + ( + [{"search_tool_name": "doomed-search", "litellm_params": {}}], + [], + ) + ) + + async def _read_db(**_): + snapshot = next(snapshots) + if not stale_read_started.is_set(): + stale_read_started.set() + await fresh_write_committed.wait() + return snapshot + + monkeypatch.setattr(proxy_server, "llm_router", fake_router) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr( + "litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db", + _read_db, + ) + + stale = asyncio.create_task(pc.reload_search_tools_from_db()) + await stale_read_started.wait() + deleter = asyncio.create_task(pc.reload_search_tools_from_db()) + await asyncio.sleep(0) + fresh_write_committed.set() + await asyncio.gather(stale, deleter) + + assert fake_router.search_tools == [] + + @pytest.mark.asyncio async def test_ProxyConfig_reload_search_tools_from_db_noops_without_prisma(monkeypatch): from litellm.proxy import proxy_server From b687fe2b50ff6662d45d6b930ab0f8bfcca3bbcb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:03:32 -0700 Subject: [PATCH 08/21] fix(prompts): propagate PATCHed prompt templates to every worker and pod --- litellm/proxy/prompts/prompt_endpoints.py | 33 +++------ litellm/proxy/prompts/prompt_registry.py | 17 +++++ litellm/proxy/proxy_server.py | 2 +- .../prompts/test_prompt_endpoints_crud.py | 69 ++++++++++++++++++- .../proxy/prompts/test_prompt_registry.py | 67 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 48 +++++++++++++ 6 files changed, 207 insertions(+), 29 deletions(-) create mode 100644 tests/test_litellm/proxy/prompts/test_prompt_registry.py diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index a289ed7cbfb..cebfd5022a8 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -1025,15 +1025,8 @@ async def delete_prompt( raise HTTPException(status_code=500, detail=str(e)) -def _reload_prompt_in_registry( - registry: "InMemoryPromptRegistry", versioned_id: str, updated_prompt_spec: PromptSpec -) -> PromptSpec: - """Remove stale entry and re-initialize the prompt in the in-memory registry.""" - if versioned_id in registry.IN_MEMORY_PROMPTS: - del registry.IN_MEMORY_PROMPTS[versioned_id] - if versioned_id in registry.prompt_id_to_custom_prompt: - del registry.prompt_id_to_custom_prompt[versioned_id] - initialized: Final = registry.initialize_prompt(prompt=updated_prompt_spec, config_file_path=None) +def _reload_prompt_in_registry(registry: "InMemoryPromptRegistry", updated_prompt_spec: PromptSpec) -> PromptSpec: + initialized: Final = registry.reload_prompt(prompt=updated_prompt_spec) if initialized is None: raise HTTPException(status_code=500, detail="Failed to patch prompt") return initialized @@ -1123,25 +1116,15 @@ async def patch_prompt( detail="Cannot update config prompts.", ) - # Use existing prompt from memory or build from DB row for field merging - if existing_prompt: - current_litellm_params = existing_prompt.litellm_params - current_prompt_info = existing_prompt.prompt_info - else: - current_spec: Final = create_versioned_prompt_spec(db_prompt=target_row) - current_litellm_params = current_spec.litellm_params - current_prompt_info = current_spec.prompt_info + current_spec: Final = create_versioned_prompt_spec(db_prompt=target_row) - # Update fields if provided updated_litellm_params: Final = ( - request.litellm_params if request.litellm_params is not None else current_litellm_params + request.litellm_params if request.litellm_params is not None else current_spec.litellm_params ) - updated_prompt_info: Final = request.prompt_info if request.prompt_info is not None else current_prompt_info - - # Ensure we have valid litellm_params - if updated_litellm_params is None: - raise HTTPException(status_code=400, detail="litellm_params cannot be None") + updated_prompt_info: Final = ( + request.prompt_info if request.prompt_info is not None else current_spec.prompt_info + ) # Build update data dict update_data: Final[dict[str, str]] = { @@ -1165,7 +1148,7 @@ async def patch_prompt( updated_prompt_spec: Final = create_versioned_prompt_spec(db_prompt=updated_prompt_db_entry) - return _reload_prompt_in_registry(IN_MEMORY_PROMPT_REGISTRY, versioned_id, updated_prompt_spec) + return _reload_prompt_in_registry(IN_MEMORY_PROMPT_REGISTRY, updated_prompt_spec) except HTTPException as e: raise e diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index 695bdabfe83..ec7c98a068a 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -155,6 +155,23 @@ class InMemoryPromptRegistry: return parsed_prompt + def reload_prompt(self, prompt: PromptSpec) -> PromptSpec | None: + import litellm + + stale_callback: Final = self.prompt_id_to_custom_prompt.pop(prompt.prompt_id, None) + self.IN_MEMORY_PROMPTS.pop(prompt.prompt_id, None) + if stale_callback is not None: + litellm.logging_callback_manager.remove_callback_from_all_lists(stale_callback) + return self.initialize_prompt(prompt=prompt) + + def sync_prompt_from_db(self, prompt: PromptSpec) -> PromptSpec | None: + existing: Final = self.IN_MEMORY_PROMPTS.get(prompt.prompt_id) + if existing is None: + return self.initialize_prompt(prompt=prompt) + if existing.litellm_params == prompt.litellm_params and existing.prompt_info == prompt.prompt_info: + return existing + return self.reload_prompt(prompt=prompt) + def get_prompt_by_id(self, prompt_id: str) -> PromptSpec | None: """ Get a prompt by its ID from memory diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e55b254ab8e..8c08342f491 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7237,7 +7237,7 @@ class ProxyConfig: for prompt in prompts_in_db: # Convert DB object to dict and create versioned prompt_id prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) - IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=prompt_spec) + IN_MEMORY_PROMPT_REGISTRY.sync_prompt_from_db(prompt=prompt_spec) except Exception as e: verbose_proxy_logger.debug("litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - %s", e) diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py index 3e8e1e9dff8..c0be93b3dcc 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py @@ -1,3 +1,5 @@ +import json + import pytest from unittest.mock import MagicMock, AsyncMock, patch from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles @@ -8,6 +10,27 @@ from litellm.types.prompts.init_prompts import ( ) +def _db_row(content: str) -> MagicMock: + row = MagicMock() + row.id = "row-1" + row.version = 1 + row.model_dump.return_value = { + "prompt_id": "test_prompt", + "version": 1, + "environment": "development", + "created_by": None, + "litellm_params": { + "prompt_id": "test_prompt", + "prompt_integration": "dotprompt", + "prompt_data": {"content": content, "metadata": {}}, + }, + "prompt_info": {"prompt_type": "db"}, + "created_at": None, + "updated_at": None, + } + return row + + @pytest.mark.asyncio async def test_delete_prompt_success(): """ @@ -208,9 +231,7 @@ async def test_patch_prompt_row_deleted_mid_update_returns_404(): api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN ) - target_row = MagicMock() - target_row.id = "row-1" - target_row.version = 1 + target_row = _db_row("Begin every reply with AHOY") mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_prompttable.find_many = AsyncMock( @@ -246,3 +267,45 @@ async def test_patch_prompt_row_deleted_mid_update_returns_404(): exc_info.value.detail == "Prompt with ID test_prompt not found in environment development" ) + + +@pytest.mark.asyncio +async def test_patch_prompt_merges_unsent_fields_from_db_row_not_stale_memory(): + from litellm.proxy.prompts.prompt_endpoints import PatchPromptRequest, patch_prompt + + mock_user_auth = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN) + db_row = _db_row("Begin every reply with HOWDY") + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[db_row]) + mock_prisma_client.db.litellm_prompttable.update = AsyncMock(return_value=db_row) + stale_in_memory = PromptSpec( + prompt_id="test_prompt.v1", + litellm_params=PromptLiteLLMParams( + prompt_id="test_prompt", + prompt_integration="dotprompt", + prompt_data={"content": "Begin every reply with AHOY", "metadata": {}}, + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( # test-quality-ok: stubs the collaborator so the test pins what the endpoint writes and reloads + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" + ) as mock_registry, + ): + mock_registry.get_prompt_by_id.return_value = stale_in_memory + mock_registry.reload_prompt.side_effect = lambda prompt: prompt + + response = await patch_prompt( + prompt_id="test_prompt", + request=PatchPromptRequest(prompt_info=PromptInfo(prompt_type="db")), + user_api_key_dict=mock_user_auth, + ) + + written_params = json.loads(mock_prisma_client.db.litellm_prompttable.update.call_args.kwargs["data"]["litellm_params"]) + assert written_params["prompt_data"]["content"] == "Begin every reply with HOWDY" + reloaded_spec = mock_registry.reload_prompt.call_args.kwargs["prompt"] + assert reloaded_spec.prompt_id == "test_prompt.v1" + assert reloaded_spec.litellm_params.prompt_data["content"] == "Begin every reply with HOWDY" + assert response.litellm_params.prompt_data["content"] == "Begin every reply with HOWDY" diff --git a/tests/test_litellm/proxy/prompts/test_prompt_registry.py b/tests/test_litellm/proxy/prompts/test_prompt_registry.py new file mode 100644 index 00000000000..0533a6c11a8 --- /dev/null +++ b/tests/test_litellm/proxy/prompts/test_prompt_registry.py @@ -0,0 +1,67 @@ +import pytest + +import litellm +from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry +from litellm.types.prompts.init_prompts import PromptInfo, PromptLiteLLMParams, PromptSpec + + +def _db_prompt_spec(content: str) -> PromptSpec: + return PromptSpec( + prompt_id="greeting.v1", + litellm_params=PromptLiteLLMParams( + prompt_id="greeting", + prompt_integration="dotprompt", + prompt_data={"content": content, "metadata": {}}, + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + + +def _served_content(registry: InMemoryPromptRegistry) -> str: + callback = registry.get_prompt_callback_by_id("greeting.v1") + assert callback is not None + return callback.prompt_manager.get_prompt("greeting").content + + +@pytest.fixture +def isolated_callbacks(monkeypatch: pytest.MonkeyPatch) -> list: + monkeypatch.setattr(litellm, "callbacks", []) + return litellm.callbacks + + +def test_sync_prompt_from_db_reloads_row_edited_elsewhere(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY")) + stale_callback = registry.get_prompt_callback_by_id("greeting.v1") + assert _served_content(registry) == "begin every reply with AHOY" + + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY")) + + assert _served_content(registry) == "begin every reply with HOWDY" + assert registry.get_prompt_by_id("greeting.v1").litellm_params.prompt_data["content"] == "begin every reply with HOWDY" + assert stale_callback not in isolated_callbacks + assert isolated_callbacks == [registry.get_prompt_callback_by_id("greeting.v1")] + + +def test_sync_prompt_from_db_keeps_unchanged_row_in_place(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY")) + first_callback = registry.get_prompt_callback_by_id("greeting.v1") + + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY")) + + assert registry.get_prompt_callback_by_id("greeting.v1") is first_callback + assert isolated_callbacks == [first_callback] + + +def test_reload_prompt_replaces_callback_without_leaking_the_old_one(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY")) + stale_callback = registry.get_prompt_callback_by_id("greeting.v1") + + reloaded = registry.reload_prompt(prompt=_db_prompt_spec("begin every reply with HOWDY")) + + assert reloaded is not None + assert _served_content(registry) == "begin every reply with HOWDY" + assert stale_callback not in isolated_callbacks + assert len(isolated_callbacks) == 1 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index afc42e8db45..ec2b79908ec 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11195,6 +11195,54 @@ async def test_init_guardrails_in_db_snapshots_and_reconciles_under_guardrail_re assert not GUARDRAIL_RECONCILE_LOCK.locked() + +@pytest.mark.asyncio +async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeypatch): + from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr(litellm, "callbacks", []) + + def db_row(content: str) -> MagicMock: + row = MagicMock() + row.model_dump.return_value = { + "prompt_id": "greeting_sync", + "version": 1, + "environment": "development", + "created_by": None, + "litellm_params": json.dumps( + { + "prompt_id": "greeting_sync", + "prompt_integration": "dotprompt", + "prompt_data": {"content": content, "metadata": {}}, + } + ), + "prompt_info": json.dumps({"prompt_type": "db"}), + "created_at": None, + "updated_at": None, + } + return row + + def served_content() -> str: + callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_sync.v1") + assert callback is not None + return callback.prompt_manager.get_prompt("greeting_sync").content + + prisma_client = MagicMock() + try: + prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[db_row("Begin every reply with AHOY")]) + await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) + assert served_content() == "Begin every reply with AHOY" + + prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[db_row("Begin every reply with HOWDY")]) + await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) + + assert served_content() == "Begin every reply with HOWDY" + assert litellm.callbacks == [IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_sync.v1")] + finally: + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_sync") + + class TestEmbeddingsFailureHookRequestData: @pytest.mark.asyncio async def test_failure_hook_gets_post_setup_data_with_logging_obj(self): From d565860f607acc5cf61f84e5e2c36f86d9f2f734 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:15:05 -0700 Subject: [PATCH 09/21] fix(cost-map): correct gemini-live native-audio text input rate --- ...odel_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- .../test_gemini_tts_native_audio_pricing.py | 24 +++++++++++++------ 3 files changed, 21 insertions(+), 11 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 467acce9850..1e2b2c63ec1 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -20355,7 +20355,7 @@ "gemini-live-2.5-flash-preview-native-audio-09-2025": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, @@ -20399,7 +20399,7 @@ "gemini/gemini-live-2.5-flash-preview-native-audio-09-2025": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 65535, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 467acce9850..1e2b2c63ec1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -20355,7 +20355,7 @@ "gemini-live-2.5-flash-preview-native-audio-09-2025": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, @@ -20399,7 +20399,7 @@ "gemini/gemini-live-2.5-flash-preview-native-audio-09-2025": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 3e-07, + "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, "max_output_tokens": 65535, diff --git a/tests/test_litellm/test_gemini_tts_native_audio_pricing.py b/tests/test_litellm/test_gemini_tts_native_audio_pricing.py index 88fd14e436d..803e112d5c9 100644 --- a/tests/test_litellm/test_gemini_tts_native_audio_pricing.py +++ b/tests/test_litellm/test_gemini_tts_native_audio_pricing.py @@ -20,6 +20,11 @@ NATIVE_AUDIO_KEYS: Final = tuple( for suffix in ("latest", "preview-09-2025", "preview-12-2025") ) +LIVE_NATIVE_AUDIO_KEYS: Final = ( + "gemini-live-2.5-flash-preview-native-audio-09-2025", + "gemini/gemini-live-2.5-flash-preview-native-audio-09-2025", +) + FLASH_TTS_INPUT: Final = 5e-07 FLASH_TTS_AUDIO_OUTPUT: Final = 1e-05 PRO_TTS_INPUT: Final = 1e-06 @@ -45,10 +50,15 @@ PUBLISHED_RATES: Final = { "output_cost_per_token": NATIVE_AUDIO_TEXT_OUTPUT, "output_cost_per_audio_token": NATIVE_AUDIO_AUDIO_OUTPUT, } - for key in NATIVE_AUDIO_KEYS + for key in (*NATIVE_AUDIO_KEYS, *LIVE_NATIVE_AUDIO_KEYS) }, } ALL_KEYS: Final = tuple(PUBLISHED_RATES) +NATIVE_AUDIO_BILLING_CASES: Final = ( + *((key, "gemini") for key in NATIVE_AUDIO_KEYS), + ("gemini-live-2.5-flash-preview-native-audio-09-2025", "vertex_ai"), + ("gemini/gemini-live-2.5-flash-preview-native-audio-09-2025", "gemini"), +) LONG_CONTEXT_TIER_FIELDS: Final = ( "input_cost_per_token_above_200k_tokens", "output_cost_per_token_above_200k_tokens", @@ -118,8 +128,8 @@ def test_tts_audio_output_is_billed_at_the_audio_rate( assert completion_cost == pytest.approx(49 * audio_output_rate) -@pytest.mark.parametrize("model", NATIVE_AUDIO_KEYS) -def test_native_audio_output_is_billed_at_the_audio_rate(model: str, local_model_cost_map): +@pytest.mark.parametrize("model, provider", NATIVE_AUDIO_BILLING_CASES) +def test_native_audio_output_is_billed_at_the_audio_rate(model: str, provider: str, local_model_cost_map): usage: Final = Usage( prompt_tokens=377, completion_tokens=84, @@ -127,18 +137,18 @@ def test_native_audio_output_is_billed_at_the_audio_rate(model: str, local_model prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=377), completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=48, reasoning_tokens=36, text_tokens=0), ) - prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="gemini") + prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=provider) assert prompt_cost == pytest.approx(377 * NATIVE_AUDIO_TEXT_INPUT) assert completion_cost == pytest.approx(48 * NATIVE_AUDIO_AUDIO_OUTPUT + 36 * NATIVE_AUDIO_TEXT_OUTPUT) -@pytest.mark.parametrize("model", NATIVE_AUDIO_KEYS) -def test_native_audio_input_is_billed_at_the_audio_rate(model: str, local_model_cost_map): +@pytest.mark.parametrize("model, provider", NATIVE_AUDIO_BILLING_CASES) +def test_native_audio_input_is_billed_at_the_audio_rate(model: str, provider: str, local_model_cost_map): usage: Final = Usage( prompt_tokens=1000, completion_tokens=0, total_tokens=1000, prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100, audio_tokens=900), ) - prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="gemini") + prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=provider) assert prompt_cost == pytest.approx(100 * NATIVE_AUDIO_TEXT_INPUT + 900 * NATIVE_AUDIO_AUDIO_INPUT) From fb13b47ee59552bdd4fef11f1354085d6beae686 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:21:24 -0700 Subject: [PATCH 10/21] test: type the pricing test helpers --- tests/test_litellm/test_gemini_tts_native_audio_pricing.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/test_gemini_tts_native_audio_pricing.py b/tests/test_litellm/test_gemini_tts_native_audio_pricing.py index 803e112d5c9..3679c73aecd 100644 --- a/tests/test_litellm/test_gemini_tts_native_audio_pricing.py +++ b/tests/test_litellm/test_gemini_tts_native_audio_pricing.py @@ -1,4 +1,5 @@ import json +from collections.abc import Iterator from pathlib import Path from typing import Final @@ -66,13 +67,13 @@ LONG_CONTEXT_TIER_FIELDS: Final = ( ) -def _load(path: Path) -> dict: +def _load(path: Path) -> dict[str, dict[str, object]]: with open(path, encoding="utf-8") as f: return json.load(f) @pytest.fixture -def local_model_cost_map(monkeypatch): +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: original_model_cost = litellm.model_cost monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") From ce8d6f7d25ae52cf0ae1ea38d76cee94233ca982 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:31:45 -0700 Subject: [PATCH 11/21] fix(mcp): token refresh and M2M egress honor the admin-entered token URL --- litellm/proxy/_experimental/mcp_server/db.py | 4 +-- .../mcp_server/oauth2_token_cache.py | 13 +++++----- .../outbound_credentials/adapter.py | 4 +-- .../authz_code_refresher.py | 7 ++++-- .../outbound_credentials/test_adapter.py | 19 ++++++++++++++ .../test_authz_code_refresher.py | 25 +++++++++++++++++++ .../mcp_server/test_db_credentials.py | 25 +++++++++++++++++++ .../mcp_server/test_oauth2_token_cache.py | 19 ++++++++++++++ 8 files changed, 104 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 4aa08020527..cf74cbd187e 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1600,7 +1600,7 @@ async def refresh_user_oauth_token( ) -> OAuthCredentialPayload | None: """Attempt to refresh a per-user OAuth2 token using its stored refresh_token. - POSTs to ``server.token_url`` with ``grant_type=refresh_token``. + POSTs to ``server.effective_token_url`` with ``grant_type=refresh_token``. On success: persists the new credential via ``store_user_oauth_credential`` and returns the updated payload dict. @@ -1609,7 +1609,7 @@ async def refresh_user_oauth_token( stale credential and triggering re-authentication. """ refresh_token: Final[str | None] = cred.get("refresh_token") - token_url: Final[str | None] = getattr(server, "token_url", None) + token_url: Final[str | None] = getattr(server, "effective_token_url", None) or getattr(server, "token_url", None) server_id: Final[str] = getattr(server, "server_id", "") client_id: Final[str | None] = getattr(server, "client_id", None) client_secret: Final[str | None] = getattr(server, "client_secret", None) diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index c76c933c5b5..b3f1da51074 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -67,7 +67,7 @@ class MCPOAuth2TokenCache(InMemoryCache): rest of the identity rather than stored in a key.""" material: Final = "\x00".join( ( - server.token_url or "", + server.effective_token_url or "", server.client_id or "", server.client_secret or "", " ".join(server.scopes or ()), @@ -82,7 +82,7 @@ class MCPOAuth2TokenCache(InMemoryCache): @staticmethod def _has_client_credentials_config(server: "MCPServer") -> bool: - return bool(server.client_id and server.client_secret and server.token_url) + return bool(server.client_id and server.client_secret and server.effective_token_url) async def async_get_token(self, server: "MCPServer") -> str | None: """Return a valid access token, fetching or refreshing as needed. @@ -112,19 +112,20 @@ class MCPOAuth2TokenCache(InMemoryCache): return token async def _fetch_token(self, server: "MCPServer") -> tuple[str, int]: - """POST to ``token_url`` with ``grant_type=client_credentials``. + """POST to ``effective_token_url`` with ``grant_type=client_credentials``. Returns ``(access_token, ttl_seconds)`` where ttl accounts for the expiry buffer so the cache entry expires before the real token does. """ client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) - if not server.client_id or not server.client_secret or not server.token_url: + token_url: Final = server.effective_token_url + if not server.client_id or not server.client_secret or not token_url: raise ValueError( f"MCP server '{server.server_id}' missing required OAuth2 fields: " f"client_id={bool(server.client_id)}, " f"client_secret={bool(server.client_secret)}, " - f"token_url={bool(server.token_url)}" + f"token_url={bool(token_url)}" ) token_request: Final = build_upstream_oauth2_token_request( @@ -146,7 +147,7 @@ class MCPOAuth2TokenCache(InMemoryCache): ) try: - response: Final = await client.post(server.token_url, data=data, headers=token_request.headers or None) + response: Final = await client.post(token_url, data=data, headers=token_request.headers or None) response.raise_for_status() except httpx.HTTPStatusError as exc: raise ValueError( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index be8ec1b8eb3..98e239b1d1d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -142,7 +142,7 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec: config=ClientCredentialsConfig( client_id=server.client_id, client_secret=SecretStr(server.client_secret) if server.client_secret else None, - token_url=server.token_url, + token_url=server.effective_token_url, scopes=tuple(server.scopes or ()), audience=server.audience, upstream_resource=resolve_upstream_resource(server), @@ -163,7 +163,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None: normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is forwarded only when the operator set it; a missing one is omitted, not derived. """ - endpoint: Final = server.token_exchange_endpoint or server.token_url + endpoint: Final = server.token_exchange_endpoint or server.effective_token_url if not server.client_id or not server.client_secret: return None profile: Final[Literal["rfc8693", "entra_obo"]] = ( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py index 6ea5756d43d..92bd30694af 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/authz_code_refresher.py @@ -88,7 +88,10 @@ class AuthorizationCodeRefresher: if token.refresh_token is None: return None server: Final = self._server_lookup(server_id) - if server is None or not server.token_url: + if server is None: + return None + token_url: Final = server.effective_token_url + if not token_url: return None try: @@ -106,7 +109,7 @@ class AuthorizationCodeRefresher: "refresh_token": token.refresh_token, **token_request.body, } - body: Final = await self._token_endpoint(server.token_url, form, token_request.headers) + body: Final = await self._token_endpoint(token_url, form, token_request.headers) if body is None: return None access_token: Final = body.get("access_token") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index e336bdc80c2..0020dbf8d61 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -579,3 +579,22 @@ def test_id_jag_honors_explicit_subject_token_type(): def test_id_jag_half_configured_defers_to_v1(server): # A half-configured server must defer (None) rather than 500 at IdJagConfig construction. assert to_server_spec(server) is None + + +def test_client_credentials_uses_admin_entered_token_url_when_issuer_yield_empties_resolved(): + """A pinned issuer empties the resolved token_url while configured_token_url keeps the + admin-entered value; the M2M spec must carry it so egress can mint.""" + spec = to_server_spec( + _server( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + url="https://up.example.com/mcp", + token_url=None, + configured_token_url="https://idp.example.com/token", + client_id="cid", + client_secret="csec", + ) + ) + assert spec is not None + assert isinstance(spec.config, ClientCredentialsConfig) + assert spec.config.token_url == "https://idp.example.com/token" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py index ab414d1e8a4..bb2f2ff8b02 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_authz_code_refresher.py @@ -20,8 +20,10 @@ class _Server: upstream_resource=None, url=None, server_id="srv", + configured_token_url=None, ): self.token_url = token_url + self.configured_token_url = configured_token_url self.client_id = client_id self.client_secret = client_secret self.token_endpoint_auth_method = token_endpoint_auth_method @@ -29,6 +31,10 @@ class _Server: self.url = url self.server_id = server_id + @property + def effective_token_url(self): + return self.token_url or self.configured_token_url + def _lookup(server): return lambda server_id: server @@ -262,3 +268,22 @@ async def test_returned_scope_overrides_prior_when_present(): assert token is not None assert token.scopes == ("read",) # a present scope replaces the prior grant assert persisted[0][5] == ("read",) + + +@pytest.mark.asyncio +async def test_refresh_uses_admin_entered_token_url_when_issuer_yield_empties_resolved(): + """A pinned issuer empties the resolved token_url while configured_token_url keeps the + admin-entered value; the refresh grant must POST there instead of silently failing.""" + posted = [] + refresher = _refresher( + server=_Server(token_url=None, configured_token_url="https://idp.example.com/token"), + body={"access_token": "new-at", "expires_in": 3600}, + post_sink=posted, + ) + token = await refresher.refresh( + "alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt") + ) + + assert token is not None + assert token.access_token == "new-at" + assert posted[0][0] == "https://idp.example.com/token" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 50248e95ffa..4d9142ad4c5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -1337,3 +1337,28 @@ def test_mcp_oauth_token_identity_changes_when_only_upstream_resource_is_edited( assert mcp_oauth_token_identity(set_to_explicit) == mcp_oauth_token_identity( _identity_server(credentials={**creds, "upstream_resource": "api://audience-one"}) ) + + +@pytest.mark.asyncio +async def test_refresh_user_oauth_token_uses_admin_entered_token_url_when_issuer_yield_empties_resolved(monkeypatch): + """A pinned issuer empties the resolved token_url while configured_token_url keeps the + admin-entered value; the silent per-user refresh must POST there instead of bailing.""" + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv-1", + name="test", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + client_secret="csec", + token_url=None, + configured_token_url="https://idp.example.com/token", + ) + result, captured = await _run_refresh(monkeypatch, server) + + assert result is not None + assert captured["url"] == "https://idp.example.com/token" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py index 72589fd8b3e..b1aa16a30c0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py @@ -392,3 +392,22 @@ async def test_invalidate_clears_every_identity_for_a_server(): assert refetched == "tok-after-invalidate" assert mock_client.post.call_count == 3 + + +@pytest.mark.asyncio +async def test_m2m_mint_uses_admin_entered_token_url_when_issuer_yield_empties_resolved(): + """A pinned issuer empties the resolved token_url while configured_token_url keeps the + admin-entered value; the client_credentials mint must POST there instead of raising.""" + server = _server(token_url=None, configured_token_url="https://auth.example.com/token") + cache = MCPOAuth2TokenCache() + mock_client = AsyncMock() + mock_client.post.return_value = _token_response("m2m-token-configured") + + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client", + return_value=mock_client, + ): + result = await cache.async_get_token(server) + + assert result == "m2m-token-configured" + assert mock_client.post.call_args[0][0] == "https://auth.example.com/token" From 54b57575d716d8ba55932a60dca00cb21b0ce701 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:34:25 -0700 Subject: [PATCH 12/21] test: restore model cost map via monkeypatch --- .../test_gemini_tts_native_audio_pricing.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/tests/test_litellm/test_gemini_tts_native_audio_pricing.py b/tests/test_litellm/test_gemini_tts_native_audio_pricing.py index 3679c73aecd..28fc248d5b2 100644 --- a/tests/test_litellm/test_gemini_tts_native_audio_pricing.py +++ b/tests/test_litellm/test_gemini_tts_native_audio_pricing.py @@ -74,15 +74,11 @@ def _load(path: Path) -> dict[str, dict[str, object]]: @pytest.fixture def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: - original_model_cost = litellm.model_cost monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield litellm.get_model_info.cache_clear() - try: - yield - finally: - litellm.model_cost = original_model_cost - litellm.get_model_info.cache_clear() @pytest.mark.parametrize("model", ALL_KEYS) From 6df307fef86fd07a73c4ec85f97cea5b62afd068 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:49:31 -0700 Subject: [PATCH 13/21] fix(prompts): validate a prompt replacement before swapping and isolate per-row sync failures --- litellm/proxy/prompts/prompt_registry.py | 40 +++++++++++------- litellm/proxy/proxy_server.py | 12 ++++-- .../proxy/prompts/test_prompt_registry.py | 23 ++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 42 +++++++++++++++++++ 4 files changed, 98 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index 49d61ef70a2..d4342773a85 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -118,7 +118,16 @@ class InMemoryPromptRegistry: verbose_proxy_logger.debug("prompt_id already exists in IN_MEMORY_PROMPTS") return self.IN_MEMORY_PROMPTS[prompt_id] - custom_prompt_callback: CustomPromptManagement | None = None + parsed_prompt, custom_prompt_callback = self._build_prompt_callback(prompt=prompt) + litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback) + + # store references to the prompt in memory + self.IN_MEMORY_PROMPTS[prompt_id] = parsed_prompt + self.prompt_id_to_custom_prompt[prompt_id] = custom_prompt_callback + + return parsed_prompt + + def _build_prompt_callback(self, prompt: PromptSpec) -> tuple[PromptSpec, CustomPromptManagement]: litellm_params_data: Final = prompt.litellm_params verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data) @@ -132,17 +141,17 @@ class InMemoryPromptRegistry: raise ValueError("prompt_integration is required") initializer: Final = prompt_initializer_registry.get(prompt_integration) - - if initializer: - custom_prompt_callback = initializer(litellm_params, prompt) - if not isinstance(custom_prompt_callback, CustomPromptManagement): - raise ValueError(f"CustomPromptManagement is required, got {type(custom_prompt_callback)}") - litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback) - else: + if initializer is None: raise ValueError(f"Unsupported prompt: {prompt_integration}") + custom_prompt_callback: Final = initializer(litellm_params, prompt) + if not isinstance(custom_prompt_callback, CustomPromptManagement): + raise ValueError( # noqa: TRY004 # prompt endpoints map ValueError to HTTP 400; keep the existing contract + f"CustomPromptManagement is required, got {type(custom_prompt_callback)}" + ) + parsed_prompt: Final = PromptSpec( - prompt_id=prompt_id, + prompt_id=prompt.prompt_id, litellm_params=litellm_params, prompt_info=prompt.prompt_info or PromptInfo(prompt_type="config"), created_at=prompt.created_at, @@ -151,21 +160,20 @@ class InMemoryPromptRegistry: environment=prompt.environment, created_by=prompt.created_by, ) - - # store references to the prompt in memory - self.IN_MEMORY_PROMPTS[prompt_id] = parsed_prompt - self.prompt_id_to_custom_prompt[prompt_id] = custom_prompt_callback - - return parsed_prompt + return parsed_prompt, custom_prompt_callback def reload_prompt(self, prompt: PromptSpec) -> PromptSpec | None: import litellm + parsed_prompt, new_callback = self._build_prompt_callback(prompt=prompt) stale_callback: Final = self.prompt_id_to_custom_prompt.pop(prompt.prompt_id, None) self.IN_MEMORY_PROMPTS.pop(prompt.prompt_id, None) if stale_callback is not None: litellm.logging_callback_manager.remove_callback_from_all_lists(stale_callback) - return self.initialize_prompt(prompt=prompt) + litellm.logging_callback_manager.add_litellm_callback(new_callback) + self.IN_MEMORY_PROMPTS[prompt.prompt_id] = parsed_prompt + self.prompt_id_to_custom_prompt[prompt.prompt_id] = new_callback + return parsed_prompt def sync_prompt_from_db(self, prompt: PromptSpec) -> PromptSpec | None: existing: Final = self.IN_MEMORY_PROMPTS.get(prompt.prompt_id) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 88ef62ecb2d..afc29255d58 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7242,9 +7242,15 @@ class ProxyConfig: try: prompts_in_db: Final[Sequence[object]] = await PromptRepository(prisma_client).table.find_many() for prompt in prompts_in_db: - # Convert DB object to dict and create versioned prompt_id - prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) - IN_MEMORY_PROMPT_REGISTRY.sync_prompt_from_db(prompt=prompt_spec) + try: + prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) + IN_MEMORY_PROMPT_REGISTRY.sync_prompt_from_db(prompt=prompt_spec) + except Exception as prompt_sync_error: # noqa: BLE001 # one poisoned row must not block syncing the remaining prompts + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - failed to sync prompt %s: %s", + getattr(prompt, "prompt_id", None), + prompt_sync_error, + ) except Exception as e: verbose_proxy_logger.debug("litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - %s", e) diff --git a/tests/test_litellm/proxy/prompts/test_prompt_registry.py b/tests/test_litellm/proxy/prompts/test_prompt_registry.py index 0533a6c11a8..47f1ba13627 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_registry.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_registry.py @@ -65,3 +65,26 @@ def test_reload_prompt_replaces_callback_without_leaking_the_old_one(isolated_ca assert _served_content(registry) == "begin every reply with HOWDY" assert stale_callback not in isolated_callbacks assert len(isolated_callbacks) == 1 + + +def test_reload_prompt_keeps_the_old_template_when_the_replacement_fails(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY")) + old_callback = registry.get_prompt_callback_by_id("greeting.v1") + + broken = PromptSpec( + prompt_id="greeting.v1", + litellm_params=PromptLiteLLMParams( + prompt_id="greeting", + prompt_integration="does_not_exist", + prompt_data={"content": "begin every reply with HOWDY", "metadata": {}}, + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + + with pytest.raises(ValueError, match="Unsupported prompt"): + registry.reload_prompt(prompt=broken) + + assert registry.get_prompt_callback_by_id("greeting.v1") is old_callback + assert _served_content(registry) == "begin every reply with AHOY" + assert isolated_callbacks == [old_callback] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7eea7e0652b..42cff844513 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11327,6 +11327,48 @@ async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeyp IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_sync") +@pytest.mark.asyncio +async def test_init_prompts_in_db_syncs_remaining_rows_when_one_row_fails(monkeypatch): + from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr(litellm, "callbacks", []) + + def db_row(prompt_id: str, integration: str) -> MagicMock: + row = MagicMock() + row.model_dump.return_value = { + "prompt_id": prompt_id, + "version": 1, + "environment": "development", + "created_by": None, + "litellm_params": json.dumps( + { + "prompt_id": prompt_id, + "prompt_integration": integration, + "prompt_data": {"content": "Begin every reply with AHOY", "metadata": {}}, + } + ), + "prompt_info": json.dumps({"prompt_type": "db"}), + "created_at": None, + "updated_at": None, + } + return row + + prisma_client = MagicMock() + try: + prisma_client.db.litellm_prompttable.find_many = AsyncMock( + return_value=[db_row("broken_sync", "does_not_exist"), db_row("healthy_sync", "dotprompt")] + ) + await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) + + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id("broken_sync.v1") is None + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("healthy_sync.v1") is not None + assert litellm.callbacks == [IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("healthy_sync.v1")] + finally: + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("healthy_sync") + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("broken_sync") + + class TestEmbeddingsFailureHookRequestData: @pytest.mark.asyncio async def test_failure_hook_gets_post_setup_data_with_logging_obj(self): From ab1b7bf3b68850de00ff3686bb1edbbbd6482c1a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:53:37 -0700 Subject: [PATCH 14/21] fix(cost): price gemini-live-2.5-flash-native-audio realtime sessions The GA vertex model had no cost map entry, and the realtime cost handler accepted the router's price-less auto-registered deployment entry for the session.created model at zero-defaulted rates, so sessions billed 0.0 even when base_model pointed at the priced preview key. Adds the GA entry at its published rates and makes the handler fall through zero-defaulted candidates unless their cost map entry explicitly declares pricing. --- litellm/cost_calculator.py | 80 ++++++++++++---- ...odel_prices_and_context_window_backup.json | 43 +++++++++ model_prices_and_context_window.json | 43 +++++++++ tests/test_litellm/test_cost_calculator.py | 91 +++++++++++++++++++ 4 files changed, 239 insertions(+), 18 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 4cd292b2416..128060815f6 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2,6 +2,7 @@ ## File for 'response_cost' calculation in Logging import logging import time +from collections.abc import Sequence from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, Literal, cast @@ -2370,6 +2371,61 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed" +def _candidate_realtime_token_costs( + model_name: str, + combined_usage_object: Usage, + custom_llm_provider: str, + data_residency: str | None, +) -> tuple[float, float] | None: + try: + return generic_cost_per_token( + model=model_name, + usage=combined_usage_object, + custom_llm_provider=custom_llm_provider, + data_residency=data_residency, + ) + except Exception: + return None + + +def _cost_map_entry_declares_pricing(model_name: str, custom_llm_provider: str) -> bool: + entries: Final = ( + litellm.model_cost.get(model_name), + litellm.model_cost.get(f"{custom_llm_provider}/{model_name}"), + ) + return any(entry is not None and any("cost_per" in field for field in entry) for entry in entries) + + +def _first_priced_realtime_token_costs( + potential_model_names: Sequence[str | None], + combined_usage_object: Usage, + custom_llm_provider: str, + data_residency: str | None, +) -> tuple[float, float]: + candidate_costs: Final = ( + (model_name, costs) + for model_name in potential_model_names + if model_name is not None + and ( + costs := _candidate_realtime_token_costs( + model_name=model_name, + combined_usage_object=combined_usage_object, + custom_llm_provider=custom_llm_provider, + data_residency=data_residency, + ) + ) + is not None + ) + return next( + ( + costs + for model_name, costs in candidate_costs + if sum(costs) > 0 or _cost_map_entry_declares_pricing(model_name, custom_llm_provider) + ), + (0.0, 0.0), + ) + + def handle_realtime_stream_cost_calculation( results: OpenAIRealtimeStreamList, combined_usage_object: Usage, @@ -2394,24 +2450,12 @@ def handle_realtime_stream_cost_calculation( potential_model_names.append(received_model) potential_model_names.append(litellm_model_name) - input_cost_per_token = 0.0 - output_cost_per_token = 0.0 - - for model_name in potential_model_names: - try: - if model_name is None: - continue - _input_cost_per_token, _output_cost_per_token = generic_cost_per_token( - model=model_name, - usage=combined_usage_object, - custom_llm_provider=custom_llm_provider, - data_residency=data_residency, - ) - except Exception: - continue - input_cost_per_token += _input_cost_per_token - output_cost_per_token += _output_cost_per_token - break # exit if we find a valid model + input_cost_per_token, output_cost_per_token = _first_priced_realtime_token_costs( + potential_model_names=potential_model_names, + combined_usage_object=combined_usage_object, + custom_llm_provider=custom_llm_provider, + data_residency=data_residency, + ) transcription_cost: Final = ( handle_realtime_transcription_cost_calculation( results=results, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a6da2c1fb09..38098fa7082 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -20370,6 +20370,49 @@ }, "supports_image_size": false }, + "gemini-live-2.5-flash-native-audio": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_endpoints": [ + "/vertex_ai/live" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035, + "search_context_size_high": 0.035 + }, + "gemini_native_audio": true + }, "gemini-live-2.5-flash-preview-native-audio-09-2025": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 3e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a6da2c1fb09..38098fa7082 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -20370,6 +20370,49 @@ }, "supports_image_size": false }, + "gemini-live-2.5-flash-native-audio": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_endpoints": [ + "/vertex_ai/live" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035, + "search_context_size_high": 0.035 + }, + "gemini_native_audio": true + }, "gemini-live-2.5-flash-preview-native-audio-09-2025": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 3e-06, diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index dc2fbe3ed73..cfd07171556 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4116,3 +4116,94 @@ def test_every_one_hour_cache_write_rate_is_double_its_input_rate(): } assert deviations == {} + + +def test_gemini_live_native_audio_ga_realtime_cost(_local_model_cost_map: None) -> None: + """Regression for https://github.com/BerriAI/litellm/issues/31087: realtime sessions on the + GA vertex model gemini-live-2.5-flash-native-audio must bill at its published rates instead + of logging zero spend because only the preview-09-2025 key existed in the cost map.""" + from litellm.types.utils import CompletionTokensDetailsWrapper + + results: OpenAIRealtimeStreamList = [ + {"type": "session.created", "session": {"model": "gemini-live-2.5-flash-native-audio"}}, + ] + combined_usage_object = Usage( + prompt_tokens=8, + completion_tokens=25, + total_tokens=33, + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=8, audio_tokens=0), + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=2, audio_tokens=23), + ) + + cost = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=combined_usage_object, + custom_llm_provider="vertex_ai", + litellm_model_name="vertex_ai/gemini-live-2.5-flash-native-audio", + ) + + expected_cost = 8 * 5e-07 + 2 * 2e-06 + 23 * 1.2e-05 + assert cost == pytest.approx(expected_cost, rel=1e-9) + + +def test_realtime_priceless_deployment_entry_falls_through_to_priced_model( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """Regression for https://github.com/BerriAI/litellm/issues/31087: the router registers every + deployment's backend key into litellm.model_cost without price fields, and the realtime cost + handler used to accept that zero-defaulted entry for the session.created model and stop, so a + configured base_model never priced the session.""" + monkeypatch.setitem( + litellm.model_cost, + "vertex_ai/some-unmapped-live-model", + {"litellm_provider": "vertex_ai", "mode": "realtime"}, + ) + priced_model = "vertex_ai/gemini-live-2.5-flash-preview-native-audio-09-2025" + priced_entry = litellm.model_cost["gemini-live-2.5-flash-preview-native-audio-09-2025"] + + results: OpenAIRealtimeStreamList = [ + {"type": "session.created", "session": {"model": "some-unmapped-live-model"}}, + ] + combined_usage_object = Usage(prompt_tokens=8, completion_tokens=25, total_tokens=33) + + cost = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=combined_usage_object, + custom_llm_provider="vertex_ai", + litellm_model_name=priced_model, + ) + + expected_cost = 8 * priced_entry["input_cost_per_token"] + 25 * priced_entry["output_cost_per_token"] + assert cost == pytest.approx(expected_cost, rel=1e-9) + assert cost > 0 + + +def test_realtime_explicitly_free_session_model_still_bills_zero( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """A session model whose cost map entry explicitly declares zero rates is genuinely free, so + the handler must keep billing it at zero instead of falling through to a priced fallback.""" + monkeypatch.setitem( + litellm.model_cost, + "vertex_ai/free-live-model", + { + "litellm_provider": "vertex_ai", + "mode": "realtime", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + ) + + results: OpenAIRealtimeStreamList = [ + {"type": "session.created", "session": {"model": "free-live-model"}}, + ] + combined_usage_object = Usage(prompt_tokens=8, completion_tokens=25, total_tokens=33) + + cost = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=combined_usage_object, + custom_llm_provider="vertex_ai", + litellm_model_name="vertex_ai/gemini-live-2.5-flash-preview-native-audio-09-2025", + ) + + assert cost == 0.0 From 0243c5dee487f37cd7e9e1f7902bd1353e34eff6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:25:15 -0700 Subject: [PATCH 15/21] fix(model_prices): correct gemini-3.5-flash-lite flex cache-read pricing --- ...odel_prices_and_context_window_backup.json | 6 +-- model_prices_and_context_window.json | 6 +-- .../llm_cost_calc/test_llm_cost_calc_utils.py | 38 +++++++++++++++++++ 3 files changed, 44 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a6da2c1fb09..cead28e18aa 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -20147,7 +20147,7 @@ "gemini-3.5-flash-lite": { "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, - "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_flex": 1.5e-08, "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, @@ -22488,7 +22488,7 @@ }, "gemini/gemini-3.5-flash-lite": { "cache_read_input_token_cost": 3e-08, - "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_flex": 1.5e-08, "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, @@ -41987,7 +41987,7 @@ "vertex_ai/gemini-3.5-flash-lite": { "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, - "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_flex": 1.5e-08, "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a6da2c1fb09..cead28e18aa 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -20147,7 +20147,7 @@ "gemini-3.5-flash-lite": { "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, - "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_flex": 1.5e-08, "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, @@ -22488,7 +22488,7 @@ }, "gemini/gemini-3.5-flash-lite": { "cache_read_input_token_cost": 3e-08, - "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_flex": 1.5e-08, "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, @@ -41987,7 +41987,7 @@ "vertex_ai/gemini-3.5-flash-lite": { "deprecation_date": "2027-07-21", "cache_read_input_token_cost": 3e-08, - "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_flex": 1.5e-08, "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, 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 e13643ed6ce..18c3014c35d 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 @@ -3377,6 +3377,44 @@ def test_generic_cost_per_token_gemini_35_flash_lite(_local_model_cost_map): assert completion_cost == pytest.approx(0.00125) +GEMINI_35_FLASH_LITE_SERVICE_TIER_PRICING = [ + (None, 3e-07, 2.5e-06, 3e-08), + ("flex", 1.5e-07, 1.25e-06, 1.5e-08), + ("priority", 5.4e-07, 4.5e-06, 5e-08), +] + + +@pytest.mark.parametrize( + "service_tier,input_rate,output_rate,cache_read_rate", GEMINI_35_FLASH_LITE_SERVICE_TIER_PRICING +) +@pytest.mark.parametrize( + "model", + ["gemini-3.5-flash-lite", "gemini/gemini-3.5-flash-lite", "vertex_ai/gemini-3.5-flash-lite"], +) +def test_gemini_35_flash_lite_service_tier_pricing( + model, service_tier, input_rate, output_rate, cache_read_rate, _local_model_cost_map +): + """Regression: Vertex publishes flash-lite Flex/Batch context caching at $0.015/M + (1.5e-08/token), so flex cache reads must not be billed at the 2e-08 rate the map + used to carry.""" + usage = Usage( + prompt_tokens=1_000, + completion_tokens=500, + total_tokens=1_500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=200, text_tokens=800), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model=model.split("/")[-1], + usage=usage, + custom_llm_provider=model.split("/")[0] if "/" in model else "gemini", + service_tier=service_tier, + ) + + assert prompt_cost == pytest.approx(800 * input_rate + 200 * cache_read_rate, rel=1e-9) + assert completion_cost == pytest.approx(500 * output_rate, rel=1e-9) + + @pytest.mark.parametrize( "service_tier,input_rate,cache_read_rate,cache_write_rate,output_rate", [ From 93e7e8d98012ff4a217b743c803989fe91b8c2a5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:26:44 -0700 Subject: [PATCH 16/21] fix(mcp): token exchange rejoins discovery for a clientless DCR bridge missing its registration endpoint --- .../mcp_server/discoverable_endpoints.py | 14 ++++- .../mcp_server/test_discoverable_endpoints.py | 59 +++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 46feeab4dc3..93b85edd88d 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -967,7 +967,7 @@ async def exchange_token_with_server( if grant_type not in ("authorization_code", "refresh_token"): raise HTTPException(status_code=400, detail="Unsupported grant_type") - resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, lambda s: s.effective_token_url) + resolved_server: Final = await _server_with_oauth_endpoints(mcp_server, _token_flow_needed_endpoint) token_url: Final = resolved_server.effective_token_url if token_url is None: raise HTTPException( @@ -1664,6 +1664,18 @@ def _register_flow_needed_endpoint(mcp_server: MCPServer) -> str | None: return mcp_server.effective_authorization_url +def _token_flow_needed_endpoint(mcp_server: MCPServer) -> str | None: + """The token exchange's deferred-discovery join gate. The exchange's relay-vs-callback arm + (:func:`_dcr_bridge_relays_client_registration`) reads the registration url, so a clientless + DCR bridge rebuilt without its discovered registration endpoint must keep joining discovery + even when the token url already resolves; skipping it would select the gateway-callback arm + and the upstream would reject the code over a redirect_uri mismatch. Every other shape only + needs the token url.""" + if mcp_server.is_dcr_bridge and not mcp_server.client_id and mcp_server.effective_registration_url is None: + return None + return mcp_server.effective_token_url + + async def register_client_with_server( request: Request, mcp_server: MCPServer, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index b1a99b498e5..04fb17ee6aa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -247,6 +247,65 @@ async def test_register_route_bridge_missing_registration_url_joins_discovery(): assert json.loads(response.body.decode("utf-8"))["client_id"] == "generated-client" +@pytest.mark.asyncio +async def test_token_route_bridge_missing_registration_url_joins_discovery(): + """A clientless DCR bridge rebuilt with an admin-entered token url but without its discovered + registration endpoint must rejoin discovery at the exchange: the relay-vs-callback arm hinges + on the registration url, so skipping discovery would swap the client's own redirect_uri for + the gateway callback and the upstream would reject the code.""" + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="bridge-partial-token-metadata", + name="bridge_partial_token_metadata", + server_name="bridge_partial_token_metadata", + alias="bridge_partial_token_metadata", + url="https://mcp.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.true_passthrough, + dcr_bridge=True, + client_id=None, + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + ) + global_mcp_server_manager.registry[server.server_id] = server + global_mcp_server_manager._set_oauth_discovery_deferred(server.server_id, True) + request = _mock_callback_request("https://litellm.example.com/") + fake_http_response = MagicMock() + fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"} + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + + with ( + patch.object( # test-quality-ok: innermost discovery seam on a module-global manager; the route-to-exchange join under test stays real + global_mcp_server_manager, + "_discover_oauth_metadata_for_server", + new=AsyncMock(return_value=_resolved_oauth_metadata()), + ) as discovery, + patch.object( # test-quality-ok: keeps the upstream token POST off the network so its redirect_uri arm can be asserted + discoverable_endpoints, + "get_async_httpx_client", + new=lambda llm_provider: fake_http_client, + ), + ): + response = await discoverable_endpoints.token_endpoint( + request=request, + grant_type="authorization_code", + code="upstream-code", + redirect_uri="https://client.example.com/cb", + client_id="dcr-client-id", + mcp_server_name=server.server_name, + ) + + discovery.assert_awaited_once_with(server) + assert response.status_code == 200 + assert fake_http_client.post.await_args.kwargs["data"]["redirect_uri"] == "https://client.example.com/cb" + + @pytest.fixture def trust_xff(): """Force ``IPAddressUtils.is_request_from_trusted_proxy`` to True. From b48bff7b5444711678504ec069cd13ad4ac769be Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:27:28 -0700 Subject: [PATCH 17/21] fix(cost_calculator): require real values when detecting declared realtime pricing --- litellm/cost_calculator.py | 5 ++++- tests/test_litellm/test_cost_calculator.py | 25 +++++++++++++++++----- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 128060815f6..9ac2b425961 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2393,7 +2393,10 @@ def _cost_map_entry_declares_pricing(model_name: str, custom_llm_provider: str) litellm.model_cost.get(model_name), litellm.model_cost.get(f"{custom_llm_provider}/{model_name}"), ) - return any(entry is not None and any("cost_per" in field for field in entry) for entry in entries) + return any( + entry is not None and any("cost_per" in field and value is not None for field, value in entry.items()) + for entry in entries + ) def _first_priced_realtime_token_costs( diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index cfd07171556..6d292cb5f94 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4146,17 +4146,32 @@ def test_gemini_live_native_audio_ga_realtime_cost(_local_model_cost_map: None) assert cost == pytest.approx(expected_cost, rel=1e-9) +@pytest.mark.parametrize( + "priceless_entry", + [ + {"litellm_provider": "vertex_ai", "mode": "realtime"}, + { + "litellm_provider": "vertex_ai", + "mode": "realtime", + "input_cost_per_token": None, + "output_cost_per_token": None, + "input_cost_per_audio_token": None, + }, + ], + ids=["registered_without_price_fields", "registered_with_none_valued_price_fields"], +) def test_realtime_priceless_deployment_entry_falls_through_to_priced_model( - _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, priceless_entry: dict ) -> None: """Regression for https://github.com/BerriAI/litellm/issues/31087: the router registers every - deployment's backend key into litellm.model_cost without price fields, and the realtime cost - handler used to accept that zero-defaulted entry for the session.created model and stop, so a - configured base_model never priced the session.""" + deployment's backend key into litellm.model_cost without price fields (and merges a None-valued + ModelInfo skeleton into mapped entries), and the realtime cost handler used to accept that + zero-defaulted entry for the session.created model and stop, so a configured base_model never + priced the session.""" monkeypatch.setitem( litellm.model_cost, "vertex_ai/some-unmapped-live-model", - {"litellm_provider": "vertex_ai", "mode": "realtime"}, + priceless_entry, ) priced_model = "vertex_ai/gemini-live-2.5-flash-preview-native-audio-09-2025" priced_entry = litellm.model_cost["gemini-live-2.5-flash-preview-native-audio-09-2025"] From e7b843d69b1ed8d23228ef5012a5aba9a1164a0f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:33:02 -0700 Subject: [PATCH 18/21] test: trim realtime cost test docstrings to one line --- tests/test_litellm/test_cost_calculator.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 6d292cb5f94..3e127a23aa6 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4119,9 +4119,7 @@ def test_every_one_hour_cache_write_rate_is_double_its_input_rate(): def test_gemini_live_native_audio_ga_realtime_cost(_local_model_cost_map: None) -> None: - """Regression for https://github.com/BerriAI/litellm/issues/31087: realtime sessions on the - GA vertex model gemini-live-2.5-flash-native-audio must bill at its published rates instead - of logging zero spend because only the preview-09-2025 key existed in the cost map.""" + """Regression for https://github.com/BerriAI/litellm/issues/31087.""" from litellm.types.utils import CompletionTokensDetailsWrapper results: OpenAIRealtimeStreamList = [ @@ -4163,11 +4161,7 @@ def test_gemini_live_native_audio_ga_realtime_cost(_local_model_cost_map: None) def test_realtime_priceless_deployment_entry_falls_through_to_priced_model( _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, priceless_entry: dict ) -> None: - """Regression for https://github.com/BerriAI/litellm/issues/31087: the router registers every - deployment's backend key into litellm.model_cost without price fields (and merges a None-valued - ModelInfo skeleton into mapped entries), and the realtime cost handler used to accept that - zero-defaulted entry for the session.created model and stop, so a configured base_model never - priced the session.""" + """Regression for https://github.com/BerriAI/litellm/issues/31087 (router-registered priceless entries).""" monkeypatch.setitem( litellm.model_cost, "vertex_ai/some-unmapped-live-model", @@ -4196,8 +4190,6 @@ def test_realtime_priceless_deployment_entry_falls_through_to_priced_model( def test_realtime_explicitly_free_session_model_still_bills_zero( _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch ) -> None: - """A session model whose cost map entry explicitly declares zero rates is genuinely free, so - the handler must keep billing it at zero instead of falling through to a priced fallback.""" monkeypatch.setitem( litellm.model_cost, "vertex_ai/free-live-model", From bd75c38e84ea064866252024e3e9ecbce05ee668 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:48:32 -0700 Subject: [PATCH 19/21] fix(model_prices): scope flash-lite flex cache-read cut to vertex entries --- ...odel_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- .../llm_cost_calc/test_llm_cost_calc_utils.py | 39 ++++++++++++------- 3 files changed, 26 insertions(+), 17 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cead28e18aa..c6ba668a3c0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22488,7 +22488,7 @@ }, "gemini/gemini-3.5-flash-lite": { "cache_read_input_token_cost": 3e-08, - "cache_read_input_token_cost_flex": 1.5e-08, + "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cead28e18aa..c6ba668a3c0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -22488,7 +22488,7 @@ }, "gemini/gemini-3.5-flash-lite": { "cache_read_input_token_cost": 3e-08, - "cache_read_input_token_cost_flex": 1.5e-08, + "cache_read_input_token_cost_flex": 2e-08, "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, 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 18c3014c35d..d7f757b5852 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 @@ -3377,26 +3377,26 @@ def test_generic_cost_per_token_gemini_35_flash_lite(_local_model_cost_map): assert completion_cost == pytest.approx(0.00125) -GEMINI_35_FLASH_LITE_SERVICE_TIER_PRICING = [ - (None, 3e-07, 2.5e-06, 3e-08), - ("flex", 1.5e-07, 1.25e-06, 1.5e-08), - ("priority", 5.4e-07, 4.5e-06, 5e-08), +GEMINI_35_FLASH_LITE_TIER_RATES_BY_SURFACE = [ + ("gemini", None, 3e-07, 2.5e-06, 3e-08), + ("gemini", "flex", 1.5e-07, 1.25e-06, 2e-08), + ("gemini", "priority", 5.4e-07, 4.5e-06, 5e-08), + ("vertex_ai", None, 3e-07, 2.5e-06, 3e-08), + ("vertex_ai", "flex", 1.5e-07, 1.25e-06, 1.5e-08), + ("vertex_ai", "priority", 5.4e-07, 4.5e-06, 5e-08), ] @pytest.mark.parametrize( - "service_tier,input_rate,output_rate,cache_read_rate", GEMINI_35_FLASH_LITE_SERVICE_TIER_PRICING -) -@pytest.mark.parametrize( - "model", - ["gemini-3.5-flash-lite", "gemini/gemini-3.5-flash-lite", "vertex_ai/gemini-3.5-flash-lite"], + "custom_llm_provider,service_tier,input_rate,output_rate,cache_read_rate", + GEMINI_35_FLASH_LITE_TIER_RATES_BY_SURFACE, ) def test_gemini_35_flash_lite_service_tier_pricing( - model, service_tier, input_rate, output_rate, cache_read_rate, _local_model_cost_map + custom_llm_provider, service_tier, input_rate, output_rate, cache_read_rate, _local_model_cost_map ): - """Regression: Vertex publishes flash-lite Flex/Batch context caching at $0.015/M - (1.5e-08/token), so flex cache reads must not be billed at the 2e-08 rate the map - used to carry.""" + """Regression: Vertex publishes flash-lite flex context caching at $0.015/M while the + Gemini API publishes $0.02/M, so vertex_ai flex cache reads must bill 1.5e-08/token + instead of the 2e-08 the map used to carry, without disturbing the Gemini API rate.""" usage = Usage( prompt_tokens=1_000, completion_tokens=500, @@ -3405,9 +3405,9 @@ def test_gemini_35_flash_lite_service_tier_pricing( ) prompt_cost, completion_cost = generic_cost_per_token( - model=model.split("/")[-1], + model="gemini-3.5-flash-lite", usage=usage, - custom_llm_provider=model.split("/")[0] if "/" in model else "gemini", + custom_llm_provider=custom_llm_provider, service_tier=service_tier, ) @@ -3415,6 +3415,15 @@ def test_gemini_35_flash_lite_service_tier_pricing( assert completion_cost == pytest.approx(500 * output_rate, rel=1e-9) +def test_gemini_35_flash_lite_flex_cache_read_map_entries(_local_model_cost_map): + """Each map entry carries its own surface's published flex cache-read rate: the bare + and vertex_ai keys are the Vertex surface at $0.015/M, the gemini key is the Gemini + API surface at $0.02/M.""" + assert litellm.model_cost["gemini-3.5-flash-lite"]["cache_read_input_token_cost_flex"] == 1.5e-08 + assert litellm.model_cost["vertex_ai/gemini-3.5-flash-lite"]["cache_read_input_token_cost_flex"] == 1.5e-08 + assert litellm.model_cost["gemini/gemini-3.5-flash-lite"]["cache_read_input_token_cost_flex"] == 2e-08 + + @pytest.mark.parametrize( "service_tier,input_rate,cache_read_rate,cache_write_rate,output_rate", [ From 1e1c23107685efaa2a899274d02c2cc9d1014196 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:52:09 -0700 Subject: [PATCH 20/21] fix(model_prices): bill gemini -latest/preview alias cache reads at 10% of input --- ...odel_prices_and_context_window_backup.json | 12 ++--- model_prices_and_context_window.json | 12 ++--- .../llms/gemini/test_cost_calculator.py | 49 +++++++++++++++++++ 3 files changed, 61 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ad19818a5ce..7e3e1005bc0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -20332,7 +20332,7 @@ "supports_image_size": false }, "gemini-2.5-flash-preview-09-2025": { - "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", @@ -20469,7 +20469,7 @@ }, "gemini-2.5-flash-lite-preview-06-17": { "deprecation_date": "2025-11-18", - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -22062,7 +22062,7 @@ "supports_image_size": false }, "gemini/gemini-2.5-flash-preview-09-2025": { - "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost": 3e-08, "deprecation_date": "2026-02-17", "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -22111,7 +22111,7 @@ "supports_image_size": false }, "gemini/gemini-flash-latest": { - "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", @@ -22158,7 +22158,7 @@ "google_maps_grounding_cost_per_query": 0.025 }, "gemini/gemini-flash-lite-latest": { - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -22206,7 +22206,7 @@ }, "gemini/gemini-2.5-flash-lite-preview-06-17": { "deprecation_date": "2025-11-18", - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ad19818a5ce..7e3e1005bc0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -20332,7 +20332,7 @@ "supports_image_size": false }, "gemini-2.5-flash-preview-09-2025": { - "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-language-models", @@ -20469,7 +20469,7 @@ }, "gemini-2.5-flash-lite-preview-06-17": { "deprecation_date": "2025-11-18", - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-language-models", @@ -22062,7 +22062,7 @@ "supports_image_size": false }, "gemini/gemini-2.5-flash-preview-09-2025": { - "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost": 3e-08, "deprecation_date": "2026-02-17", "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -22111,7 +22111,7 @@ "supports_image_size": false }, "gemini/gemini-flash-latest": { - "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", @@ -22158,7 +22158,7 @@ "google_maps_grounding_cost_per_query": 0.025 }, "gemini/gemini-flash-lite-latest": { - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -22206,7 +22206,7 @@ }, "gemini/gemini-2.5-flash-lite-preview-06-17": { "deprecation_date": "2025-11-18", - "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost": 1e-08, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index ba503d635a0..2f6f9a556c0 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -361,3 +361,52 @@ def test_gemini_image_generation_cost_no_web_search_when_absent(monkeypatch): ) assert cost_zero == cost_none + + +@pytest.mark.parametrize( + "model,custom_llm_provider,expected_cache_read_cost", + [ + ("gemini/gemini-flash-latest", "gemini", 3e-08), + ("gemini/gemini-flash-lite-latest", "gemini", 1e-08), + ("gemini/gemini-2.5-flash-preview-09-2025", "gemini", 3e-08), + ("gemini/gemini-2.5-flash-lite-preview-06-17", "gemini", 1e-08), + ("vertex_ai/gemini-2.5-flash-preview-09-2025", "vertex_ai", 3e-08), + ("vertex_ai/gemini-2.5-flash-lite-preview-06-17", "vertex_ai", 1e-08), + ], +) +def test_flash_alias_cache_read_is_ten_percent_of_input( + monkeypatch, model, custom_llm_provider, expected_cache_read_cost +): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + + model_info = litellm.get_model_info( + model=model, custom_llm_provider=custom_llm_provider + ) + + assert model_info["cache_read_input_token_cost"] == expected_cache_read_cost + assert model_info["cache_read_input_token_cost"] == pytest.approx( + 0.10 * model_info["input_cost_per_token"] + ) + + +@pytest.mark.parametrize( + "prefixed,bare", + [ + ("gemini/gemini-flash-latest", "gemini-flash-latest"), + ("gemini/gemini-flash-lite-latest", "gemini-flash-lite-latest"), + ], +) +def test_flash_latest_alias_spellings_price_identically(monkeypatch, prefixed, bare): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + + prefixed_entry = litellm.model_cost[prefixed] + bare_entry = litellm.model_cost[bare] + + for cost_key in ( + "input_cost_per_token", + "output_cost_per_token", + "cache_read_input_token_cost", + ): + assert prefixed_entry[cost_key] == bare_entry[cost_key] From 5461bb3b48925a6a64e585c0be3ddb177b0ba707 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 17:00:24 -0700 Subject: [PATCH 21/21] fix(prompts): sync only the newest row when environments share a versioned prompt id --- litellm/proxy/proxy_server.py | 28 ++++++++-- tests/test_litellm/proxy/test_proxy_server.py | 51 +++++++++++++++++++ 2 files changed, 76 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index afc29255d58..b1ffce9c15a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7239,16 +7239,38 @@ class ProxyConfig: from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY from litellm.types.prompts.init_prompts import PromptSpec + def parse_row(db_prompt: object) -> PromptSpec | None: + try: + return self._get_prompt_spec_for_db_prompt(db_prompt=db_prompt) + except Exception as row_error: # noqa: BLE001 # a malformed row must not block syncing the remaining prompts + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - failed to parse prompt row %s: %s", + getattr(db_prompt, "prompt_id", None), + row_error, + ) + return None + try: prompts_in_db: Final[Sequence[object]] = await PromptRepository(prisma_client).table.find_many() - for prompt in prompts_in_db: + parsed_specs: Final[tuple[PromptSpec, ...]] = tuple( + spec for row in prompts_in_db if (spec := parse_row(row)) is not None + ) + newest_spec_per_id: Final[Mapping[str, PromptSpec]] = MappingProxyType( + { + spec.prompt_id: spec + for spec in sorted( + parsed_specs, + key=lambda s: s.updated_at.timestamp() if s.updated_at else float("-inf"), + ) + } + ) + for prompt_spec in newest_spec_per_id.values(): try: - prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) IN_MEMORY_PROMPT_REGISTRY.sync_prompt_from_db(prompt=prompt_spec) except Exception as prompt_sync_error: # noqa: BLE001 # one poisoned row must not block syncing the remaining prompts verbose_proxy_logger.exception( "litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - failed to sync prompt %s: %s", - getattr(prompt, "prompt_id", None), + prompt_spec.prompt_id, prompt_sync_error, ) except Exception as e: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 42cff844513..fbf71829abc 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11369,6 +11369,57 @@ async def test_init_prompts_in_db_syncs_remaining_rows_when_one_row_fails(monkey IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("broken_sync") +@pytest.mark.asyncio +async def test_init_prompts_in_db_serves_the_newest_row_when_environments_collide_on_a_versioned_id(monkeypatch): + from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr(litellm, "callbacks", []) + + def db_row(environment: str, content: str, updated_at: datetime) -> MagicMock: + row = MagicMock() + row.model_dump.return_value = { + "prompt_id": "greeting_env", + "version": 1, + "environment": environment, + "created_by": None, + "litellm_params": json.dumps( + { + "prompt_id": "greeting_env", + "prompt_integration": "dotprompt", + "prompt_data": {"content": content, "metadata": {}}, + } + ), + "prompt_info": json.dumps({"prompt_type": "db"}), + "created_at": None, + "updated_at": updated_at, + } + return row + + freshly_patched = db_row( + "production", "Begin every reply with HOWDY", datetime(2026, 8, 26, 12, 0, tzinfo=timezone.utc) + ) + stale_sibling = db_row( + "development", "Begin every reply with AHOY", datetime(2026, 8, 26, 11, 0, tzinfo=timezone.utc) + ) + + prisma_client = MagicMock() + try: + prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[freshly_patched, stale_sibling]) + await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) + + first_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_env.v1") + assert first_callback is not None + assert first_callback.prompt_manager.get_prompt("greeting_env").content == "Begin every reply with HOWDY" + + await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) + + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_env.v1") is first_callback + assert litellm.callbacks == [first_callback] + finally: + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_env") + + class TestEmbeddingsFailureHookRequestData: @pytest.mark.asyncio async def test_failure_hook_gets_post_setup_data_with_logging_obj(self):