From 48572c9516abc0585c3133aa02254e4b06706f01 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 20 Jul 2026 19:20:34 -0700 Subject: [PATCH 01/12] fix(mcp): attach resolved OAuth credentials to OpenAPI spec_path tool calls --- .../mcp_server/mcp_server_manager.py | 88 ++++++++++++++++++- .../mcp_server/openapi_to_mcp_generator.py | 25 +++++- .../proxy/_experimental/mcp_server/server.py | 14 +++ .../mcp_server/test_mcp_hook_extra_headers.py | 86 ++++++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 45 ++++++++++ .../test_openapi_to_mcp_generator.py | 59 +++++++++++++ .../mcp_server/test_openapi_tool_auth.py | 83 +++++++++++++++++ 7 files changed, 396 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 8f30071eb5d..3e591083ba2 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -610,6 +610,34 @@ def _passthrough_token_from_mcp_auth_header( return None +async def _materialize_auth_headers(auth: httpx.Auth | None) -> dict[str, str] | None: + """Extract the header a resolved ``httpx.Auth`` would set, as a plain dict, or None. + + OpenAPI tool closures egress through ``AsyncHTTPHandler`` methods that accept headers but no + ``auth``, so a resolved credential must be materialized into a header value. Driving one step + of the auth's own flow (against a throwaway request that is never sent) keeps this generic + across every auth shape without per-class branching; ``header_name`` is the resolver-arm + convention for "this auth sets a header" (``NoOpAuth`` has none and yields nothing to apply). + The materialized value is point-in-time: flow behaviors past the first request, like the M2M + one-shot 401 refetch, do not apply on this arm. + """ + if auth is None: + return None + header_name = getattr(auth, "header_name", None) + if not isinstance(header_name, str) or not header_name: + return None + probe = httpx.Request("GET", "http://localhost/") + flow = auth.async_auth_flow(probe) + try: + first_request = await flow.__anext__() + except StopAsyncIteration: + return None + finally: + await flow.aclose() + header_value = first_request.headers.get(header_name) + return {header_name: header_value} if header_value else None + + def _consumes_caller_authorization(server: MCPServer) -> bool: """True when this server's egress forwards the caller's request-wide ``Authorization`` upstream: the client-forwarded token modes, legacy OAuth pass-through, and legacy upstream-delegated @@ -4705,6 +4733,54 @@ class MCPServerManager: ) return oauth2_headers + async def resolve_openapi_upstream_auth( + self, + *, + mcp_server: MCPServer, + oauth2_headers: dict[str, str] | None, + raw_headers: dict[str, str] | None, + mcp_auth_header: str | dict[str, str] | None, + user_api_key_auth: UserAPIKeyAuth | None, + forwarded_headers: dict[str, str] | None, + ) -> tuple[dict[str, str] | None, dict[str, str] | None]: + """Resolve the gateway-owned upstream credential for a spec_path (OpenAPI) tool call. + + OpenAPI tools egress through a plain httpx call assembled from ContextVars, never through + ``_create_mcp_client``, so the v2 resolver graft there does not run for them and a resolved + credential (authorization_code's stored per-user token, client_credentials' minted M2M + token, token_exchange's exchanged token, passthrough's forwarded caller token) must be + materialized into headers here. Returns ``(resolved_auth_headers, forwarded_headers)``: + the resolved headers are authoritative over every other Authorization source (the same + rule ``_resolve_v2_auth`` applies on the MCPClient path) and ``forwarded_headers`` comes + back with any header the resolver claimed already dropped. Unmigrated (v1) servers resolve + through the stored-token lookup instead, and a missing per-user credential raises the same + discovery challenge the MCPClient path serves, rather than egressing unauthenticated. + """ + spec = to_server_spec(mcp_server) + if spec is None: + stored_headers = await self._resolve_oauth2_headers_for_tool_call( + mcp_server, oauth2_headers, user_api_key_auth + ) + return stored_headers, forwarded_headers + + subject_token: str | None = None + if isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)): + subject_token = self._extract_bearer_token(oauth2_headers, raw_headers) + elif isinstance(spec.config, PassthroughConfig): + inbound_token, forwarded_headers = _take_forwarded_authorization(forwarded_headers) + per_server_token = _passthrough_token_from_mcp_auth_header(mcp_auth_header) + subject_token = per_server_token if per_server_token is not None else inbound_token + + resolved_auth, forwarded_headers = await self._resolve_v2_auth( + server=mcp_server, + spec=spec, + provider=self._cred_provider, + subject_token=subject_token, + user_api_key_auth=user_api_key_auth, + extra_headers=forwarded_headers, + ) + return await _materialize_auth_headers(resolved_auth), forwarded_headers + async def _gather_openapi_tool_tasks( self, tasks: list[Any], @@ -4813,22 +4889,32 @@ class MCPServerManager: auth_header_value = ( _format_byok_openapi_auth_header(mcp_server, mcp_auth_header) if mcp_auth_header else None ) - forwarded_headers = _openapi_forwarded_extra_headers(mcp_server, raw_headers, user_api_key_auth) + resolved_auth_headers, forwarded_headers = await self.resolve_openapi_upstream_auth( + mcp_server=mcp_server, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + mcp_auth_header=mcp_auth_header, + user_api_key_auth=user_api_key_auth, + forwarded_headers=_openapi_forwarded_extra_headers(mcp_server, raw_headers, user_api_key_auth), + ) async def _call_openapi_via_handler(): from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, _request_extra_headers, + _request_resolved_auth_headers, ) auth_token = _request_auth_header.set(auth_header_value) extra_token = _request_extra_headers.set(forwarded_headers) + resolved_token = _request_resolved_auth_headers.set(resolved_auth_headers) try: async with self._limit_outbound_concurrency(mcp_server): return await self._call_openapi_tool_handler(mcp_server, name, arguments) finally: _request_auth_header.reset(auth_token) _request_extra_headers.reset(extra_token) + _request_resolved_auth_headers.reset(resolved_token) tasks.append(asyncio.create_task(_call_openapi_via_handler())) else: diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 1ee300be718..0b795057837 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -62,6 +62,14 @@ _request_extra_headers: contextvars.ContextVar[Optional[Dict[str, str]]] = conte "_request_extra_headers", default=None ) +# Per-request headers carrying the gateway-resolved upstream credential +# (stored per-user OAuth token, minted M2M token, exchanged OBO token). +# Set from MCPServerManager.resolve_openapi_upstream_auth; authoritative +# over every other Authorization source in _merge_openapi_tool_request_headers. +_request_resolved_auth_headers: contextvars.ContextVar[dict[str, str] | None] = contextvars.ContextVar( + "_request_resolved_auth_headers", default=None +) + def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str: """Ensure path params cannot introduce directory traversal.""" @@ -294,10 +302,15 @@ def _merge_openapi_tool_request_headers( """Merge static closure headers with per-request ContextVar overrides. Precedence (highest to lowest): - 1. ``_request_auth_header`` — BYOK override of ``Authorization`` - 2. ``static_headers`` — operator-configured headers baked into the + 1. ``_request_resolved_auth_headers`` — the gateway-resolved upstream + credential (stored per-user OAuth token, minted M2M token, + exchanged OBO token). The resolver is authoritative: a BYOK or + forwarded ``Authorization`` must not shadow it, mirroring + ``_resolve_v2_auth`` on the MCPClient path + 2. ``_request_auth_header`` — BYOK override of ``Authorization`` + 3. ``static_headers`` — operator-configured headers baked into the tool closure at registration time - 3. ``_request_extra_headers`` — per-request headers forwarded from + 4. ``_request_extra_headers`` — per-request headers forwarded from the MCP caller (allowlisted by ``MCPServer.extra_headers``) This matches the existing MCP invariant in @@ -323,6 +336,12 @@ def _merge_openapi_tool_request_headers( del effective_headers[existing] effective_headers["Authorization"] = override_auth + resolved_auth_headers = _request_resolved_auth_headers.get() or {} + for name, value in resolved_auth_headers.items(): + for existing in [k for k in effective_headers if k.lower() == name.lower()]: + del effective_headers[existing] + effective_headers[name] = value + return effective_headers diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a8ab0937124..75faca3c914 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -376,6 +376,7 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, _request_extra_headers, + _request_resolved_auth_headers, ) from litellm.proxy._experimental.mcp_server.sse_transport import SseServerTransport from litellm.proxy._experimental.mcp_server.tool_registry import ( @@ -2785,13 +2786,26 @@ if MCP_AVAILABLE: forwarded_headers = {} forwarded_headers[header_name] = value + resolved_auth_headers: dict[str, str] | None = None + if mcp_server: + resolved_auth_headers, forwarded_headers = await global_mcp_server_manager.resolve_openapi_upstream_auth( + mcp_server=mcp_server, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + mcp_auth_header=mcp_auth_header, + user_api_key_auth=user_api_key_auth, + forwarded_headers=forwarded_headers, + ) + _auth_token = _request_auth_header.set(auth_header_value) _extra_token = _request_extra_headers.set(forwarded_headers) + _resolved_token = _request_resolved_auth_headers.set(resolved_auth_headers) try: local_content = await _handle_local_mcp_tool(name, arguments) finally: _request_auth_header.reset(_auth_token) _request_extra_headers.reset(_extra_token) + _request_resolved_auth_headers.reset(_resolved_token) response = CallToolResult(content=cast(Any, local_content), isError=False) # Try managed MCP server tool (pass the full prefixed name) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 73486fe0b6a..ca5b7914cca 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1033,3 +1033,89 @@ class TestResolveByokMcpAuthHeader: check_mock.assert_awaited_once_with(server, user_auth) assert result == "caller-header" + + +class TestOpenApiResolvedUpstreamAuth: + """LIT-4629: spec_path servers egress through plain httpx, so the manager's OpenAPI arm must + materialize the v2-resolved credential into the `_request_resolved_auth_headers` ContextVar; + before the fix the resolved token never reached the upstream API.""" + + def _oauth_server(self, **overrides: Any) -> MCPServer: + fields: Dict[str, Any] = dict( + server_id="srv-sheets", + name="google_sheets", + server_name="google_sheets", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + spec_path="https://example.com/sheets-openapi.yaml", + ) + fields.update(overrides) + return MCPServer(**fields) + + @pytest.mark.asyncio + async def test_call_tool_openapi_injects_v2_resolved_token_contextvar(self): + """The managed spec_path arm resolves the v2 credential and sets the ContextVar; kills + the mutant that drops the resolve_openapi_upstream_auth call in call_tool.""" + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _request_resolved_auth_headers, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + + manager = MCPServerManager() + server = self._oauth_server() + user_auth = UserAPIKeyAuth(user_id="alice", api_key="sk-user") + captured: Dict[str, Any] = {} + + async def fake_openapi_handler(_server, _name, _arguments): + captured["resolved"] = _request_resolved_auth_headers.get() + return MagicMock() + + with patch.object(manager, "_resolve_mcp_server_for_tool_call", return_value=server): + with patch.object( + manager._cred_provider, + "resolve_credentials", + new=AsyncMock(return_value=Ok(StaticHeaderAuth("Bearer stored-user-token"))), + ): + with patch.object(manager, "_call_openapi_tool_handler", side_effect=fake_openapi_handler): + await manager.call_tool( + server_name=server.server_name, + name="get_values", + arguments={}, + user_api_key_auth=user_auth, + ) + + assert captured["resolved"] == {"Authorization": "Bearer stored-user-token"} + assert _request_resolved_auth_headers.get() is None + + @pytest.mark.asyncio + async def test_call_tool_openapi_m2m_missing_token_url_fails_closed(self): + """A url-less M2M spec server with no token_url must fail with a typed error instead of + egressing unauthenticated (the pre-#32259 silent failure this arm previously preserved). + Drives the real adapter/resolver chain: ClientCredentialsConfig with missing grant fields + resolves to a misconfigured CredError, raised as an HTTPException.""" + from fastapi import HTTPException + + manager = MCPServerManager() + server = self._oauth_server( + oauth2_flow="client_credentials", + client_id="m2m-client", + client_secret="m2m-secret", + token_url=None, + ) + called = AsyncMock() + + with patch.object(manager, "_resolve_mcp_server_for_tool_call", return_value=server): + with patch.object(manager, "_call_openapi_tool_handler", new=called): + with pytest.raises(HTTPException): + await manager.call_tool( + server_name=server.server_name, + name="get_values", + arguments={}, + user_api_key_auth=UserAPIKeyAuth(user_id="alice", api_key="sk-user"), + ) + + called.assert_not_awaited() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 03f91260955..576b9f4f139 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -8891,3 +8891,48 @@ async def test_resolve_toolset_tool_permissions_single_db_fetch_across_checks(): assert first == {"server-a": ["lookup_status"]} assert second == first list_toolsets_mock.assert_awaited_once() + + +class TestMaterializeAuthHeaders: + """_materialize_auth_headers drives one step of a resolved httpx.Auth's own flow to turn it + into a header dict for the OpenAPI egress arm, which sends plain headers and cannot carry an + httpx.Auth. Generic across auth shapes via the resolver-arm header_name convention.""" + + @pytest.mark.asyncio + async def test_static_header_auth_materializes_its_header(self): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _materialize_auth_headers, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + + headers = await _materialize_auth_headers(StaticHeaderAuth("Bearer stored-token")) + assert headers == {"Authorization": "Bearer stored-token"} + + @pytest.mark.asyncio + async def test_client_credentials_bearer_auth_materializes_bearer(self): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _materialize_auth_headers, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( + ClientCredentialsBearerAuth, + ) + + async def _refetch(_stale: str): + return None + + headers = await _materialize_auth_headers(ClientCredentialsBearerAuth("m2m-token", _refetch)) + assert headers == {"Authorization": "Bearer m2m-token"} + + @pytest.mark.asyncio + async def test_noop_and_none_materialize_to_none(self): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + _materialize_auth_headers, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + NoOpAuth, + ) + + assert await _materialize_auth_headers(None) is None + assert await _materialize_auth_headers(NoOpAuth()) is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 39f3c767220..7bcacb3ff4a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -17,6 +17,7 @@ import pytest from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, _request_extra_headers, + _request_resolved_auth_headers, _resolve_param_list, _resolve_ref, build_input_schema, @@ -1207,3 +1208,61 @@ class TestRequestExtraHeaders: call_args = async_client.get.call_args headers_sent = call_args[1]["headers"] assert "X-TOKEN" not in headers_sent + + @pytest.mark.asyncio + async def test_resolved_auth_headers_win_over_every_other_authorization_source(self): + """The gateway-resolved credential (stored per-user OAuth / minted M2M token) is + authoritative: it must override the BYOK override, static headers, and forwarded caller + headers on the Authorization name, case-insensitively, mirroring _resolve_v2_auth's rule + on the MCPClient path. Without this, a spec_path oauth2 server's completed OAuth flow + stores a token that never reaches the upstream API (LIT-4629).""" + operation = {} + func = create_tool_function( + path="/secure", + method="get", + operation=operation, + base_url="https://api.example.com", + headers={"authorization": "Bearer static-operator"}, + ) + + with patch(GET_ASYNC_CLIENT_TARGET) as mock_client: + async_client = _create_mock_client("get", "secure-data") + mock_client.return_value = async_client + + extra_token = _request_extra_headers.set({"Authorization": "Bearer caller-forwarded"}) + auth_token = _request_auth_header.set("Bearer byok-credential") + resolved_token = _request_resolved_auth_headers.set({"Authorization": "Bearer resolved-oauth"}) + try: + result = await func() + finally: + _request_auth_header.reset(auth_token) + _request_extra_headers.reset(extra_token) + _request_resolved_auth_headers.reset(resolved_token) + + assert result == "secure-data" + headers_sent = async_client.get.call_args[1]["headers"] + authorization_values = [v for k, v in headers_sent.items() if k.lower() == "authorization"] + assert authorization_values == ["Bearer resolved-oauth"] + + @pytest.mark.asyncio + async def test_resolved_auth_headers_not_leaked_between_calls(self): + """After resetting the resolved-auth ContextVar, subsequent calls send no credential.""" + operation = {} + func = create_tool_function( + path="/data", + method="get", + operation=operation, + base_url="https://api.example.com", + ) + + with patch(GET_ASYNC_CLIENT_TARGET) as mock_client: + async_client = _create_mock_client("get", "ok") + mock_client.return_value = async_client + + token = _request_resolved_auth_headers.set({"Authorization": "Bearer resolved-oauth"}) + _request_resolved_auth_headers.reset(token) + + await func() + + headers_sent = async_client.get.call_args[1]["headers"] + assert "Authorization" not in headers_sent diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py index 3ad01e9c3ec..1e4349c3143 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -218,3 +218,86 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable(): assert exc.value.status_code == 503 pre_call.assert_not_awaited() handle_local.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_openapi_local_tool_injects_resolved_oauth_token(): + """LIT-4629: the local-registry (OpenAPI) dispatch is the primary egress for spec_path + tools, and before the fix it dropped the gateway-resolved OAuth credential entirely, so a + user's completed OAuth flow stored a token that never reached the upstream API. The resolved + credential must land in the `_request_resolved_auth_headers` ContextVar the tool closure + reads. Kills the mutant that deletes the resolve_openapi_upstream_auth call in server.py.""" + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _request_resolved_auth_headers, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + from litellm.types.mcp import MCPAuth, MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + user = UserAPIKeyAuth( + api_key="sk-user", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + oauth_server = MCPServer( + server_id="srv-sheets", + name="google_sheets", + server_name="google_sheets", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + spec_path="https://example.com/sheets-openapi.yaml", + ) + + fake_tool = MagicMock() + fake_tool.name = "get_values" + captured: dict = {} + + async def handle_local(_name, _arguments): + captured["resolved"] = _request_resolved_auth_headers.get() + return [] + + with ( + patch.object( + mcp_module.global_mcp_server_manager, + "_get_mcp_server_from_tool_name", + return_value=oauth_server, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "pre_call_tool_check", + new=AsyncMock(return_value={}), + ), + patch.object( + mcp_module.global_mcp_tool_registry, + "get_tool", + return_value=fake_tool, + ), + patch.object( + mcp_module.global_mcp_server_manager._cred_provider, + "resolve_credentials", + new=AsyncMock(return_value=Ok(StaticHeaderAuth("Bearer stored-user-token"))), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + new=handle_local, + ), + patch( + "litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed", + return_value=True, + ), + ): + await mcp_module.execute_mcp_tool( + name="get_values", + arguments={}, + allowed_mcp_servers=[oauth_server], + start_time=datetime.now(timezone.utc), + user_api_key_auth=user, + ) + + assert captured["resolved"] == {"Authorization": "Bearer stored-user-token"} + assert _request_resolved_auth_headers.get() is None From 040aa9d8960fef0f1cf2b3f4b56fa2c9592cb690 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 20 Jul 2026 19:33:46 -0700 Subject: [PATCH 02/12] fix(mcp): never promote caller oauth2 headers to the resolved credential on the v1 arm --- .../mcp_server/mcp_server_manager.py | 16 +++- .../mcp_server/test_mcp_hook_extra_headers.py | 77 +++++++++++++++++++ 2 files changed, 89 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 3e591083ba2..ff714b8de22 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -4755,12 +4755,19 @@ class MCPServerManager: back with any header the resolver claimed already dropped. Unmigrated (v1) servers resolve through the stored-token lookup instead, and a missing per-user credential raises the same discovery challenge the MCPClient path serves, rather than egressing unauthenticated. + + The resolved headers carry only credentials the gateway itself resolved (a stored per-user + token, a minted or exchanged token). Caller-supplied ``oauth2_headers`` are never promoted + into them: on the v2 arm they feed only subject-token extraction (the designed RFC 8693 + input), and on the v1 arm their presence disables the stored lookup entirely, so a + caller's gateway credential can never displace a per-server BYOK header or leak upstream + as the resolved credential. """ spec = to_server_spec(mcp_server) if spec is None: - stored_headers = await self._resolve_oauth2_headers_for_tool_call( - mcp_server, oauth2_headers, user_api_key_auth - ) + if oauth2_headers: + return None, forwarded_headers + stored_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) return stored_headers, forwarded_headers subject_token: str | None = None @@ -4872,6 +4879,7 @@ class MCPServerManager: ) tasks.append(during_hook_task) + caller_oauth2_headers = oauth2_headers oauth2_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, oauth2_headers, user_api_key_auth) # For OpenAPI servers, call the tool handler directly instead of via MCP client @@ -4891,7 +4899,7 @@ class MCPServerManager: ) resolved_auth_headers, forwarded_headers = await self.resolve_openapi_upstream_auth( mcp_server=mcp_server, - oauth2_headers=oauth2_headers, + oauth2_headers=caller_oauth2_headers, raw_headers=raw_headers, mcp_auth_header=mcp_auth_header, user_api_key_auth=user_api_key_auth, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index ca5b7914cca..b56a12db5b1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1119,3 +1119,80 @@ class TestOpenApiResolvedUpstreamAuth: ) called.assert_not_awaited() + + @pytest.mark.asyncio + async def test_caller_oauth2_headers_never_become_resolved_for_byok_server(self): + """Greptile P1 regression: BYOK servers defer to v1 (to_server_spec None), and the v1 arm + must never promote caller-supplied oauth2 headers into the resolved-auth slot, where they + would override the per-server BYOK credential and leak the caller's gateway Authorization + upstream.""" + manager = MCPServerManager() + server = MCPServer( + server_id="byok-spec", + name="byok_spec", + server_name="byok_spec", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.api_key, + spec_path="https://example.com/openapi.yaml", + is_byok=True, + ) + + resolved, forwarded = await manager.resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers={"Authorization": "Bearer sk-litellm-gateway-key"}, + raw_headers=None, + mcp_auth_header="user-byok-key", + user_api_key_auth=UserAPIKeyAuth(user_id="alice", api_key="sk-user"), + forwarded_headers=None, + ) + + assert resolved is None + assert forwarded is None + + @pytest.mark.asyncio + async def test_v1_server_threads_stored_headers_only_without_caller_headers(self): + """The v1 (unmigrated) arm resolves the stored per-user token only when the caller sent no + oauth2 headers of their own; with caller headers present the stored lookup is skipped and + nothing is promoted to resolved.""" + manager = MCPServerManager() + server = MCPServer( + server_id="v1-spec", + name="v1_spec", + server_name="v1_spec", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + spec_path="https://example.com/openapi.yaml", + delegate_auth_to_upstream=True, + ) + stored = {"Authorization": "Bearer stored-v1-token"} + user_auth = UserAPIKeyAuth(user_id="alice", api_key="sk-user") + + with patch.object( + manager, "_resolve_oauth2_headers_for_tool_call", new=AsyncMock(return_value=stored) + ) as lookup: + resolved, _ = await manager.resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers=None, + raw_headers=None, + mcp_auth_header=None, + user_api_key_auth=user_auth, + forwarded_headers=None, + ) + assert resolved == stored + lookup.assert_awaited_once_with(server, None, user_auth) + + with patch.object( + manager, "_resolve_oauth2_headers_for_tool_call", new=AsyncMock(return_value=stored) + ) as lookup: + resolved, _ = await manager.resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers={"Authorization": "Bearer caller-supplied"}, + raw_headers=None, + mcp_auth_header=None, + user_api_key_auth=user_auth, + forwarded_headers=None, + ) + assert resolved is None + lookup.assert_not_awaited() From 8441ff3a6caff5945bff024ba641035ca3bbb5a9 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 20 Jul 2026 19:40:59 -0700 Subject: [PATCH 03/12] style(mcp): wrap the resolve_openapi_upstream_auth call to the 120 col limit --- litellm/proxy/_experimental/mcp_server/server.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 75faca3c914..396dd6c7dc7 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2788,7 +2788,10 @@ if MCP_AVAILABLE: resolved_auth_headers: dict[str, str] | None = None if mcp_server: - resolved_auth_headers, forwarded_headers = await global_mcp_server_manager.resolve_openapi_upstream_auth( + ( + resolved_auth_headers, + forwarded_headers, + ) = await global_mcp_server_manager.resolve_openapi_upstream_auth( mcp_server=mcp_server, oauth2_headers=oauth2_headers, raw_headers=raw_headers, From 01d624e860a936f2664eb958f7942a9ebccb380e Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 21 Jul 2026 12:01:17 -0700 Subject: [PATCH 04/12] fix(ui): add tooltip to the Active key status badge (#34109) --- .../components/VirtualKeysPage/VirtualKeysTable.test.tsx | 7 ++++++- .../src/components/VirtualKeysPage/keyTableColumns.tsx | 6 +++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 513054aae7a..420a1213a8d 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -498,8 +498,13 @@ describe("Status column reflects blocked / expiry / scim metadata", () => { renderWithProviders(); + const tag = await screen.findByTestId(`key-status-${mockKey.token_id}`); + expect(tag).toHaveTextContent("Active"); + + const user = userEvent.setup(); + await user.hover(tag); await waitFor(() => { - expect(screen.getByTestId(`key-status-${mockKey.token_id}`)).toHaveTextContent("Active"); + expect(screen.getByText(/not blocked and has not expired/i)).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index 133ff89a898..901e878ee5f 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -46,7 +46,11 @@ const getKeyStatus = (key: KeyResponse): KeyStatus => { if (!Number.isNaN(expiresAt) && expiresAt < Date.now()) { return { tone: "warning", label: "Expired", tooltip: "This key has passed its expiry date." }; } - return { tone: "success", label: "Active" }; + return { + tone: "success", + label: "Active", + tooltip: "This key is not blocked and has not expired.", + }; }; const UserPopoverCell = ({ From fcd236097ecfb36eda5489b3bded4513b6209086 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 21 Jul 2026 12:36:21 -0700 Subject: [PATCH 05/12] fix(interactions): add queued to the Interaction status enum (#34135) Google added a queued value to Interaction.status in the live Interactions OpenAPI spec, so the compliance canary test_status_enum_values started failing on every open PR. The exact-match assertion is deliberate; it is how we find out the spec moved, so this adds the new value rather than loosening the check, and mirrors it into the generated Status enums so InteractionStatus stays truthful. --- litellm/types/interactions/generated.py | 2 ++ tests/test_litellm/interactions/test_openapi_compliance.py | 1 + 2 files changed, 3 insertions(+) diff --git a/litellm/types/interactions/generated.py b/litellm/types/interactions/generated.py index 793cc02ff17..4a1ef5ed696 100644 --- a/litellm/types/interactions/generated.py +++ b/litellm/types/interactions/generated.py @@ -173,6 +173,7 @@ class Status1(Enum): cancelled = "cancelled" incomplete = "incomplete" budget_exceeded = "budget_exceeded" + queued = "queued" class InteractionStatusUpdate(BaseModel): @@ -341,6 +342,7 @@ class Status3(Enum): CANCELLED = "cancelled" INCOMPLETE = "incomplete" BUDGET_EXCEEDED = "budget_exceeded" + QUEUED = "queued" class ModelOption(RootModel[str]): diff --git a/tests/test_litellm/interactions/test_openapi_compliance.py b/tests/test_litellm/interactions/test_openapi_compliance.py index 209e99895db..11b08fa45a8 100644 --- a/tests/test_litellm/interactions/test_openapi_compliance.py +++ b/tests/test_litellm/interactions/test_openapi_compliance.py @@ -194,6 +194,7 @@ class TestResponseCompliance: "cancelled", "incomplete", "budget_exceeded", + "queued", ] assert status_prop["enum"] == expected_statuses print(f"✓ Status enum values: {expected_statuses}") From ae2f276d19f486f726264e393f94104119762d40 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 21 Jul 2026 12:56:56 -0700 Subject: [PATCH 06/12] ci(image-scan): match Python packages against CPE data (#34136) grype defaults match.python.using-cpes to false, so PyPI packages are matched only against the GitHub Advisory Database. When a CVE is published to NVD but its GHSA has not propagated to the global advisory database, the scan reports clean even though grype's own database already carries the NVD record with the correct version ranges. The pypdf CVEs (CVE-2026-59935 / 59936 / 59937 / 59938, analyzed in NVD since 2026-07-08) are the case that exposed this; their GHSA IDs are still repo-level and return 404 from the global advisory API, so the ecosystem matcher has nothing to match on. Enabling CPE matching for Python closes that gap. Measured against a v1.91.1 build the finding count goes from 28 to 38; the additions are mostly actionable, and the few cross-product CPE collisions cannot fail the build because --only-fixed drops the ones carrying no fix version and the remainder land below the --fail-on high threshold. --- .github/workflows/image-scan.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index 90ede5a653f..8d791ca5bc7 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -58,6 +58,8 @@ jobs: # free OSS, run as a pinned, checksum-verified binary; no GitHub Action # dependency and no vendor SaaS callout. - name: Scan image for fixable HIGH/CRITICAL CVEs + env: + GRYPE_MATCH_PYTHON_USING_CPES: "true" run: | "$RUNNER_TEMP/grype" litellm-image-scan:${{ github.sha }} \ --only-fixed \ From 257ada88cc7cbc2069a7d5f9430e46a2fd6577ae Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 21 Jul 2026 13:01:18 -0700 Subject: [PATCH 07/12] chore(deps): bump pypdf to 6.14.2 and pyasn1 to 0.6.4 (#34148) Both are lock-only moves. pypdf stays inside the existing >=6.12.0,<7.0 constraint and pyasn1 is transitive, so pyproject.toml is unchanged. pypdf 6.13.3 carries CVE-2026-59935 / 59936 / 59937 / 59938, resolved across 6.14.0 through 6.14.2. pyasn1 0.6.3 carries CVE-2026-59884 / 59885 / 59886, resolved in 0.6.4. All seven are resource-exhaustion issues reachable through parsing untrusted input; pypdf is used for page text extraction in the RAG ingestion file parser. Scanning the lock before and after with CPE matching enabled takes the count for these two packages from seven to zero. --- uv.lock | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/uv.lock b/uv.lock index 1dfa2c1201c..cee24aca330 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-07-15T21:54:47.972166Z" +exclude-newer = "2026-07-18T19:44:23.519632Z" exclude-newer-span = "P3D" [manifest] @@ -6816,11 +6816,11 @@ wheels = [ [[package]] name = "pyasn1" -version = "0.6.3" +version = "0.6.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/5c/5f/6583902b6f79b399c9c40674ac384fd9cd77805f9e6205075f828ef11fb2/pyasn1-0.6.3.tar.gz", hash = "sha256:697a8ecd6d98891189184ca1fa05d1bb00e2f84b5977c481452050549c8a72cf", size = 148685, upload-time = "2026-03-17T01:06:53.382Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a4/9a/23310166d960def5897e91fe20e5b724601b02a22e84ba1f94232c0b7f67/pyasn1-0.6.4.tar.gz", hash = "sha256:9c447d8431c947fe4c8febc4ed9e760bc29011a5b01e5c74b67025bd9fb8ce81", size = 151262, upload-time = "2026-07-09T01:12:33.988Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5d/a0/7d793dce3fa811fe047d6ae2431c672364b462850c6235ae306c0efd025f/pyasn1-0.6.3-py3-none-any.whl", hash = "sha256:a80184d120f0864a52a073acc6fc642847d0be408e7c7252f31390c0f4eadcde", size = 83997, upload-time = "2026-03-17T01:06:52.036Z" }, + { url = "https://files.pythonhosted.org/packages/9a/3b/6163796d69c3977d1e4287bea4a6979161cbbdd170ebb430511e8e1999ce/pyasn1-0.6.4-py3-none-any.whl", hash = "sha256:deda9277cfd454080ec40b207fb6df82206a3a2688735233cdcd8d3d565f088b", size = 84410, upload-time = "2026-07-09T01:12:32.92Z" }, ] [[package]] @@ -7140,14 +7140,14 @@ wheels = [ [[package]] name = "pypdf" -version = "6.13.3" +version = "6.14.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/17/18/9947cc201af9ccf76720fd3347bf4f70eb882ce3fcf4cb05f7443e4cf871/pypdf-6.13.3.tar.gz", hash = "sha256:f3cb822769725f1bac658c406cfc9460399043f3750c2d3e4650e0a85eacabd7", size = 6484063, upload-time = "2026-06-17T15:22:00.898Z" } +sdist = { url = "https://files.pythonhosted.org/packages/03/72/7dfd5ff1c9c37de97a731701f51af091325f123d9d4270361c9c69e4431f/pypdf-6.14.2.tar.gz", hash = "sha256:7873f502fe4385e79539b21d872392dc0c4e3714327c15881cbc7fbfd1f95b25", size = 6491182, upload-time = "2026-06-23T14:18:30.859Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/94/56/2967e621598987905fb8cdfadd8f8de6b5c68c9351f0523c4df8409f28f1/pypdf-6.13.3-py3-none-any.whl", hash = "sha256:c6e3f86afb625791510b02ad5480e94b63970bb957df75d44657c282ecc52224", size = 347288, upload-time = "2026-06-17T15:21:59.512Z" }, + { url = "https://files.pythonhosted.org/packages/49/e6/136aa8993a2ae7214e0b0ef2edaa0d2e08d1d4e4982635b08a835ff31ec8/pypdf-6.14.2-py3-none-any.whl", hash = "sha256:3f07891af76dc002657e04993ab9b4de81de29f9013b9761d0b7968bff12e946", size = 349514, upload-time = "2026-06-23T14:18:28.867Z" }, ] [[package]] From 062e58fb1d652cd468cdecc7f6d13c56b9905d9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 13:03:47 -0700 Subject: [PATCH 08/12] fix(a2a): accept semver protocolVersion values like 0.3.0 in agent cards --- litellm/proxy/a2a/agent_card.py | 25 +++++++--- litellm/proxy/a2a/version_convert.py | 17 ++++--- litellm/proxy/agent_endpoints/endpoints.py | 3 +- .../test_litellm/proxy/a2a/test_agent_card.py | 48 +++++++++++++++++++ .../proxy/a2a/test_version_convert.py | 10 ++++ .../proxy/agent_endpoints/test_endpoints.py | 41 ++++++++++++++++ 6 files changed, 128 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/a2a/agent_card.py b/litellm/proxy/a2a/agent_card.py index e97ab4a01ae..129ad8f0a96 100644 --- a/litellm/proxy/a2a/agent_card.py +++ b/litellm/proxy/a2a/agent_card.py @@ -8,22 +8,35 @@ and uses LiteLLM auth. """ from copy import deepcopy -from typing import Any, Dict, List, Mapping +from typing import Any, Dict, List, Literal, Mapping + +SupportedA2AVersion = Literal["0.3", "1.0"] # Protocol versions LiteLLM can serve to A2A clients. The admin pins one per agent; # responses are normalized to it regardless of the upstream agent's own version. -SUPPORTED_A2A_PROTOCOL_VERSIONS = ("0.3", "1.0") +SUPPORTED_A2A_PROTOCOL_VERSIONS: tuple[SupportedA2AVersion, ...] = ("0.3", "1.0") # Default served version when the agent card does not pin one. LITELLM_A2A_PROTOCOL_VERSION = "1.0" +def normalize_protocol_version(version: object) -> SupportedA2AVersion | None: + """Map a raw ``protocolVersion`` value to the supported canonical major.minor version. + + Semver strings the A2A spec and Google a2a-sdk emit (e.g. ``"0.3.0"``, ``"1.0.1"``) + canonicalize to their major.minor (``"0.3"``, ``"1.0"``). Anything outside the + supported set, including non-strings, yields ``None``. + """ + if not isinstance(version, str): + return None + major_minor = ".".join(version.split(".")[:2]) + return next((supported for supported in SUPPORTED_A2A_PROTOCOL_VERSIONS if supported == major_minor), None) + + def resolve_served_protocol_version(card: Mapping[str, Any] | None) -> str: """Return the validated protocol version an agent card pins, else the default.""" - version = card.get("protocolVersion") if card else None - if version in SUPPORTED_A2A_PROTOCOL_VERSIONS: - return version - return LITELLM_A2A_PROTOCOL_VERSION + normalized = normalize_protocol_version(card.get("protocolVersion") if card else None) + return normalized if normalized is not None else LITELLM_A2A_PROTOCOL_VERSION # Security scheme exposed by the LiteLLM-fronted agent card. Always replaces diff --git a/litellm/proxy/a2a/version_convert.py b/litellm/proxy/a2a/version_convert.py index e8f49e6f6a9..9de33a0966a 100644 --- a/litellm/proxy/a2a/version_convert.py +++ b/litellm/proxy/a2a/version_convert.py @@ -30,6 +30,7 @@ from typing import Callable, Literal, Union from pydantic import BaseModel from litellm._logging import verbose_proxy_logger +from litellm.proxy.a2a.agent_card import normalize_protocol_version A2AVersion = Literal["0.3", "1.0"] RequestId = Union[str, int, None] @@ -103,16 +104,14 @@ def normalize_request_params(params: JsonDict, served: A2AVersion, *, method: st def _detect_card_version(card: JsonDict) -> A2AVersion: """Infer the wire version of an agent card dict. - ``protocolVersion`` is the authoritative indicator; fall back to presence of - ``supportedInterfaces`` (a 1.0-only field) only when the explicit field is absent. - Cards that set ``protocolVersion: "0.3"`` or carry neither signal are treated as 0.3. + ``protocolVersion`` is the authoritative indicator; semver values normalize to + their major.minor (``"0.3.0"`` -> ``"0.3"``). Fall back to presence of + ``supportedInterfaces`` (a 1.0-only field) only when the explicit field is + absent or unrecognized; cards carrying neither signal are treated as 0.3. """ - pv = card.get("protocolVersion") - if pv == "1.0": - return "1.0" - if pv == "0.3": - return "0.3" - # No protocolVersion field: use structural heuristic. + normalized = normalize_protocol_version(card.get("protocolVersion")) + if normalized is not None: + return normalized return "1.0" if "supportedInterfaces" in card else "0.3" diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index a7ceffed97b..2421f270974 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKey from litellm.proxy.a2a.agent_card import ( SUPPORTED_A2A_PROTOCOL_VERSIONS, merge_agent_card, + normalize_protocol_version, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user @@ -51,7 +52,7 @@ def _proxy_base_url(http_request: Request) -> str: def _validate_protocol_version(upstream_card: Mapping[str, Any] | None) -> None: """Reject an agent card pinning an unsupported A2A protocol version.""" version = upstream_card.get("protocolVersion") if upstream_card else None - if version is not None and version not in SUPPORTED_A2A_PROTOCOL_VERSIONS: + if version is not None and normalize_protocol_version(version) is None: raise HTTPException( status_code=400, detail=( diff --git a/tests/test_litellm/proxy/a2a/test_agent_card.py b/tests/test_litellm/proxy/a2a/test_agent_card.py index d302bde7895..32e211b45be 100644 --- a/tests/test_litellm/proxy/a2a/test_agent_card.py +++ b/tests/test_litellm/proxy/a2a/test_agent_card.py @@ -1,10 +1,14 @@ """Unit tests for the pure merge logic in litellm/proxy/a2a/agent_card.py.""" +import pytest + from litellm.proxy.a2a.agent_card import ( LITELLM_A2A_PROTOCOL_VERSION, LITELLM_SECURITY_REQUIREMENTS, LITELLM_SECURITY_SCHEMES, merge_agent_card, + normalize_protocol_version, + resolve_served_protocol_version, ) PROXY_URL = "https://proxy.example/a2a/agent-xyz" @@ -205,3 +209,47 @@ def test_strips_additional_interfaces_to_prevent_backend_url_leak(): ] merged = merge_agent_card(upstream, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE) assert "additionalInterfaces" not in merged + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("0.3", "0.3"), + ("0.3.0", "0.3"), + ("1.0", "1.0"), + ("1.0.0", "1.0"), + ("1.0.1", "1.0"), + ("0.2.6", None), + ("2.0", None), + ("0.30", None), + ("garbage", None), + ("", None), + (None, None), + (1.0, None), + ], +) +def test_normalize_protocol_version(raw, expected): + assert normalize_protocol_version(raw) == expected + + +def test_resolve_served_protocol_version_canonicalizes_semver_pins(): + assert resolve_served_protocol_version({"protocolVersion": "0.3.0"}) == "0.3" + assert resolve_served_protocol_version({"protocolVersion": "1.0.0"}) == "1.0" + assert resolve_served_protocol_version({"protocolVersion": "0.3"}) == "0.3" + assert resolve_served_protocol_version({"protocolVersion": "1.0"}) == "1.0" + + +def test_resolve_served_protocol_version_falls_back_for_unsupported(): + assert ( + resolve_served_protocol_version({"protocolVersion": "0.2.6"}) + == LITELLM_A2A_PROTOCOL_VERSION + ) + assert resolve_served_protocol_version(None) == LITELLM_A2A_PROTOCOL_VERSION + + +def test_serves_semver_pinned_protocol_version_as_major_minor(): + card = _full_upstream_card() + card["protocolVersion"] = "0.3.0" + merged = merge_agent_card(card, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE) + assert merged["protocolVersion"] == "0.3" + assert merged["supportedInterfaces"][0]["protocolVersion"] == "0.3" diff --git a/tests/test_litellm/proxy/a2a/test_version_convert.py b/tests/test_litellm/proxy/a2a/test_version_convert.py index f3c51ca6b72..7eb5debb792 100644 --- a/tests/test_litellm/proxy/a2a/test_version_convert.py +++ b/tests/test_litellm/proxy/a2a/test_version_convert.py @@ -313,3 +313,13 @@ def test_agent_card_with_0_3_pin_and_supported_interfaces_is_lowered(): def test_agent_card_same_version_passthrough(): card = _extended_card_1_0() assert normalize_agent_card(card, "1.0") is card + + +def test_detect_card_version_normalizes_semver_protocol_version(): + from litellm.proxy.a2a.version_convert import _detect_card_version + + assert _detect_card_version({"protocolVersion": "1.0.0"}) == "1.0" + assert ( + _detect_card_version({"protocolVersion": "0.3.0", "supportedInterfaces": []}) + == "0.3" + ) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 3740c01b7fc..d4228f799d5 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -540,6 +540,47 @@ class TestAgentRBACProxyAdmin: assert resp.status_code == 200 +class TestAgentProtocolVersionValidation: + """Registration accepts spec-default semver protocolVersion values and still + rejects genuinely unsupported versions.""" + + @pytest.fixture(autouse=True) + def _setup(self, monkeypatch): + self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN) + self.mock_registry = MagicMock() + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry) + + def _create_agent_with_protocol_version(self, protocol_version: str): + config = _sample_agent_config() + config["agent_card_params"]["protocolVersion"] = protocol_version + with patch("litellm.proxy.proxy_server.prisma_client"): + self.mock_registry.get_agent_by_name = MagicMock(return_value=None) + self.mock_registry.add_agent_to_db = AsyncMock( + return_value=_sample_agent_response() + ) + self.mock_registry.register_agent = MagicMock() + return self.admin_client.post( + "/v1/agents", + json=config, + headers={"Authorization": "Bearer k"}, + ) + + def test_semver_protocol_version_registers_and_stores_major_minor(self): + resp = self._create_agent_with_protocol_version("0.3.0") + assert resp.status_code == 200 + stored_card = self.mock_registry.add_agent_to_db.await_args.kwargs["agent"][ + "agent_card_params" + ] + assert stored_card["protocolVersion"] == "0.3" + assert stored_card["supportedInterfaces"][0]["protocolVersion"] == "0.3" + + def test_unsupported_protocol_version_is_rejected(self): + resp = self._create_agent_with_protocol_version("0.2.6") + assert resp.status_code == 400 + assert "Unsupported protocolVersion '0.2.6'" in resp.json()["detail"] + self.mock_registry.add_agent_to_db.assert_not_awaited() + + class TestCheckAgentManagementPermission: """Unit tests for the _check_agent_management_permission helper.""" From 231021153190870454cab5ac4c0bf5a0fbd5c46e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 13:24:58 -0700 Subject: [PATCH 09/12] fix(a2a): reject malformed protocolVersion suffixes while keeping semver prereleases --- litellm/proxy/a2a/agent_card.py | 18 ++++++++++++++---- .../test_litellm/proxy/a2a/test_agent_card.py | 7 +++++++ .../proxy/agent_endpoints/test_endpoints.py | 6 ++++++ 3 files changed, 27 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/a2a/agent_card.py b/litellm/proxy/a2a/agent_card.py index 129ad8f0a96..29a689a32de 100644 --- a/litellm/proxy/a2a/agent_card.py +++ b/litellm/proxy/a2a/agent_card.py @@ -7,6 +7,7 @@ the base; specific fields are replaced so all traffic flows through the proxy and uses LiteLLM auth. """ +import re from copy import deepcopy from typing import Any, Dict, List, Literal, Mapping @@ -20,16 +21,25 @@ SUPPORTED_A2A_PROTOCOL_VERSIONS: tuple[SupportedA2AVersion, ...] = ("0.3", "1.0" LITELLM_A2A_PROTOCOL_VERSION = "1.0" +_PROTOCOL_VERSION_PATTERN = re.compile( + r"^(\d+\.\d+)(?:\.\d+(?:-[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?(?:\+[0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*)?)?$" +) + + def normalize_protocol_version(version: object) -> SupportedA2AVersion | None: """Map a raw ``protocolVersion`` value to the supported canonical major.minor version. - Semver strings the A2A spec and Google a2a-sdk emit (e.g. ``"0.3.0"``, ``"1.0.1"``) - canonicalize to their major.minor (``"0.3"``, ``"1.0"``). Anything outside the - supported set, including non-strings, yields ``None``. + Accepts the bare major.minor convention of the 1.0 spec (``"0.3"``, ``"1.0"``) and the + full semver forms older SDKs emit (``"0.3.0"``, ``"1.0.1"``, including prerelease and + build suffixes like ``"0.3.0-rc1"``). Malformed strings, versions outside the + supported set, and non-strings yield ``None``. """ if not isinstance(version, str): return None - major_minor = ".".join(version.split(".")[:2]) + match = _PROTOCOL_VERSION_PATTERN.match(version) + if match is None: + return None + major_minor = match.group(1) return next((supported for supported in SUPPORTED_A2A_PROTOCOL_VERSIONS if supported == major_minor), None) diff --git a/tests/test_litellm/proxy/a2a/test_agent_card.py b/tests/test_litellm/proxy/a2a/test_agent_card.py index 32e211b45be..dfa848e335e 100644 --- a/tests/test_litellm/proxy/a2a/test_agent_card.py +++ b/tests/test_litellm/proxy/a2a/test_agent_card.py @@ -219,9 +219,16 @@ def test_strips_additional_interfaces_to_prevent_backend_url_leak(): ("1.0", "1.0"), ("1.0.0", "1.0"), ("1.0.1", "1.0"), + ("0.3.0-rc1", "0.3"), + ("1.0.0-rc.1+build.5", "1.0"), ("0.2.6", None), ("2.0", None), ("0.30", None), + ("0.3.garbage", None), + ("0.3.", None), + ("1.0.not-semver", None), + ("0.3.0.0", None), + ("0.3-rc1", None), ("garbage", None), ("", None), (None, None), diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index d4228f799d5..bcd3333baf9 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -580,6 +580,12 @@ class TestAgentProtocolVersionValidation: assert "Unsupported protocolVersion '0.2.6'" in resp.json()["detail"] self.mock_registry.add_agent_to_db.assert_not_awaited() + def test_malformed_protocol_version_is_rejected(self): + resp = self._create_agent_with_protocol_version("0.3.garbage") + assert resp.status_code == 400 + assert "Unsupported protocolVersion '0.3.garbage'" in resp.json()["detail"] + self.mock_registry.add_agent_to_db.assert_not_awaited() + class TestCheckAgentManagementPermission: """Unit tests for the _check_agent_management_permission helper.""" From efa997dfe0c043c8755403f55964e7283426e3c6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 21 Jul 2026 13:35:01 -0700 Subject: [PATCH 10/12] feat(budgets): add configurable budget_reset_time of day (#31007) Budgets reset at midnight in the configured timezone with no way to control the time of day, so a drained daily budget surfaces as an overnight incident. Add a litellm_settings.budget_reset_time option (e.g. "12:00") that shifts day/week/month resets to a configurable wall-clock time in the existing timezone, so the end of the budget window lands during business hours. The reset time is parsed once into an immutable BudgetResetSettings and injected into the reset job (constructor) and computation, rather than read from a module-level global at call time. A malformed value fails fast at startup. Sub-day durations ignore the offset. Unset preserves midnight resets. --- litellm/litellm_core_utils/duration_parser.py | 142 ++++---- .../proxy/common_utils/reset_budget_job.py | 112 ++++-- litellm/proxy/common_utils/timezone_utils.py | 71 +++- litellm/proxy/proxy_server.py | 13 +- .../test_proxy_budget_reset.py | 33 +- .../test_duration_parser.py | 119 +++++- .../common_utils/test_reset_budget_job.py | 338 ++++++------------ .../proxy/common_utils/test_timezone_utils.py | 82 ++++- 8 files changed, 567 insertions(+), 343 deletions(-) diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 438ff5600ba..79036367652 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -7,8 +7,8 @@ duration_in_seconds is used in diff parts of the code base, example """ import re -import time -from datetime import datetime, timedelta, timezone, tzinfo +import time as time_module +from datetime import datetime, time, timedelta, timezone, tzinfo from typing import Optional, Tuple from zoneinfo import ZoneInfo @@ -61,7 +61,7 @@ def duration_in_seconds(duration: str) -> int: elif unit == "w": return value * 604800 elif unit == "mo": - now = time.time() + now = time_module.time() current_time = datetime.fromtimestamp(now) # Calculate target month and year, handling overflow past December @@ -94,12 +94,17 @@ def duration_in_seconds(duration: str) -> int: raise ValueError(f"Unsupported duration unit, passed duration: {duration}") -def get_next_standardized_reset_time(duration: str, current_time: datetime, timezone_str: str = "UTC") -> datetime: +def get_next_standardized_reset_time( + duration: str, + current_time: datetime, + timezone_str: str = "UTC", + reset_time_of_day: time = time(0, 0), +) -> datetime: """ Get the next standardized reset time based on the duration. All durations will reset at predictable intervals, aligned from the current time: - - Nd: If N=1, reset at next midnight; if N>1, reset every N days from now + - Nd: If N=1, reset at the next `reset_time_of_day`; if N>1, reset every N days from now - Nh: Every N hours, aligned to hour boundaries (e.g., 1:00, 2:00) - Nm: Every N minutes, aligned to minute boundaries (e.g., 1:05, 1:10) - Ns: Every N seconds, aligned to second boundaries @@ -108,12 +113,15 @@ def get_next_standardized_reset_time(duration: str, current_time: datetime, time - duration: Duration string (e.g. "30s", "30m", "30h", "30d") - current_time: Current datetime - timezone_str: Timezone string (e.g. "UTC", "US/Eastern", "Asia/Kolkata") + - reset_time_of_day: Wall-clock time the reset lands on for day/week/month + durations (defaults to midnight). Ignored for sub-day durations, where a + time-of-day is meaningless. Returns: - Next reset time at a standardized interval in the specified timezone """ # Set up timezone and normalize current time - current_time, tz = _setup_timezone(current_time, timezone_str) + current_time, _ = _setup_timezone(current_time, timezone_str) # Parse duration value, unit = _parse_duration(duration) @@ -126,9 +134,9 @@ def get_next_standardized_reset_time(duration: str, current_time: datetime, time # Handle different time units if unit == "d": - return _handle_day_reset(current_time, base_midnight, value, tz) + return _handle_day_reset(current_time, base_midnight, value, reset_time_of_day) elif unit == "w": - return _handle_day_reset(current_time, base_midnight, value * 7, tz) + return _handle_day_reset(current_time, base_midnight, value * 7, reset_time_of_day) elif unit == "h": return _handle_hour_reset(current_time, base_midnight, value) elif unit == "m": @@ -136,7 +144,7 @@ def get_next_standardized_reset_time(duration: str, current_time: datetime, time elif unit == "s": return _handle_second_reset(current_time, base_midnight, value) elif unit == "mo": - return _handle_month_reset(current_time, base_midnight, value) + return _handle_month_reset(current_time, base_midnight, value, reset_time_of_day) else: # Unrecognized unit, default to next midnight return base_midnight + timedelta(days=1) @@ -175,46 +183,58 @@ def _parse_duration(duration: str) -> Tuple[Optional[int], Optional[str]]: return int(value), unit -def _handle_day_reset(current_time: datetime, base_midnight: datetime, value: int, tz: tzinfo) -> datetime: +def _apply_time_of_day(dt: datetime, reset_time_of_day: time) -> datetime: + """Set the wall-clock time of `dt` to `reset_time_of_day`, keeping its date and tzinfo.""" + return dt.replace( + hour=reset_time_of_day.hour, + minute=reset_time_of_day.minute, + second=reset_time_of_day.second, + microsecond=reset_time_of_day.microsecond, + ) + + +def _next_occurrence( + boundary_midnight: datetime, + reset_time_of_day: time, + current_time: datetime, + period: timedelta, +) -> datetime: + """Place the reset at `reset_time_of_day` on the boundary day, rolling forward one + `period` if that instant has already passed (or is exactly now).""" + candidate = _apply_time_of_day(boundary_midnight, reset_time_of_day) + if candidate <= current_time: + return candidate + period + return candidate + + +def _first_of_next_month(first_of_month: datetime) -> datetime: + """Given the 1st of some month, return the 1st of the following month.""" + if first_of_month.month == 12: + return first_of_month.replace(year=first_of_month.year + 1, month=1) + return first_of_month.replace(month=first_of_month.month + 1) + + +def _handle_day_reset( + current_time: datetime, + base_midnight: datetime, + value: int, + reset_time_of_day: time, +) -> datetime: """Handle day-based reset times.""" # Handle zero value - immediate expiration if value == 0: return current_time - if value == 1: # Daily reset at midnight - return base_midnight + timedelta(days=1) - elif value == 7: # Weekly reset on Monday at midnight + if value == 1: # Daily reset at the configured time of day + return _next_occurrence(base_midnight, reset_time_of_day, current_time, timedelta(days=1)) + elif value == 7: # Weekly reset on Monday at the configured time of day days_until_monday = (7 - current_time.weekday()) % 7 - if days_until_monday == 0: # If today is Monday - days_until_monday = 7 - return base_midnight + timedelta(days=days_until_monday) - elif value == 30: # Monthly reset on 1st at midnight - # Get 1st of next month at midnight - if current_time.month == 12: - next_reset = datetime( - year=current_time.year + 1, - month=1, - day=1, - hour=0, - minute=0, - second=0, - microsecond=0, - tzinfo=tz, - ) - else: - next_reset = datetime( - year=current_time.year, - month=current_time.month + 1, - day=1, - hour=0, - minute=0, - second=0, - microsecond=0, - tzinfo=tz, - ) - return next_reset - else: # Custom day value - next interval is value days from current - return current_time.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=value) + upcoming_monday = base_midnight + timedelta(days=days_until_monday) + return _next_occurrence(upcoming_monday, reset_time_of_day, current_time, timedelta(days=7)) + elif value == 30: # Monthly reset on 1st at the configured time of day + return _handle_month_reset(current_time, base_midnight, 1, reset_time_of_day) + else: # Custom day value - next interval is value days from the start of today + return _apply_time_of_day(base_midnight + timedelta(days=value), reset_time_of_day) def _handle_hour_reset(current_time: datetime, base_midnight: datetime, value: int) -> datetime: @@ -316,36 +336,30 @@ def _handle_second_reset(current_time: datetime, base_midnight: datetime, value: return current_time.replace(hour=next_hour, minute=next_minute, second=next_second, microsecond=0) -def _handle_month_reset(current_time: datetime, base_midnight: datetime, value: int) -> datetime: +def _handle_month_reset( + current_time: datetime, + base_midnight: datetime, + value: int, + reset_time_of_day: time, +) -> datetime: """ - Handle monthly reset times. For monthly resets, we always reset at the start of the next month. + Handle monthly reset times. Resets land on the 1st at `reset_time_of_day`; if the + 1st of the current month at that time has already passed, roll to the 1st of next month. Args: current_time: Current datetime base_midnight: Midnight of current day value: Number of months (currently only supports 1 month resets) + reset_time_of_day: Wall-clock time the reset lands on Returns: - datetime: First day of next month at midnight + datetime: First day of the next reset month at `reset_time_of_day` """ if value != 1: raise ValueError("Monthly resets currently only support 1 month intervals") - # Get the first day of next month - if current_time.month == 12: - next_month = 1 - next_year = current_time.year + 1 - else: - next_month = current_time.month + 1 - next_year = current_time.year - - return datetime( - year=next_year, - month=next_month, - day=1, - hour=0, - minute=0, - second=0, - microsecond=0, - tzinfo=current_time.tzinfo, - ) + first_of_this_month = base_midnight.replace(day=1) + candidate = _apply_time_of_day(first_of_this_month, reset_time_of_day) + if candidate <= current_time: + return _apply_time_of_day(_first_of_next_month(first_of_this_month), reset_time_of_day) + return candidate diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index e758420ee37..23a5b8f9c53 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -13,6 +13,11 @@ from litellm.proxy._types import ( LiteLLM_UserTable, LiteLLM_VerificationToken, ) +from litellm.proxy.common_utils.timezone_utils import ( + BudgetResetSettings, + compute_budget_reset_at, + get_budget_reset_settings, +) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import ( @@ -32,9 +37,15 @@ class ResetBudgetJob: Resets the budget for all the keys, users, and teams that need it """ - def __init__(self, proxy_logging_obj: ProxyLogging, prisma_client: PrismaClient): + def __init__( + self, + proxy_logging_obj: ProxyLogging, + prisma_client: PrismaClient, + reset_settings: BudgetResetSettings | None = None, + ): self.proxy_logging_obj: ProxyLogging = proxy_logging_obj self.prisma_client: PrismaClient = prisma_client + self.reset_settings: BudgetResetSettings = reset_settings or get_budget_reset_settings() async def reset_budget( self, @@ -237,7 +248,7 @@ class ResetBudgetJob: if budgets_to_reset is not None and len(budgets_to_reset) > 0: for budget in budgets_to_reset: - budget = await ResetBudgetJob._reset_budget_reset_at_date(budget, now) + budget = await ResetBudgetJob._reset_budget_reset_at_date(budget, now, self.reset_settings) await self.prisma_client.update_data( query_type="update_many", @@ -442,7 +453,11 @@ class ResetBudgetJob: if keys_to_reset is not None and len(keys_to_reset) > 0: for key in keys_to_reset: try: - updated_key = await ResetBudgetJob._reset_budget_for_key(key=key, current_time=now) + updated_key = await ResetBudgetJob._reset_budget_for_key( + key=key, + current_time=now, + reset_settings=self.reset_settings, + ) if updated_key is not None: updated_keys.append(updated_key) else: @@ -513,7 +528,11 @@ class ResetBudgetJob: if users_to_reset is not None and len(users_to_reset) > 0: for user in users_to_reset: try: - updated_user = await ResetBudgetJob._reset_budget_for_user(user=user, current_time=now) + updated_user = await ResetBudgetJob._reset_budget_for_user( + user=user, + current_time=now, + reset_settings=self.reset_settings, + ) if updated_user is not None: updated_users.append(updated_user) else: @@ -588,7 +607,11 @@ class ResetBudgetJob: if teams_to_reset is not None and len(teams_to_reset) > 0: for team in teams_to_reset: try: - updated_team = await ResetBudgetJob._reset_budget_for_team(team=team, current_time=now) + updated_team = await ResetBudgetJob._reset_budget_for_team( + team=team, + current_time=now, + reset_settings=self.reset_settings, + ) if updated_team is not None: updated_teams.append(updated_team) else: @@ -655,10 +678,9 @@ class ResetBudgetJob: counter_key: str, spend_counter_cache: Any, now: datetime, + reset_settings: BudgetResetSettings, ) -> bool: """Reset a single budget window if expired. Returns True if the window was reset.""" - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - reset_at_str = window.get("reset_at") if not reset_at_str: return False @@ -671,7 +693,9 @@ class ResetBudgetJob: await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=0.0) except Exception as redis_err: verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err) - window["reset_at"] = get_budget_reset_time(budget_duration=window["budget_duration"]).isoformat() + window["reset_at"] = compute_budget_reset_at( + budget_duration=window["budget_duration"], settings=reset_settings + ).isoformat() return True async def reset_budget_windows(self) -> None: @@ -703,7 +727,13 @@ class ResetBudgetJob: changed = False for window in windows: counter_key = f"spend:key:{row['token']}:window:{window['budget_duration']}" - if await ResetBudgetJob._reset_expired_window(window, counter_key, spend_counter_cache, now): + if await ResetBudgetJob._reset_expired_window( + window, + counter_key, + spend_counter_cache, + now, + self.reset_settings, + ): changed = True if changed: await VerificationTokenRepository(self.prisma_client).table.update( @@ -726,7 +756,13 @@ class ResetBudgetJob: changed = False for window in windows: counter_key = f"spend:team:{row['team_id']}:window:{window['budget_duration']}" - if await ResetBudgetJob._reset_expired_window(window, counter_key, spend_counter_cache, now): + if await ResetBudgetJob._reset_expired_window( + window, + counter_key, + spend_counter_cache, + now, + self.reset_settings, + ): changed = True if changed: await TeamRepository(self.prisma_client).table.update( @@ -741,6 +777,7 @@ class ResetBudgetJob: item: Union[LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken], current_time: datetime, item_type: Literal["key", "team", "user"], + reset_settings: BudgetResetSettings, ): """ In-place, updates spend=0, and sets budget_reset_at to current_time + budget_duration @@ -755,24 +792,40 @@ class ResetBudgetJob: try: item.spend = 0.0 if hasattr(item, "budget_duration") and item.budget_duration is not None: - from litellm.proxy.common_utils.timezone_utils import ( - get_budget_reset_time, + item.budget_reset_at = compute_budget_reset_at( + budget_duration=item.budget_duration, settings=reset_settings ) - - item.budget_reset_at = get_budget_reset_time(budget_duration=item.budget_duration) return item except Exception as e: verbose_proxy_logger.exception("Error resetting budget for %s: %s. Item: %s", item_type, e, item) raise e @staticmethod - async def _reset_budget_for_team(team: LiteLLM_TeamTable, current_time: datetime) -> Optional[LiteLLM_TeamTable]: - await ResetBudgetJob._reset_budget_common(item=team, current_time=current_time, item_type="team") + async def _reset_budget_for_team( + team: LiteLLM_TeamTable, + current_time: datetime, + reset_settings: BudgetResetSettings, + ) -> LiteLLM_TeamTable | None: + await ResetBudgetJob._reset_budget_common( + item=team, + current_time=current_time, + item_type="team", + reset_settings=reset_settings, + ) return team @staticmethod - async def _reset_budget_for_user(user: LiteLLM_UserTable, current_time: datetime) -> Optional[LiteLLM_UserTable]: - await ResetBudgetJob._reset_budget_common(item=user, current_time=current_time, item_type="user") + async def _reset_budget_for_user( + user: LiteLLM_UserTable, + current_time: datetime, + reset_settings: BudgetResetSettings, + ) -> LiteLLM_UserTable | None: + await ResetBudgetJob._reset_budget_common( + item=user, + current_time=current_time, + item_type="user", + reset_settings=reset_settings, + ) return user @staticmethod @@ -788,15 +841,15 @@ class ResetBudgetJob: @staticmethod async def _reset_budget_reset_at_date( - budget: LiteLLM_BudgetTableFull, current_time: datetime + budget: LiteLLM_BudgetTableFull, + current_time: datetime, + reset_settings: BudgetResetSettings, ) -> LiteLLM_BudgetTableFull: try: if budget.budget_duration is not None: - from litellm.proxy.common_utils.timezone_utils import ( - get_budget_reset_time, + budget.budget_reset_at = compute_budget_reset_at( + budget_duration=budget.budget_duration, settings=reset_settings ) - - budget.budget_reset_at = get_budget_reset_time(budget_duration=budget.budget_duration) except Exception as e: verbose_proxy_logger.exception("Error resetting budget_reset_at for budget: %s. Item: %s", e, budget) raise e @@ -804,7 +857,14 @@ class ResetBudgetJob: @staticmethod async def _reset_budget_for_key( - key: LiteLLM_VerificationToken, current_time: datetime - ) -> Optional[LiteLLM_VerificationToken]: - await ResetBudgetJob._reset_budget_common(item=key, current_time=current_time, item_type="key") + key: LiteLLM_VerificationToken, + current_time: datetime, + reset_settings: BudgetResetSettings, + ) -> LiteLLM_VerificationToken | None: + await ResetBudgetJob._reset_budget_common( + item=key, + current_time=current_time, + item_type="key", + reset_settings=reset_settings, + ) return key diff --git a/litellm/proxy/common_utils/timezone_utils.py b/litellm/proxy/common_utils/timezone_utils.py index 32f9f47d519..a50daf40144 100644 --- a/litellm/proxy/common_utils/timezone_utils.py +++ b/litellm/proxy/common_utils/timezone_utils.py @@ -1,10 +1,47 @@ -from datetime import datetime, timezone +from datetime import datetime, time, timezone + +from pydantic import BaseModel, ConfigDict import litellm from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time -def get_budget_reset_timezone(): +class BudgetResetSettings(BaseModel): + """Immutable, validated settings that govern when budgets reset. + + Parsed once from `litellm_settings` and injected into consumers (the reset + job, management endpoints) so reset times never depend on reaching into + module-level globals at call time. + """ + + model_config = ConfigDict(frozen=True) + + timezone: str = "UTC" + reset_time_of_day: time = time(0, 0) + + +def parse_budget_reset_time(raw: object) -> time: + """Parse a `budget_reset_time` config value (e.g. "12:00") into a `time`. + + Falls back to midnight when unset; raises a clear error on a malformed value + so a bad config fails loudly at startup instead of silently resetting at midnight. + """ + if raw is None or raw == "": + return time(0, 0) + if not isinstance(raw, str): + raise ValueError(f"Invalid budget_reset_time {raw!r}; must be a quoted 24-hour 'HH:MM' string, e.g. \"12:00\"") + for fmt in ("%H:%M", "%H:%M:%S"): + try: + parsed = datetime.strptime(raw, fmt) + return time(hour=parsed.hour, minute=parsed.minute, second=parsed.second) + except ValueError: + continue + raise ValueError( + f"Invalid budget_reset_time {raw!r}; expected a 24-hour 'HH:MM' or 'HH:MM:SS' string, e.g. \"12:00\"" + ) + + +def get_budget_reset_timezone() -> str: """ Get the budget reset timezone from litellm_settings. Falls back to UTC if not specified. @@ -15,15 +52,29 @@ def get_budget_reset_timezone(): return getattr(litellm, "timezone", None) or "UTC" -def get_budget_reset_time(budget_duration: str) -> datetime: - """ - Get the budget reset time based on the configured timezone. - Falls back to UTC if not specified. - """ +def get_budget_reset_settings() -> BudgetResetSettings: + """Build validated reset settings from litellm_settings. Raises on a malformed + `budget_reset_time`, which lets the proxy fail fast at startup.""" + return BudgetResetSettings( + timezone=get_budget_reset_timezone(), + reset_time_of_day=parse_budget_reset_time(getattr(litellm, "budget_reset_time", None)), + ) - reset_at = get_next_standardized_reset_time( + +def compute_budget_reset_at(budget_duration: str, settings: BudgetResetSettings) -> datetime: + """Compute the next reset time for a budget duration using injected settings.""" + return get_next_standardized_reset_time( duration=budget_duration, current_time=datetime.now(timezone.utc), - timezone_str=get_budget_reset_timezone(), + timezone_str=settings.timezone, + reset_time_of_day=settings.reset_time_of_day, ) - return reset_at + + +def get_budget_reset_time(budget_duration: str) -> datetime: + """Get the budget reset time using the globally-configured timezone and reset time. + + Thin wrapper over `compute_budget_reset_at` for callers that don't yet receive + `BudgetResetSettings` by injection (creation/update endpoints, startup backfill). + """ + return compute_budget_reset_at(budget_duration, get_budget_reset_settings()) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3b40abed19e..50fd85a3932 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -319,7 +319,10 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.common_utils.timezone_utils import ( + get_budget_reset_settings, + get_budget_reset_time, +) from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, get_management_object_ttl, @@ -4597,6 +4600,13 @@ class ProxyConfig: litellm.json_logs = True litellm._turn_on_json() verbose_proxy_logger.debug(f"{blue_color_code} Enabled JSON logging via config{reset_color_code}") + elif key == "budget_reset_time": + from litellm.proxy.common_utils.timezone_utils import ( + parse_budget_reset_time, + ) + + parse_budget_reset_time(value) + setattr(litellm, key, value) else: verbose_proxy_logger.debug( f"{blue_color_code} setting litellm.{key}={_redact_general_setting_value(key, value, is_full_admin=False)}{reset_color_code}" @@ -7868,6 +7878,7 @@ class ProxyStartupEvent: budget_reset_job = ResetBudgetJob( proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client, + reset_settings=get_budget_reset_settings(), ) scheduler.add_job( diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 5c96eb619bf..44da3ea06a0 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -30,6 +30,7 @@ def _attrify(d: dict): None)` (et al), which returns None for plain dicts — that would silently skip the row. """ + class _AttrDict(dict): def __getattr__(self, k): try: @@ -120,9 +121,11 @@ async def test_reset_budget_keys_partial_failure(): key1, key2, key3, key4, key5, key6 = ( _attrify(k) for k in [key1, key2, key3, key4, key5, key6] ) - prisma_client.get_data = AsyncMock(return_value=[key1, key2, key3, key4, key5, key6]) + prisma_client.get_data = AsyncMock( + return_value=[key1, key2, key3, key4, key5, key6] + ) - async def fake_reset_key(key, current_time): + async def fake_reset_key(key, current_time, reset_settings=None): if key["id"] == "key1": # Simulate a failure on key1 (for example, this might be due to an invariant check) raise Exception("Simulated failure for key1") @@ -207,9 +210,11 @@ async def test_reset_budget_users_partial_failure(): user1, user2, user3, user4, user5, user6 = ( _attrify(u) for u in [user1, user2, user3, user4, user5, user6] ) - prisma_client.get_data = AsyncMock(return_value=[user1, user2, user3, user4, user5, user6]) + prisma_client.get_data = AsyncMock( + return_value=[user1, user2, user3, user4, user5, user6] + ) - async def fake_reset_user(user, current_time): + async def fake_reset_user(user, current_time, reset_settings=None): if user["id"] == "user1": raise Exception("Simulated failure for user1") else: @@ -397,7 +402,7 @@ async def test_reset_budget_teams_partial_failure(): team1, team2 = _attrify(team1), _attrify(team2) prisma_client.get_data = AsyncMock(return_value=[team1, team2]) - async def fake_reset_team(team, current_time): + async def fake_reset_team(team, current_time, reset_settings=None): if team["id"] == "team1": raise Exception("Simulated failure for team1") else: @@ -513,14 +518,14 @@ async def test_reset_budget_continues_other_categories_on_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_key(key, current_time): + async def fake_reset_key(key, current_time, reset_settings=None): key["spend"] = 0.0 key["budget_reset_at"] = ( current_time + timedelta(seconds=key["budget_duration"]) ).isoformat() return key - async def fake_reset_user(user, current_time): + async def fake_reset_user(user, current_time, reset_settings=None): if user["id"] == "user1": raise Exception("Simulated failure for user1") user["spend"] = 0.0 @@ -529,7 +534,7 @@ async def test_reset_budget_continues_other_categories_on_failure(): ).isoformat() return user - async def fake_reset_team(team, current_time): + async def fake_reset_team(team, current_time, reset_settings=None): team["spend"] = 0.0 team["budget_reset_at"] = ( current_time + timedelta(seconds=team["budget_duration"]) @@ -632,7 +637,7 @@ async def test_service_logger_keys_success(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_key(key, current_time): + async def fake_reset_key(key, current_time, reset_settings=None): key["spend"] = 0.0 key["budget_reset_at"] = ( current_time + timedelta(seconds=key["budget_duration"]) @@ -688,7 +693,7 @@ async def test_service_logger_keys_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_key(key, current_time): + async def fake_reset_key(key, current_time, reset_settings=None): if key["id"] == "key1": raise Exception("Simulated failure for key1") key["spend"] = 0.0 @@ -750,7 +755,7 @@ async def test_service_logger_users_success(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_user(user, current_time): + async def fake_reset_user(user, current_time, reset_settings=None): user["spend"] = 0.0 user["budget_reset_at"] = ( current_time + timedelta(seconds=user["budget_duration"]) @@ -802,7 +807,7 @@ async def test_service_logger_users_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_user(user, current_time): + async def fake_reset_user(user, current_time, reset_settings=None): if user["id"] == "user1": raise Exception("Simulated failure for user1") user["spend"] = 0.0 @@ -863,7 +868,7 @@ async def test_service_logger_teams_success(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_team(team, current_time): + async def fake_reset_team(team, current_time, reset_settings=None): team["spend"] = 0.0 team["budget_reset_at"] = ( current_time + timedelta(seconds=team["budget_duration"]) @@ -915,7 +920,7 @@ async def test_service_logger_teams_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_team(team, current_time): + async def fake_reset_team(team, current_time, reset_settings=None): if team["id"] == "team1": raise Exception("Simulated failure for team1") team["spend"] = 0.0 diff --git a/tests/test_litellm/litellm_core_utils/test_duration_parser.py b/tests/test_litellm/litellm_core_utils/test_duration_parser.py index 3e4446c6672..b6b617610a8 100644 --- a/tests/test_litellm/litellm_core_utils/test_duration_parser.py +++ b/tests/test_litellm/litellm_core_utils/test_duration_parser.py @@ -1,5 +1,5 @@ import unittest -from datetime import datetime, timezone +from datetime import datetime, time, timezone from zoneinfo import ZoneInfo from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time @@ -199,5 +199,122 @@ class TestStandardizedResetTime(unittest.TestCase): self.assertEqual(result, expected) +class TestResetTimeOfDay(unittest.TestCase): + """A configurable reset_time_of_day shifts day/week/month resets off midnight.""" + + def test_daily_reset_before_offset_is_today(self): + now = datetime(2023, 5, 15, 8, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1d", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 15, 12, 0, 0, tzinfo=timezone.utc)) + + def test_daily_reset_after_offset_is_tomorrow(self): + now = datetime(2023, 5, 15, 14, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1d", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 16, 12, 0, 0, tzinfo=timezone.utc)) + + def test_daily_reset_exactly_at_offset_rolls_forward(self): + now = datetime(2023, 5, 15, 12, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1d", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 16, 12, 0, 0, tzinfo=timezone.utc)) + + def test_daily_reset_with_seconds_offset(self): + now = datetime(2023, 5, 15, 8, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1d", now, "UTC", reset_time_of_day=time(9, 30, 15) + ) + self.assertEqual(result, datetime(2023, 5, 15, 9, 30, 15, tzinfo=timezone.utc)) + + def test_offset_applies_in_configured_timezone(self): + # 2023-05-15 22:30 UTC == 2023-05-16 01:30 in Jerusalem (IDT, UTC+3), + # so the next noon-Jerusalem reset is 2023-05-16 12:00 IDT. + now = datetime(2023, 5, 15, 22, 30, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1d", now, "Asia/Jerusalem", reset_time_of_day=time(12, 0) + ) + jerusalem = result.astimezone(ZoneInfo("Asia/Jerusalem")) + self.assertEqual( + (jerusalem.year, jerusalem.month, jerusalem.day), (2023, 5, 16) + ) + self.assertEqual(jerusalem.hour, 12) + self.assertEqual(jerusalem.minute, 0) + + def test_weekly_reset_lands_on_monday_at_offset(self): + wednesday = datetime(2023, 5, 17, 15, 45, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "7d", wednesday, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 22, 12, 0, 0, tzinfo=timezone.utc)) + + def test_weekly_reset_today_is_monday_before_offset_is_today(self): + monday_morning = datetime(2023, 5, 22, 9, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "7d", monday_morning, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 22, 12, 0, 0, tzinfo=timezone.utc)) + + def test_weekly_reset_today_is_monday_after_offset_is_next_week(self): + monday_afternoon = datetime(2023, 5, 22, 15, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "7d", monday_afternoon, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 29, 12, 0, 0, tzinfo=timezone.utc)) + + def test_monthly_30d_lands_on_first_at_offset(self): + now = datetime(2023, 5, 15, 10, 30, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "30d", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 6, 1, 12, 0, 0, tzinfo=timezone.utc)) + + def test_monthly_1mo_today_is_first_before_offset_is_today(self): + now = datetime(2023, 5, 1, 9, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1mo", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 1, 12, 0, 0, tzinfo=timezone.utc)) + + def test_monthly_year_rollover_at_offset(self): + now = datetime(2023, 12, 15, 9, 0, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "1mo", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)) + + def test_custom_day_reset_applies_offset(self): + now = datetime(2023, 5, 15, 10, 30, 0, tzinfo=timezone.utc) + result = get_next_standardized_reset_time( + "3d", now, "UTC", reset_time_of_day=time(12, 0) + ) + self.assertEqual(result, datetime(2023, 5, 18, 12, 0, 0, tzinfo=timezone.utc)) + + def test_sub_day_durations_ignore_offset(self): + base = datetime(2023, 5, 15, 15, 20, 30, tzinfo=timezone.utc) + self.assertEqual( + get_next_standardized_reset_time( + "2h", base, "UTC", reset_time_of_day=time(12, 0) + ), + datetime(2023, 5, 15, 16, 0, 0, tzinfo=timezone.utc), + ) + self.assertEqual( + get_next_standardized_reset_time( + "30m", base, "UTC", reset_time_of_day=time(12, 0) + ), + datetime(2023, 5, 15, 15, 30, 0, tzinfo=timezone.utc), + ) + + def test_default_offset_is_midnight(self): + now = datetime(2023, 5, 15, 10, 30, 0, tzinfo=timezone.utc) + self.assertEqual( + get_next_standardized_reset_time("1d", now, "UTC"), + datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc), + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 5e348b1bb7e..be5bc74c385 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -5,25 +5,23 @@ import sys import time import types from datetime import datetime, timedelta, timezone +from datetime import time as dt_time from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock import pytest -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob +from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings from litellm.proxy.utils import ProxyLogging # Mock classes for testing class MockLiteLLMTeamMembership: - async def update_many( - self, where: Dict[str, Any], data: Dict[str, Any] - ) -> Dict[str, Any]: + async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: # Mock the update_many method for litellm_teammembership return {"count": 1} @@ -32,9 +30,7 @@ class MockLiteLLMVerificationToken: def __init__(self): self.update_many_calls: List[Dict[str, Any]] = [] - async def update_many( - self, where: Dict[str, Any], data: Dict[str, Any] - ) -> Dict[str, Any]: + async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: self.update_many_calls.append({"where": where, "data": data}) return {"count": 1} @@ -52,9 +48,7 @@ class MockLiteLLMOrganizationTable: self.find_many_calls.append({"where": where}) return self._find_many_results - async def update_many( - self, where: Dict[str, Any], data: Dict[str, Any] - ) -> Dict[str, Any]: + async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: self.update_many_calls.append({"where": where, "data": data}) return {"count": 1} @@ -72,9 +66,7 @@ class MockLiteLLMTagTable: self.find_many_calls.append({"where": where}) return self._find_many_results - async def update_many( - self, where: Dict[str, Any], data: Dict[str, Any] - ) -> Dict[str, Any]: + async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: self.update_many_calls.append({"where": where, "data": data}) return {"count": 1} @@ -110,9 +102,7 @@ class MockBatcher: _self._outer = outer def update(_self, where, data): - _self._outer.calls.append( - {"table": _self._table_name, "where": where, "data": data} - ) + _self._outer.calls.append({"table": _self._table_name, "where": where, "data": data}) self.litellm_verificationtoken = _Table("key", self) self.litellm_usertable = _Table("user", self) @@ -172,11 +162,7 @@ class MockPrismaClient: return [item for item in data if hasattr(item, "budget_reset_at")] # Handle specific filtering for enduser table queries - if ( - table_name == "enduser" - and query_type == "find_all" - and "budget_id_list" in kwargs - ): + if table_name == "enduser" and query_type == "find_all" and "budget_id_list" in kwargs: budget_id_list = kwargs["budget_id_list"] # Return endusers that match the budget IDs return [ @@ -188,11 +174,7 @@ class MockPrismaClient: ] # Handle key queries with expires and reset_at - if ( - table_name == "key" - and query_type == "find_all" - and ("expires" in kwargs or "reset_at" in kwargs) - ): + if table_name == "key" and query_type == "find_all" and ("expires" in kwargs or "reset_at" in kwargs): return [item for item in data if hasattr(item, "budget_reset_at")] return data @@ -227,9 +209,7 @@ def mock_proxy_logging(): @pytest.fixture def reset_budget_job(mock_prisma_client, mock_proxy_logging): - return ResetBudgetJob( - proxy_logging_obj=mock_proxy_logging, prisma_client=mock_prisma_client - ) + return ResetBudgetJob(proxy_logging_obj=mock_proxy_logging, prisma_client=mock_prisma_client) # Helper function to run async tests @@ -270,6 +250,40 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client): assert set(write["data"].keys()) == {"spend", "budget_reset_at"} +def test_reset_budget_for_key_honors_injected_reset_time(mock_prisma_client, mock_proxy_logging): + """Injected BudgetResetSettings drives the written reset time end to end (DI, no globals). + + Before the configurable-reset-time change this wrote a midnight reset_at (hour 0); + with noon injected it must write a noon reset_at. + """ + job = ResetBudgetJob( + proxy_logging_obj=mock_proxy_logging, + prisma_client=mock_prisma_client, + reset_settings=BudgetResetSettings(timezone="UTC", reset_time_of_day=dt_time(12, 0)), + ) + now = datetime.now(timezone.utc) + test_key = type( + "LiteLLM_VerificationToken", + (), + { + "spend": 100.0, + "budget_duration": "1d", + "budget_reset_at": now, + "id": "test-key-noon", + "token": "tok-noon", + }, + ) + mock_prisma_client.data["key"] = [test_key] + + asyncio.run(job.reset_budget_for_litellm_keys()) + + key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"] + assert len(key_writes) == 1 + reset_at = key_writes[0]["data"]["budget_reset_at"].astimezone(timezone.utc) + assert reset_at.hour == 12 + assert reset_at.minute == 0 + + def test_reset_budget_for_user(reset_budget_job, mock_prisma_client): # Setup test data with timezone-aware datetime now = datetime.now(timezone.utc) @@ -486,11 +500,7 @@ def test_reset_budget_for_keys_linked_to_budgets(reset_budget_job, mock_prisma_c budgets_to_reset = [test_budget] # Run the method - asyncio.run( - reset_budget_job.reset_budget_for_keys_linked_to_budgets( - budgets_to_reset=budgets_to_reset - ) - ) + asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset)) # Verify that update_many was called on litellm_verificationtoken calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls @@ -531,11 +541,7 @@ def test_reset_budget_for_keys_linked_to_budgets_excludes_keys_with_own_budget_d budgets_to_reset = [test_budget] - asyncio.run( - reset_budget_job.reset_budget_for_keys_linked_to_budgets( - budgets_to_reset=budgets_to_reset - ) - ) + asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset)) calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls assert len(calls) == 1 @@ -548,17 +554,13 @@ def test_reset_budget_for_keys_linked_to_budgets_excludes_keys_with_own_budget_d assert call["where"]["budget_id"] == {"in": ["7d-budget-tier"]} -def test_reset_budget_for_keys_linked_to_budgets_empty( - reset_budget_job, mock_prisma_client -): +def test_reset_budget_for_keys_linked_to_budgets_empty(reset_budget_job, mock_prisma_client): """ Test that when there are no budgets to reset, no update is performed on the verification token table. """ # Run with empty list - asyncio.run( - reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=[]) - ) + asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=[])) # Verify no update_many calls were made calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls @@ -584,11 +586,7 @@ def test_reset_budget_for_orgs_linked_to_budgets(reset_budget_job, mock_prisma_c }, ) - asyncio.run( - reset_budget_job.reset_budget_for_orgs_linked_to_budgets( - budgets_to_reset=[test_budget] - ) - ) + asyncio.run(reset_budget_job.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=[test_budget])) calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls assert len(calls) == 1 @@ -598,16 +596,12 @@ def test_reset_budget_for_orgs_linked_to_budgets(reset_budget_job, mock_prisma_c assert call["data"]["spend"] == 0 -def test_reset_budget_for_orgs_linked_to_budgets_empty( - reset_budget_job, mock_prisma_client -): +def test_reset_budget_for_orgs_linked_to_budgets_empty(reset_budget_job, mock_prisma_client): """ Test that when there are no budgets to reset, no update is performed on the organization table. """ - asyncio.run( - reset_budget_job.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=[]) - ) + asyncio.run(reset_budget_job.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=[])) calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls assert len(calls) == 0 @@ -631,11 +625,7 @@ def test_reset_budget_for_tags_linked_to_budgets(reset_budget_job, mock_prisma_c }, ) - asyncio.run( - reset_budget_job.reset_budget_for_tags_linked_to_budgets( - budgets_to_reset=[test_budget] - ) - ) + asyncio.run(reset_budget_job.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=[test_budget])) calls = mock_prisma_client.db.litellm_tagtable.update_many_calls assert len(calls) == 1 @@ -645,16 +635,12 @@ def test_reset_budget_for_tags_linked_to_budgets(reset_budget_job, mock_prisma_c assert call["data"]["spend"] == 0 -def test_reset_budget_for_tags_linked_to_budgets_empty( - reset_budget_job, mock_prisma_client -): +def test_reset_budget_for_tags_linked_to_budgets_empty(reset_budget_job, mock_prisma_client): """ Test that when there are no budgets to reset, no update is performed on the tag table. """ - asyncio.run( - reset_budget_job.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=[]) - ) + asyncio.run(reset_budget_job.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=[])) calls = mock_prisma_client.db.litellm_tagtable.update_many_calls assert len(calls) == 0 @@ -668,9 +654,7 @@ def test_reset_budget_for_tags_linked_to_budgets_empty( ], ids=["30d-calendar-month", "1mo-calendar-month", "1d-next-midnight"], ) -def test_reset_budget_reset_at_date_calendar_aligned( - budget_duration, expected_day, expected_month -): +def test_reset_budget_reset_at_date_calendar_aligned(budget_duration, expected_day, expected_month): """ Verify that _reset_budget_reset_at_date produces calendar-aligned reset times (matching get_budget_reset_time), not sliding-window offsets. @@ -694,7 +678,7 @@ def test_reset_budget_reset_at_date_calendar_aligned( with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt: mock_dt.now.return_value = fixed_now mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now)) + asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings())) assert test_budget.budget_reset_at.day == expected_day assert test_budget.budget_reset_at.month == expected_month @@ -724,7 +708,7 @@ def test_reset_budget_reset_at_date_7d_next_monday(): with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt: mock_dt.now.return_value = fixed_now mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now)) + asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings())) # Next Monday after Wednesday June 14 is June 19 assert test_budget.budget_reset_at.day == 19 @@ -749,7 +733,7 @@ def test_reset_budget_reset_at_date_none_duration(): }, ) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, now)) + asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, now, BudgetResetSettings())) assert test_budget.budget_reset_at == original_reset_at @@ -773,7 +757,7 @@ def test_reset_budget_reset_at_date_none_reset_at(): with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt: mock_dt.now.return_value = fixed_now mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now)) + asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings())) # Should be set to 1st of next month (July 1) assert test_budget.budget_reset_at is not None @@ -781,9 +765,7 @@ def test_reset_budget_reset_at_date_none_reset_at(): assert test_budget.budget_reset_at.month == 7 -def test_budget_table_reset_also_resets_linked_keys( - reset_budget_job, mock_prisma_client -): +def test_budget_table_reset_also_resets_linked_keys(reset_budget_job, mock_prisma_client): """ Integration-style test: when reset_budget_for_litellm_budget_table runs, it should also reset spend for keys linked to the expiring budget tiers @@ -818,9 +800,7 @@ def test_budget_table_reset_also_resets_linked_keys( assert calls[0]["data"]["spend"] == 0 -def test_budget_table_reset_also_resets_linked_orgs( - reset_budget_job, mock_prisma_client -): +def test_budget_table_reset_also_resets_linked_orgs(reset_budget_job, mock_prisma_client): """ Integration-style test: when reset_budget_for_litellm_budget_table runs, it should also reset spend for orgs linked to the expiring budget tiers @@ -853,9 +833,7 @@ def test_budget_table_reset_also_resets_linked_orgs( assert calls[0]["data"]["spend"] == 0 -def test_budget_table_reset_also_resets_linked_tags( - reset_budget_job, mock_prisma_client -): +def test_budget_table_reset_also_resets_linked_tags(reset_budget_job, mock_prisma_client): """ Integration-style test: when reset_budget_for_litellm_budget_table runs, it should also reset spend for tags linked to the expiring budget tiers. @@ -887,9 +865,7 @@ def test_budget_table_reset_also_resets_linked_tags( assert calls[0]["data"]["spend"] == 0 -def test_reset_budget_resets_endusers_with_null_budget_id( - reset_budget_job, mock_prisma_client -): +def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock_prisma_client): """ When litellm.max_end_user_budget_id is configured and that budget is being reset, end users with budget_id=NULL should also have their spend @@ -959,17 +935,13 @@ def test_reset_budget_resets_endusers_with_null_budget_id( mock_prisma_client.data["enduser"] = [enduser_with_budget] # Set up the DB mock for NULL-budget-id end users - mock_prisma_client.db.litellm_endusertable.set_find_many_results( - [enduser_no_budget_row] - ) + mock_prisma_client.db.litellm_endusertable.set_find_many_results([enduser_no_budget_row]) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) # Both end users should have been reset updated = mock_prisma_client.updated_data["enduser"] - assert ( - len(updated) == 2 - ), f"Expected 2 endusers reset (1 explicit + 1 implicit), got {len(updated)}" + assert len(updated) == 2, f"Expected 2 endusers reset (1 explicit + 1 implicit), got {len(updated)}" user_ids = {u.user_id for u in updated} assert "enduser-explicit" in user_ids @@ -986,9 +958,7 @@ def test_reset_budget_resets_endusers_with_null_budget_id( litellm.max_end_user_budget_id = None -def test_reset_budget_skips_null_budget_id_endusers_when_default_not_configured( - reset_budget_job, mock_prisma_client -): +def test_reset_budget_skips_null_budget_id_endusers_when_default_not_configured(reset_budget_job, mock_prisma_client): """ When litellm.max_end_user_budget_id is NOT configured, end users with budget_id=NULL should NOT be fetched or reset. @@ -1073,20 +1043,14 @@ def test_reset_budget_for_team_members_preserves_total_spend(): mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_teammembership.update_many = AsyncMock( - return_value={"count": 1} - ) + mock_prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1}) - job = ResetBudgetJob( - proxy_logging_obj=MagicMock(), prisma_client=mock_prisma_client - ) + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=mock_prisma_client) asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) mock_prisma_client.db.litellm_teammembership.update_many.assert_called_once() - call_kwargs = ( - mock_prisma_client.db.litellm_teammembership.update_many.call_args.kwargs - ) + call_kwargs = mock_prisma_client.db.litellm_teammembership.update_many.call_args.kwargs assert call_kwargs["where"]["budget_id"]["in"] == ["budget-1"] assert call_kwargs["data"] == {"spend": 0} assert "total_spend" not in call_kwargs["data"] @@ -1142,9 +1106,7 @@ def test_reset_budget_windows_uses_is_not_null_filter(monkeypatch): raises `MissingRequiredValueError`. We work around it by using `query_raw` with `IS NOT NULL`. If someone reverts to the ORM filter, this test fails. """ - job, prisma_client, _ = _make_reset_budget_windows_job( - monkeypatch, key_rows=[], team_rows=[] - ) + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=[], team_rows=[]) asyncio.run(job.reset_budget_windows()) @@ -1184,15 +1146,11 @@ def test_reset_budget_windows_resets_expired_key_window(monkeypatch): # The `budget_limits` payload is re-serialized JSON with a bumped reset_at. written_windows = json.loads(call_kwargs["data"]["budget_limits"]) assert len(written_windows) == 1 - new_reset_at = datetime.fromisoformat( - written_windows[0]["reset_at"].replace("Z", "+00:00") - ).replace(tzinfo=None) + new_reset_at = datetime.fromisoformat(written_windows[0]["reset_at"].replace("Z", "+00:00")).replace(tzinfo=None) assert new_reset_at > now # The spend counter for this key+window was cleared. - spend_counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:key:sk-expired:window:1d", value=0.0 - ) + spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0) def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch): @@ -1206,9 +1164,7 @@ def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch): "budget_limits": [{"budget_duration": "1d", "reset_at": future}], } ] - job, prisma_client, _ = _make_reset_budget_windows_job( - monkeypatch, key_rows=key_rows, team_rows=[] - ) + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) asyncio.run(job.reset_budget_windows()) @@ -1237,9 +1193,7 @@ def test_reset_budget_windows_resets_expired_team_window(monkeypatch): assert call_kwargs["where"] == {"team_id": "team-expired"} assert "budget_limits" in call_kwargs["data"] - spend_counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:team:team-expired:window:30d", value=0.0 - ) + spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team:team-expired:window:30d", value=0.0) def test_reset_budget_windows_handles_string_budget_limits(monkeypatch): @@ -1252,14 +1206,10 @@ def test_reset_budget_windows_handles_string_budget_limits(monkeypatch): key_rows = [ { "token": "sk-string-limits", - "budget_limits": json.dumps( - [{"budget_duration": "1d", "reset_at": expired}] - ), + "budget_limits": json.dumps([{"budget_duration": "1d", "reset_at": expired}]), } ] - job, prisma_client, _ = _make_reset_budget_windows_job( - monkeypatch, key_rows=key_rows, team_rows=[] - ) + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) asyncio.run(job.reset_budget_windows()) @@ -1274,9 +1224,7 @@ def test_reset_budget_windows_skips_row_with_empty_budget_limits(monkeypatch): {"token": "sk-empty-list", "budget_limits": []}, {"token": "sk-empty-str", "budget_limits": ""}, ] - job, prisma_client, _ = _make_reset_budget_windows_job( - monkeypatch, key_rows=key_rows, team_rows=[] - ) + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) asyncio.run(job.reset_budget_windows()) @@ -1361,27 +1309,17 @@ def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch): ) prisma_client = MagicMock() - prisma_client.db.litellm_teammembership.find_many = AsyncMock( - return_value=[membership] - ) - prisma_client.db.litellm_teammembership.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[membership]) + prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:team_member:alice:team-x", value=0.0, ttl=60 - ) - counter_cache.redis_cache.async_set_cache.assert_any_await( - key="spend:team_member:alice:team-x", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team_member:alice:team-x", value=0.0, ttl=60) + counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:team_member:alice:team-x", value=0.0, ttl=60) -def test_reset_budget_for_keys_invalidates_redis_counter( - reset_budget_job, mock_prisma_client, monkeypatch -): +def test_reset_budget_for_keys_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): """Key budget reset must clear the Redis spend counter.""" counter_cache = _make_counter_invalidation_job(monkeypatch) @@ -1402,14 +1340,10 @@ def test_reset_budget_for_keys_invalidates_redis_counter( asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:key:sk-abc", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-abc", value=0.0, ttl=60) -def test_reset_budget_for_users_invalidates_redis_counter( - reset_budget_job, mock_prisma_client, monkeypatch -): +def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): """User budget reset must clear the Redis spend counter.""" counter_cache = _make_counter_invalidation_job(monkeypatch) @@ -1430,14 +1364,10 @@ def test_reset_budget_for_users_invalidates_redis_counter( asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:user:alice", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:user:alice", value=0.0, ttl=60) -def test_reset_budget_for_teams_invalidates_redis_counter( - reset_budget_job, mock_prisma_client, monkeypatch -): +def test_reset_budget_for_teams_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): """Team budget reset must clear the Redis spend counter.""" counter_cache = _make_counter_invalidation_job(monkeypatch) @@ -1458,9 +1388,7 @@ def test_reset_budget_for_teams_invalidates_redis_counter( asyncio.run(reset_budget_job.reset_budget_for_litellm_teams()) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:team:team-x", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team:team-x", value=0.0, ttl=60) def test_reset_does_not_zero_counter_when_db_write_fails(monkeypatch): @@ -1511,9 +1439,7 @@ def test_reset_does_not_zero_counter_when_db_write_fails(monkeypatch): batcher.commit = failing_commit prisma_client.db.batch_ = MagicMock(return_value=batcher) - job = ResetBudgetJob( - proxy_logging_obj=MockProxyLogging(), prisma_client=prisma_client - ) + job = ResetBudgetJob(proxy_logging_obj=MockProxyLogging(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_litellm_keys()) @@ -1543,8 +1469,8 @@ def test_reset_budget_for_keys_writes_only_spend_and_reset_at(reset_budget_job, "budget_duration": "30d", "budget_reset_at": now, "token": "sk-problematic", - "object_permission_id": "perm-abc", # would be rejected on update - "budget_limits": [{"max_budget": 5}], # would be rejected on update + "object_permission_id": "perm-abc", # would be rejected on update + "budget_limits": [{"max_budget": 5}], # would be rejected on update "metadata": {"some": "thing"}, }, ) @@ -1570,19 +1496,13 @@ def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monke linked_key = type("Key", (), {"token": "sk-linked"}) prisma_client = MagicMock() - prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[linked_key] - ) - prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[linked_key]) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget])) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:key:sk-linked", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-linked", value=0.0, ttl=60) def test_reset_budget_for_orgs_linked_to_budgets_invalidates_redis_counter(monkeypatch): @@ -1593,22 +1513,14 @@ def test_reset_budget_for_orgs_linked_to_budgets_invalidates_redis_counter(monke linked_org = type("Org", (), {"organization_id": "org-acme"}) prisma_client = MagicMock() - prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[linked_org] - ) - prisma_client.db.litellm_organizationtable.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[linked_org]) + prisma_client.db.litellm_organizationtable.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_orgs_linked_to_budgets([expired_budget])) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:org:org-acme", value=0.0, ttl=60 - ) - counter_cache.redis_cache.async_set_cache.assert_any_await( - key="spend:org:org-acme", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:org:org-acme", value=0.0, ttl=60) + counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:org:org-acme", value=0.0, ttl=60) def test_reset_budget_for_tags_linked_to_budgets_invalidates_redis_counter(monkeypatch): @@ -1625,12 +1537,8 @@ def test_reset_budget_for_tags_linked_to_budgets_invalidates_redis_counter(monke job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) - counter_cache.in_memory_cache.set_cache.assert_any_call( - key="spend:tag:tenant-42", value=0.0, ttl=60 - ) - counter_cache.redis_cache.async_set_cache.assert_any_await( - key="spend:tag:tenant-42", value=0.0, ttl=60 - ) + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:tag:tenant-42", value=0.0, ttl=60) + counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:tag:tenant-42", value=0.0, ttl=60) def test_reset_budget_for_tags_linked_to_budgets_invalidates_management_cache( @@ -1657,9 +1565,7 @@ def test_reset_budget_for_tags_linked_to_budgets_invalidates_management_cache( job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) - counter_cache.user_api_key_cache.async_delete_cache.assert_any_await( - key="tag:tenant-42" - ) + counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="tag:tenant-42") def test_reset_budget_for_tags_linked_to_budgets_invalidates_each_tag_management_cache( @@ -1684,8 +1590,7 @@ def test_reset_budget_for_tags_linked_to_budgets_invalidates_each_tag_management asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) deleted_keys = { - call.kwargs.get("key") - for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list + call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list } assert deleted_keys == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"} @@ -1711,19 +1616,13 @@ def test_reset_budget_for_keys_linked_to_budgets_invalidates_management_cache( linked_key = type("Key", (), {"token": "sk-linked"}) prisma_client = MagicMock() - prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[linked_key] - ) - prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[linked_key]) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget])) - counter_cache.user_api_key_cache.async_delete_cache.assert_any_await( - key="sk-linked" - ) + counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="sk-linked") def test_reset_budget_for_orgs_linked_to_budgets_invalidates_management_cache( @@ -1736,19 +1635,14 @@ def test_reset_budget_for_orgs_linked_to_budgets_invalidates_management_cache( linked_org = type("Org", (), {"organization_id": "org-acme"}) prisma_client = MagicMock() - prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[linked_org] - ) - prisma_client.db.litellm_organizationtable.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[linked_org]) + prisma_client.db.litellm_organizationtable.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_orgs_linked_to_budgets([expired_budget])) deleted_keys = { - call.kwargs.get("key") - for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list + call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list } assert deleted_keys == { "org_id:org-acme", @@ -1768,19 +1662,13 @@ def test_reset_budget_for_team_members_invalidates_management_cache(monkeypatch) ) prisma_client = MagicMock() - prisma_client.db.litellm_teammembership.find_many = AsyncMock( - return_value=[membership] - ) - prisma_client.db.litellm_teammembership.update_many = AsyncMock( - return_value={"count": 1} - ) + prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[membership]) + prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1}) job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) - counter_cache.user_api_key_cache.async_delete_cache.assert_any_await( - key="team-x_alice" - ) + counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="team-x_alice") def test_reset_budget_for_tags_linked_to_budgets_management_cache_delete_failure_still_resets( @@ -1788,9 +1676,7 @@ def test_reset_budget_for_tags_linked_to_budgets_management_cache_delete_failure ): """If ``async_delete_cache`` raises, the DB cascade must still complete.""" counter_cache = _make_counter_invalidation_job(monkeypatch) - counter_cache.user_api_key_cache.async_delete_cache = AsyncMock( - side_effect=RuntimeError("cache unavailable") - ) + counter_cache.user_api_key_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("cache unavailable")) expired_budget = type("B", (), {"budget_id": "budget-1"}) linked_tag = type("Tag", (), {"tag_name": "tenant-42"}) diff --git a/tests/test_litellm/proxy/common_utils/test_timezone_utils.py b/tests/test_litellm/proxy/common_utils/test_timezone_utils.py index 80b813226df..7f686c53c95 100644 --- a/tests/test_litellm/proxy/common_utils/test_timezone_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_timezone_utils.py @@ -1,19 +1,33 @@ import os import sys -from datetime import datetime, timezone +from datetime import datetime, time, timezone from zoneinfo import ZoneInfo +import pytest + sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path import litellm from litellm.proxy.common_utils.timezone_utils import ( + BudgetResetSettings, + compute_budget_reset_at, + get_budget_reset_settings, get_budget_reset_time, get_budget_reset_timezone, + parse_budget_reset_time, ) +def _restore_attr(obj, name, original): + if original is None: + if hasattr(obj, name): + delattr(obj, name) + else: + setattr(obj, name, original) + + def test_get_budget_reset_time(): """ Test that the budget reset time is set to the first of the next month @@ -100,3 +114,69 @@ def test_get_budget_reset_time_respects_timezone(): delattr(litellm, "timezone") else: litellm.timezone = original + + +def test_parse_budget_reset_time_hh_mm(): + assert parse_budget_reset_time("12:00") == time(12, 0) + + +def test_parse_budget_reset_time_hh_mm_ss(): + assert parse_budget_reset_time("09:30:15") == time(9, 30, 15) + + +def test_parse_budget_reset_time_unset_defaults_to_midnight(): + assert parse_budget_reset_time(None) == time(0, 0) + assert parse_budget_reset_time("") == time(0, 0) + + +def test_parse_budget_reset_time_invalid_string_raises(): + with pytest.raises(ValueError): + parse_budget_reset_time("25:00") + with pytest.raises(ValueError): + parse_budget_reset_time("noon") + + +def test_parse_budget_reset_time_non_string_raises(): + # Unquoted "12:00" in YAML parses to the int 720; it must fail loudly, + # not silently fall back to midnight. + with pytest.raises(ValueError): + parse_budget_reset_time(720) + + +def test_get_budget_reset_settings_reads_globals(): + orig_tz = getattr(litellm, "timezone", None) + orig_rt = getattr(litellm, "budget_reset_time", None) + try: + litellm.timezone = "Asia/Jerusalem" + litellm.budget_reset_time = "12:00" + settings = get_budget_reset_settings() + assert settings.timezone == "Asia/Jerusalem" + assert settings.reset_time_of_day == time(12, 0) + finally: + _restore_attr(litellm, "timezone", orig_tz) + _restore_attr(litellm, "budget_reset_time", orig_rt) + + +def test_compute_budget_reset_at_applies_offset(): + settings = BudgetResetSettings( + timezone="Asia/Jerusalem", reset_time_of_day=time(12, 0) + ) + reset_at = compute_budget_reset_at("1d", settings) + jerusalem = reset_at.astimezone(ZoneInfo("Asia/Jerusalem")) + assert jerusalem.hour == 12 + assert jerusalem.minute == 0 + assert reset_at > datetime.now(timezone.utc) + + +def test_get_budget_reset_time_honors_global_budget_reset_time(): + orig_tz = getattr(litellm, "timezone", None) + orig_rt = getattr(litellm, "budget_reset_time", None) + try: + litellm.timezone = "UTC" + litellm.budget_reset_time = "12:00" + reset_at = get_budget_reset_time(budget_duration="1d") + assert reset_at.astimezone(timezone.utc).hour == 12 + assert reset_at.astimezone(timezone.utc).minute == 0 + finally: + _restore_attr(litellm, "timezone", orig_tz) + _restore_attr(litellm, "budget_reset_time", orig_rt) From ee0028a8417e713d909a634db805df251b328d29 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 21 Jul 2026 13:35:14 -0700 Subject: [PATCH 11/12] feat(ui): surface key budget_reset_at in key info and keys table (#34113) * feat(budgets): add configurable budget_reset_time of day Budgets reset at midnight in the configured timezone with no way to control the time of day, so a drained daily budget surfaces as an overnight incident. Add a litellm_settings.budget_reset_time option (e.g. "12:00") that shifts day/week/month resets to a configurable wall-clock time in the existing timezone, so the end of the budget window lands during business hours. The reset time is parsed once into an immutable BudgetResetSettings and injected into the reset job (constructor) and computation, rather than read from a module-level global at call time. A malformed value fails fast at startup. Sub-day durations ignore the offset. Unset preserves midnight resets. * feat(ui): surface key budget_reset_at in key info and keys table --- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 7 ++ .../VirtualKeysPage/keyTableColumns.tsx | 1 - .../key_info_view.budget_display.test.tsx | 87 ++++++++++++++++++- .../components/templates/key_info_view.tsx | 12 +++ 4 files changed, 105 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 420a1213a8d..95f45ea199e 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -174,6 +174,13 @@ it("should render VirtualKeysTable component", () => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); +it("shows the Budget Reset column by default", async () => { + renderWithProviders(); + await waitFor(() => { + expect(screen.getByText("Budget Reset")).toBeInTheDocument(); + }); +}); + it("left-anchors the create-key CTA below the title, between the header and the table toolbar", () => { renderWithProviders(Create New Key} />); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index 901e878ee5f..fdbc07ee020 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -363,6 +363,5 @@ export const KEY_TABLE_HIDDEN_COLUMNS: Record = { created_by: false, updated_at: false, expires: false, - budget_reset_at: false, rate_limits: false, }; diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx index 5407d37fcf1..bab720f7517 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx @@ -1,5 +1,5 @@ import { renderWithProviders } from "../../../tests/test-utils"; -import { screen, waitFor } from "@testing-library/react"; +import { fireEvent, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import KeyInfoView from "./key_info_view"; @@ -238,3 +238,88 @@ describe("KeyInfoView overview budget display (LIT-2845)", () => { }); }); }); + +describe("KeyInfoView budget reset visibility", () => { + beforeEach(() => { + vi.mocked(useTeams).mockReturnValue({ teams: [], setTeams: vi.fn() }); + vi.mocked(useAuthorized).mockReturnValue(baseAuthorized); + }); + + const KEY_WITH_RESET = { + ...MOCK_KEY_DATA, + max_budget: 0.1, + budget_duration: "1d", + budget_reset_at: "2026-07-22T12:00:00+00:00", + } as unknown as KeyResponse; + + it("shows the next budget reset in the overview Spend card when budget_reset_at is set", async () => { + renderWithProviders( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + await waitFor(() => { + expect(screen.getByText(/^Resets Jul 22, 2026/)).toBeInTheDocument(); + }); + }); + + it("omits the reset line from the overview Spend card when budget_reset_at is null", async () => { + renderWithProviders( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + await waitFor(() => { + expect(screen.getByText(/of \$0\.10/)).toBeInTheDocument(); + }); + expect(screen.queryByText(/^Resets /)).not.toBeInTheDocument(); + }); + + it("shows the duration and next reset in the Settings tab", async () => { + renderWithProviders( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + await waitFor(() => { + expect(screen.getByRole("tab", { name: "Settings" })).toBeInTheDocument(); + }); + fireEvent.click(screen.getByRole("tab", { name: "Settings" })); + await waitFor(() => { + expect(screen.getByText("Budget Reset")).toBeInTheDocument(); + }); + expect(screen.getByText(/Every 1d, next Jul 22, 2026/)).toBeInTheDocument(); + }); + + it("shows 'Never' in the Settings tab when no reset is scheduled", async () => { + renderWithProviders( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + await waitFor(() => { + expect(screen.getByRole("tab", { name: "Settings" })).toBeInTheDocument(); + }); + fireEvent.click(screen.getByRole("tab", { name: "Settings" })); + await waitFor(() => { + expect(screen.getByText("Budget Reset")).toBeInTheDocument(); + }); + expect(screen.getByText("Budget Reset").parentElement).toHaveTextContent("Never"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index cd97f6b851d..23baf024bd1 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -534,6 +534,9 @@ export default function KeyInfoView({
${formatNumberWithCommas(currentKeyData.spend, 4)} of {budgetDisplay} + {currentKeyData.budget_reset_at && ( + Resets {formatTimestamp(currentKeyData.budget_reset_at)} + )}
@@ -751,6 +754,15 @@ export default function KeyInfoView({ +
+ Budget Reset + + {currentKeyData.budget_reset_at + ? `${currentKeyData.budget_duration ? `Every ${currentKeyData.budget_duration}, next ` : ""}${formatTimestamp(currentKeyData.budget_reset_at)}` + : "Never"} + +
+ {currentKeyData.budget_fallbacks && Object.keys(currentKeyData.budget_fallbacks).length > 0 && (
Budget Fallbacks From d2819baf0af37450616c5687a3613dee2a3e866b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 21 Jul 2026 13:41:10 -0700 Subject: [PATCH 12/12] feat(ui): add block/unblock key action to key info page (#34116) Adds a Block Key / Unblock Key action to the key info page, wired to the existing /key/block and /key/unblock endpoints which previously had no UI. The Reset Spend and Delete Key buttons move together with it into a new overflow dropdown next to Regenerate Key, and a red Blocked tag shows next to the key alias while the key is blocked. --- .../hooks/keys/useSetKeyBlockedState.test.ts | 103 ++++++++++++++++++ .../hooks/keys/useSetKeyBlockedState.ts | 45 ++++++++ .../src/components/networking.tsx | 2 +- .../templates/KeyInfoHeader.test.tsx | 94 +++++++++++++--- .../components/templates/KeyInfoHeader.tsx | 57 ++++++++-- .../components/templates/key_info_view.tsx | 65 ++++++++++- 6 files changed, 330 insertions(+), 36 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.test.ts new file mode 100644 index 00000000000..5eb3bdc105d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.test.ts @@ -0,0 +1,103 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useSetKeyBlockedState, setKeyBlockedState } from "./useSetKeyBlockedState"; +import { apiClient } from "@/components/networking"; + +vi.mock("@/components/networking", () => ({ + apiClient: { post: vi.fn() }, +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +const mockPost = vi.mocked(apiClient.post); + +const createWrapper = () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false }, mutations: { retry: false } } }); + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + return { queryClient, wrapper }; +}; + +describe("setKeyBlockedState", () => { + beforeEach(() => { + mockPost.mockReset(); + }); + + it("POSTs the key hash to /key/block when blocking", async () => { + mockPost.mockResolvedValueOnce({ blocked: true }); + + const result = await setKeyBlockedState("sk-access", { keyToken: "hashed-token", blocked: true }); + + expect(mockPost).toHaveBeenCalledWith("/key/block", { + accessToken: "sk-access", + body: { key: "hashed-token" }, + }); + expect(result).toEqual({ blocked: true }); + }); + + it("POSTs the key hash to /key/unblock when unblocking", async () => { + mockPost.mockResolvedValueOnce({ blocked: false }); + + const result = await setKeyBlockedState("sk-access", { keyToken: "hashed-token", blocked: false }); + + expect(mockPost).toHaveBeenCalledWith("/key/unblock", { + accessToken: "sk-access", + body: { key: "hashed-token" }, + }); + expect(result).toEqual({ blocked: false }); + }); + + it("falls back to the requested state when the response has no blocked field", async () => { + mockPost.mockResolvedValueOnce(null); + + const result = await setKeyBlockedState("sk-access", { keyToken: "hashed-token", blocked: true }); + + expect(result).toEqual({ blocked: true }); + }); +}); + +describe("useSetKeyBlockedState", () => { + beforeEach(() => { + mockPost.mockReset(); + mockUseAuthorized.mockReturnValue({ accessToken: "sk-access" }); + }); + + it("invalidates key queries after a successful mutation", async () => { + mockPost.mockResolvedValueOnce({ blocked: true }); + const { queryClient, wrapper } = createWrapper(); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + + const { result } = renderHook(() => useSetKeyBlockedState(), { wrapper }); + result.current.mutate({ keyToken: "hashed-token", blocked: true }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: ["keys"] }); + }); + + it("surfaces request failures as mutation errors", async () => { + mockPost.mockRejectedValueOnce(new Error("Key not found.")); + const { wrapper } = createWrapper(); + + const { result } = renderHook(() => useSetKeyBlockedState(), { wrapper }); + result.current.mutate({ keyToken: "missing", blocked: true }); + + await waitFor(() => expect(result.current.isError).toBe(true)); + expect(result.current.error?.message).toBe("Key not found."); + }); + + it("errors without an access token", async () => { + mockUseAuthorized.mockReturnValue({ accessToken: null }); + const { wrapper } = createWrapper(); + + const { result } = renderHook(() => useSetKeyBlockedState(), { wrapper }); + result.current.mutate({ keyToken: "hashed-token", blocked: true }); + + await waitFor(() => expect(result.current.isError).toBe(true)); + expect(mockPost).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.ts new file mode 100644 index 00000000000..792ef567f99 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useSetKeyBlockedState.ts @@ -0,0 +1,45 @@ +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { apiClient } from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { keyKeys } from "./useKeys"; + +export interface SetKeyBlockedStateInput { + keyToken: string; + blocked: boolean; +} + +export interface SetKeyBlockedStateResult { + blocked: boolean; +} + +interface BlockKeyResponse { + blocked?: boolean | null; +} + +export const setKeyBlockedState = async ( + accessToken: string, + { keyToken, blocked }: SetKeyBlockedStateInput, +): Promise => { + const response = await apiClient.post(blocked ? "/key/block" : "/key/unblock", { + accessToken, + body: { key: keyToken }, + }); + return { blocked: response?.blocked ?? blocked }; +}; + +export const useSetKeyBlockedState = () => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async (input) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return setKeyBlockedState(accessToken, input); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: keyKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index d44a491b840..d6e9ba5665c 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -372,7 +372,7 @@ export function getGlobalLitellmHeaderName(): string { return globalLitellmHeaderName; } -const apiClient = createApiClient({ +export const apiClient = createApiClient({ getBaseUrl: getProxyBaseUrl, getAuthHeaderName: getGlobalLitellmHeaderName, onError: handleError, diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx index 14750787609..f67f70e7df9 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.test.tsx @@ -60,22 +60,16 @@ describe("KeyInfoHeader", () => { }); describe("action buttons", () => { - it("should show Regenerate and Delete buttons by default", () => { + it("should show Regenerate button and actions dropdown by default", () => { render(); expect(screen.getByRole("button", { name: /regenerate key/i })).toBeInTheDocument(); - expect(screen.getByRole("button", { name: /delete key/i })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /more key actions/i })).toBeInTheDocument(); }); - it("should show Regenerate and Delete buttons when canModifyKey is true", () => { - render(); - expect(screen.getByRole("button", { name: /regenerate key/i })).toBeInTheDocument(); - expect(screen.getByRole("button", { name: /delete key/i })).toBeInTheDocument(); - }); - - it("should hide Regenerate and Delete buttons when canModifyKey is false", () => { + it("should hide Regenerate button and actions dropdown when canModifyKey is false", () => { render(); expect(screen.queryByRole("button", { name: /regenerate key/i })).not.toBeInTheDocument(); - expect(screen.queryByRole("button", { name: /delete key/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /more key actions/i })).not.toBeInTheDocument(); }); it("should call onRegenerate when Regenerate Key is clicked", async () => { @@ -85,13 +79,6 @@ describe("KeyInfoHeader", () => { expect(onRegenerate).toHaveBeenCalledTimes(1); }); - it("should call onDelete when Delete Key is clicked", async () => { - const onDelete = vi.fn(); - render(); - await userEvent.click(screen.getByRole("button", { name: /delete key/i })); - expect(onDelete).toHaveBeenCalledTimes(1); - }); - it("should disable Regenerate button when regenerateDisabled is true", () => { render(); expect(screen.getByRole("button", { name: /regenerate key/i })).toBeDisabled(); @@ -103,6 +90,79 @@ describe("KeyInfoHeader", () => { }); }); + describe("destructive actions dropdown", () => { + const openDropdown = async () => { + await userEvent.click(screen.getByRole("button", { name: /more key actions/i })); + }; + + it("should list Block Key, Reset Spend, and Delete Key when all handlers are provided", async () => { + render(); + await openDropdown(); + expect(await screen.findByRole("menuitem", { name: /block key/i })).toBeInTheDocument(); + expect(screen.getByRole("menuitem", { name: /reset spend/i })).toBeInTheDocument(); + expect(screen.getByRole("menuitem", { name: /delete key/i })).toBeInTheDocument(); + }); + + it("should omit Block Key and Reset Spend when their handlers are not provided", async () => { + render(); + await openDropdown(); + expect(await screen.findByRole("menuitem", { name: /delete key/i })).toBeInTheDocument(); + expect(screen.queryByRole("menuitem", { name: /block key/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("menuitem", { name: /reset spend/i })).not.toBeInTheDocument(); + }); + + it("should show Unblock Key instead of Block Key when the key is blocked", async () => { + render(); + await openDropdown(); + expect(await screen.findByRole("menuitem", { name: /unblock key/i })).toBeInTheDocument(); + expect(screen.queryByRole("menuitem", { name: /^block key/i })).not.toBeInTheDocument(); + }); + + it("should call onToggleBlocked when Block Key is clicked", async () => { + const onToggleBlocked = vi.fn(); + render(); + await openDropdown(); + await userEvent.click(await screen.findByRole("menuitem", { name: /block key/i })); + expect(onToggleBlocked).toHaveBeenCalledTimes(1); + }); + + it("should call onToggleBlocked when Unblock Key is clicked", async () => { + const onToggleBlocked = vi.fn(); + render(); + await openDropdown(); + await userEvent.click(await screen.findByRole("menuitem", { name: /unblock key/i })); + expect(onToggleBlocked).toHaveBeenCalledTimes(1); + }); + + it("should call onResetSpend when Reset Spend is clicked", async () => { + const onResetSpend = vi.fn(); + render(); + await openDropdown(); + await userEvent.click(await screen.findByRole("menuitem", { name: /reset spend/i })); + expect(onResetSpend).toHaveBeenCalledTimes(1); + }); + + it("should call onDelete when Delete Key is clicked", async () => { + const onDelete = vi.fn(); + render(); + await openDropdown(); + await userEvent.click(await screen.findByRole("menuitem", { name: /delete key/i })); + expect(onDelete).toHaveBeenCalledTimes(1); + }); + }); + + describe("blocked tag", () => { + it("should show a Blocked tag when isBlocked is true", () => { + render(); + expect(screen.getByText("Blocked")).toBeInTheDocument(); + }); + + it("should not show a Blocked tag by default", () => { + render(); + expect(screen.queryByText("Blocked")).not.toBeInTheDocument(); + }); + }); + describe("Create New Key button", () => { it("should show when onCreateNew is provided", () => { render(); diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx index c9c393352cc..d0dd782a697 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoHeader.tsx @@ -1,5 +1,6 @@ import React from "react"; -import { Button, Typography, Tooltip, Space, Divider, Flex, Popover } from "antd"; +import { Button, Typography, Tooltip, Space, Divider, Flex, Popover, Dropdown, Tag } from "antd"; +import type { MenuProps } from "antd"; import { ArrowLeftOutlined, SyncOutlined, @@ -12,6 +13,9 @@ import { SafetyCertificateOutlined, TransactionOutlined, FieldTimeOutlined, + MoreOutlined, + StopOutlined, + CheckCircleOutlined, } from "@ant-design/icons"; import LabeledField from "../common_components/LabeledField"; import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; @@ -38,6 +42,8 @@ interface KeyInfoHeaderProps { onRegenerate?: () => void; onDelete?: () => void; onResetSpend?: () => void; + onToggleBlocked?: () => void; + isBlocked?: boolean; canModifyKey?: boolean; backButtonText?: string; regenerateDisabled?: boolean; @@ -133,11 +139,33 @@ export function KeyInfoHeader({ onRegenerate, onDelete, onResetSpend, + onToggleBlocked, + isBlocked = false, canModifyKey = true, backButtonText = "Back to Keys", regenerateDisabled = false, regenerateTooltip, }: KeyInfoHeaderProps) { + const destructiveActionItems: MenuProps["items"] = [ + ...(onToggleBlocked + ? [ + isBlocked + ? { key: "unblock", label: "Unblock Key", icon: } + : { key: "block", label: "Block Key", icon: , danger: true }, + ] + : []), + ...(onResetSpend + ? [{ key: "reset-spend", label: "Reset Spend", icon: , danger: true }] + : []), + { key: "delete", label: "Delete Key", icon: , danger: true }, + ]; + + const handleDestructiveActionClick: MenuProps["onClick"] = ({ key }) => { + if (key === "block" || key === "unblock") onToggleBlocked?.(); + if (key === "reset-spend") onResetSpend?.(); + if (key === "delete") onDelete?.(); + }; + return (
{onCreateNew && ( @@ -156,9 +184,16 @@ export function KeyInfoHeader({
- - {data.keyName} - + + + {data.keyName} + + {isBlocked && ( + }> + Blocked + + )} + Key ID: {data.keyId} @@ -172,14 +207,12 @@ export function KeyInfoHeader({ - {onResetSpend && ( - - )} - + +