From 5e68a003476d0a2ecce31036ea25597cf0a549d7 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 21 Jul 2026 00:10:49 +0000 Subject: [PATCH 01/84] test(e2e): add live A2A agent e2e suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 1 + tests/e2e/a2a/a2a_client.py | 217 +++++++++++++++++++++++++ tests/e2e/a2a/conftest.py | 17 ++ tests/e2e/a2a/test_a2a_agent_e2e.py | 139 ++++++++++++++++ tests/e2e/coverage_registry/other.yaml | 6 + 5 files changed, 380 insertions(+) create mode 100644 tests/e2e/a2a/a2a_client.py create mode 100644 tests/e2e/a2a/conftest.py create mode 100644 tests/e2e/a2a/test_a2a_agent_e2e.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 47f3c74d7f1..c35fb2fa435 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -13,6 +13,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `realtime/` - realtime websocket sessions, including the pipecat audio path - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`) - `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface; not Playwright) +- `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0) - `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server only (see "MCP suite: real Datadog only" below) - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py new file mode 100644 index 00000000000..5274ec15383 --- /dev/null +++ b/tests/e2e/a2a/a2a_client.py @@ -0,0 +1,217 @@ +"""Client for the proxy's A2A (agent-to-agent) surface. + +An A2A agent is registered admin-side via POST /v1/agents with an agent card and +litellm_params; the proxy fronts it at /a2a/{id}, serving a proxy-owned agent card +at /.well-known/agent-card.json and accepting A2A JSON-RPC calls at /a2a/{id}. This +suite registers agents backed by the litellm_completion_bridge (custom_llm_provider ++ model), so message/send runs a real provider completion and comes back in the +agent's pinned A2A protocol version. The A2A request/response models are co-located +here because only this suite uses them. +""" + +from __future__ import annotations + +import warnings +from dataclasses import dataclass + +from pydantic import BaseModel, ConfigDict, Field + +from e2e_http import NoBody, Result, is_ok +from proxy_client import ProxyClient + + +class A2ACapabilities(BaseModel): + streaming: bool | None = None + push_notifications: bool | None = Field(default=None, serialization_alias="pushNotifications") + + +class A2ASkill(BaseModel): + id: str + name: str + description: str + tags: list[str] + + +class AgentCardParams(BaseModel): + """The upstream agent card an admin registers. `protocolVersion` is the field the + proxy validates against SUPPORTED_A2A_PROTOCOL_VERSIONS on registration.""" + + protocol_version: str = Field(serialization_alias="protocolVersion") + name: str + description: str + version: str + capabilities: A2ACapabilities = A2ACapabilities() + skills: list[A2ASkill] + default_input_modes: list[str] = Field(default=["text"], serialization_alias="defaultInputModes") + default_output_modes: list[str] = Field(default=["text"], serialization_alias="defaultOutputModes") + + +class A2ABridgeParams(BaseModel): + """litellm_params that route the agent through the completion bridge: an A2A + message/send is transformed into a litellm.acompletion against this provider.""" + + model_config = ConfigDict(protected_namespaces=()) + + custom_llm_provider: str + model: str + + +class AgentRegisterBody(BaseModel): + agent_name: str + agent_card_params: AgentCardParams + litellm_params: A2ABridgeParams + + +class A2ASecurityScheme(BaseModel): + type: str + scheme: str + + +class A2AInterface(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + url: str + protocol_version: str | None = Field(default=None, alias="protocolVersion") + + +class ServedAgentCard(BaseModel): + """The proxy-owned card, either nested under a registration response's + `agent_card_params` or served raw at /.well-known/agent-card.json. The proxy + rewrites `url`/`supportedInterfaces` to itself and replaces the security scheme + with its own virtual-key bearer scheme.""" + + model_config = ConfigDict(populate_by_name=True) + + protocol_version: str = Field(alias="protocolVersion") + name: str + url: str | None = None + security_schemes: dict[str, A2ASecurityScheme] | None = Field(default=None, alias="securitySchemes") + security: list[dict[str, list[str]]] | None = None + supported_interfaces: list[A2AInterface] | None = Field(default=None, alias="supportedInterfaces") + + +class AgentResponse(BaseModel): + agent_id: str + agent_name: str + agent_card_params: ServedAgentCard + + +class A2ATextPart(BaseModel): + kind: str = "text" + text: str + + +class A2AOutboundMessage(BaseModel): + role: str = "user" + parts: list[A2ATextPart] + message_id: str = Field(serialization_alias="messageId") + + +class A2AMessageSendParams(BaseModel): + message: A2AOutboundMessage + + +class A2AJsonRpcRequest(BaseModel): + jsonrpc: str = "2.0" + id: str + method: str = "message/send" + params: A2AMessageSendParams + + +class A2AResponsePart(BaseModel): + kind: str | None = None + text: str | None = None + + +class A2AResponseMessage(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + message_id: str | None = Field(default=None, alias="messageId") + role: str | None = None + parts: list[A2AResponsePart] = [] + + +class A2AResult(BaseModel): + """A message/send result. In 0.3 the message fields sit directly on the result + (`kind`/`role`/`parts`); in 1.0 they are nested under `message`. `text` reads the + agent's reply from whichever shape the served version produced.""" + + model_config = ConfigDict(populate_by_name=True) + + kind: str | None = None + role: str | None = None + message_id: str | None = Field(default=None, alias="messageId") + parts: list[A2AResponsePart] = [] + message: A2AResponseMessage | None = None + + @property + def text(self) -> str: + parts = self.message.parts if self.message is not None else self.parts + return "".join(part.text or "" for part in parts) + + @property + def is_nested_v1_shape(self) -> bool: + return self.message is not None + + +class A2AError(BaseModel): + code: int + message: str + + +class A2AResponse(BaseModel): + jsonrpc: str + id: str | None = None + result: A2AResult | None = None + error: A2AError | None = None + + +@dataclass(frozen=True, slots=True) +class A2AClient: + proxy: ProxyClient + + def register_agent(self, body: AgentRegisterBody) -> Result[AgentResponse]: + return self.proxy.transport.post( + "/v1/agents", + headers=self.proxy.transport.master, + json=body, + response_type=AgentResponse, + ) + + def get_agent(self, agent_id: str) -> Result[AgentResponse]: + return self.proxy.transport.get( + f"/v1/agents/{agent_id}", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=AgentResponse, + ) + + def delete_agent(self, agent_id: str) -> None: + result = self.proxy.transport.delete( + f"/v1/agents/{agent_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + if not is_ok(result): + warnings.warn(f"delete_agent({agent_id!r}) failed: {result}", stacklevel=2) + + def agent_card(self, agent_id: str, key: str) -> Result[ServedAgentCard]: + return self.proxy.transport.get( + f"/a2a/{agent_id}/.well-known/agent-card.json", + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=ServedAgentCard, + ) + + def send_message(self, agent_id: str, key: str, body: A2AJsonRpcRequest) -> Result[A2AResponse]: + return self.proxy.transport.post( + f"/a2a/{agent_id}", + headers=self.proxy.transport.bearer(key), + json=body, + response_type=A2AResponse, + ) + + +def build_a2a_client(proxy: ProxyClient) -> A2AClient: + return A2AClient(proxy=proxy) diff --git a/tests/e2e/a2a/conftest.py b/tests/e2e/a2a/conftest.py new file mode 100644 index 00000000000..93f3b56c8f7 --- /dev/null +++ b/tests/e2e/a2a/conftest.py @@ -0,0 +1,17 @@ +"""A2A suite's `client` fixture. + +The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker +live in the parent tests/e2e/conftest.py. A2AClient holds the shared ProxyClient, +so the `resources` fixture cleans up keys this suite creates; agents are torn down +via `resources.defer(...)` in each test. +""" + +import pytest + +from a2a_client import A2AClient, build_a2a_client +from proxy_client import ProxyClient + + +@pytest.fixture(scope="session") +def client(proxy: ProxyClient) -> A2AClient: + return build_a2a_client(proxy) diff --git a/tests/e2e/a2a/test_a2a_agent_e2e.py b/tests/e2e/a2a/test_a2a_agent_e2e.py new file mode 100644 index 00000000000..eb61ace238c --- /dev/null +++ b/tests/e2e/a2a/test_a2a_agent_e2e.py @@ -0,0 +1,139 @@ +"""A2A agents end to end, against a live proxy. + +An admin registers an agent whose card pins an A2A protocol version and whose +litellm_params route it through the completion bridge; a caller then discovers the +proxy-owned card and drives it over A2A JSON-RPC. These tests assert the recorded +state (the agent persists, a spend row lands) and the enforced behavior (the served +card points back at the proxy, message/send returns a real completion in the pinned +protocol version, and an unsupported version is refused at registration). +""" + +from __future__ import annotations + +import pytest + +from a2a_client import ( + A2ABridgeParams, + A2AClient, + A2AJsonRpcRequest, + A2AMessageSendParams, + A2AOutboundMessage, + A2ASkill, + A2ATextPart, + AgentCardParams, + AgentRegisterBody, + AgentResponse, +) +from e2e_config import unique_marker +from e2e_http import UnknownApiError, unwrap +from lifecycle import ResourceManager + +BRIDGE = A2ABridgeParams(custom_llm_provider="anthropic", model="claude-haiku-4-5") + +pytestmark = pytest.mark.e2e + + +def _register(client: A2AClient, resources: ResourceManager, protocol_version: str) -> AgentResponse: + marker = unique_marker() + body = AgentRegisterBody( + agent_name=f"e2e-a2a-{marker}", + agent_card_params=AgentCardParams( + protocol_version=protocol_version, + name=f"E2E A2A {marker}", + description="e2e agent backed by the litellm completion bridge", + version="1.0.0", + skills=[A2ASkill(id="chat", name="Chat", description="general chat", tags=["chat"])], + ), + litellm_params=BRIDGE, + ) + agent = unwrap(client.register_agent(body)) + resources.defer(lambda: client.delete_agent(agent.agent_id)) + return agent + + +def _ask(text: str) -> A2AJsonRpcRequest: + return A2AJsonRpcRequest( + id=f"e2e-{unique_marker()}", + params=A2AMessageSendParams( + message=A2AOutboundMessage(parts=[A2ATextPart(text=text)], message_id=unique_marker()) + ), + ) + + +class TestA2AAgentLifecycle: + @pytest.mark.covers("other.a2a.register.persists") + def test_register_persists(self, client: A2AClient, resources: ResourceManager) -> None: + agent = _register(client, resources, "0.3") + fetched = unwrap(client.get_agent(agent.agent_id)) + assert fetched.agent_id == agent.agent_id + assert fetched.agent_name == agent.agent_name + assert fetched.agent_card_params.protocol_version == "0.3" + + @pytest.mark.covers("other.a2a.discovery.proxy_fronted_card") + def test_discovery_card_is_proxy_fronted(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3") + card = unwrap(client.agent_card(agent.agent_id, scoped_key)) + assert card.url is not None and card.url.endswith(f"/a2a/{agent.agent_id}") + assert card.security_schemes is not None + scheme = next(iter(card.security_schemes.values())) + assert scheme.scheme == "bearer" + assert card.supported_interfaces is not None + assert card.supported_interfaces[0].url == card.url + + @pytest.mark.covers("other.a2a.message_send.bridge_invokes") + def test_message_send_runs_completion_bridge(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3") + request = _ask("Reply with exactly the word PONG and nothing else") + response = unwrap(client.send_message(agent.agent_id, scoped_key, request)) + assert response.error is None + assert response.result is not None + assert "PONG" in response.result.text.upper() + + rows = client.proxy.poll_logs_for_request_id(request.id) + assert rows, f"no spend log row landed for a2a request {request.id}" + assert rows[0].call_type == "asend_message" + assert rows[0].model == f"a2a_agent/{agent.agent_card_params.name}" + + @pytest.mark.covers("other.a2a.version.serves_pinned_0_3") + def test_pinned_v0_3_serves_flat_message_shape(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3") + request = _ask("Say hi in one word") + result = unwrap(client.send_message(agent.agent_id, scoped_key, request)).result + assert result is not None + assert not result.is_nested_v1_shape + assert result.kind == "message" + assert result.role == "agent" + assert result.text != "" + + @pytest.mark.covers("other.a2a.version.serves_pinned_1_0") + def test_pinned_v1_0_serves_nested_message_shape(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "1.0") + request = _ask("Say hi in one word") + result = unwrap(client.send_message(agent.agent_id, scoped_key, request)).result + assert result is not None + assert result.is_nested_v1_shape + assert result.message is not None + assert result.message.role == "ROLE_AGENT" + assert result.text != "" + + @pytest.mark.covers("other.a2a.register.unsupported_version_rejected") + def test_unsupported_protocol_version_rejected(self, client: A2AClient, resources: ResourceManager) -> None: + marker = unique_marker() + body = AgentRegisterBody( + agent_name=f"e2e-a2a-bad-{marker}", + agent_card_params=AgentCardParams( + protocol_version="9.9", + name=f"E2E A2A bad {marker}", + description="unsupported version", + version="1.0.0", + skills=[A2ASkill(id="chat", name="Chat", description="c", tags=["chat"])], + ), + litellm_params=BRIDGE, + ) + result = client.register_agent(body) + match result: + case UnknownApiError(status_code=status, body=detail): + assert status == 400 + assert "protocolVersion" in detail + case _: + pytest.fail(f"expected 400 for unsupported protocolVersion, got {result}") diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 6b183cbf9f3..f4d0120e085 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -28,3 +28,9 @@ - {id: other.config.overrides.audit_logged, module: other, tier: P1, area: config, assertions: [audit_logged], source: "config_override_endpoints.py:67-100", rationale: "Config override mutations audit-logged, values redacted"} - {id: other.key_mgmt.regenerate.grace_period_honored, module: other, tier: P1, area: auth, assertions: [grace_period_honored], source: "key_management_endpoints.py:4503-4560", rationale: "Old key valid during grace_period then revoked"} - {id: other.key_mgmt.spend_reset.resets_to_value, module: other, tier: P1, area: auth, assertions: [resets_to_value], source: "key_management_endpoints.py:4841", rationale: "reset_spend resets accumulated spend"} +- {id: other.a2a.register.persists, module: other, tier: P1, area: a2a, assertions: [persists], source: "agent_endpoints/endpoints.py:325-443", rationale: "POST /v1/agents registers an agent card; GET /v1/agents/{id} reads it back"} +- {id: other.a2a.register.unsupported_version_rejected, module: other, tier: P1, area: a2a, assertions: [unsupported_version_rejected], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a protocolVersion outside SUPPORTED_A2A_PROTOCOL_VERSIONS is refused with 400"} +- {id: other.a2a.discovery.proxy_fronted_card, module: other, tier: P1, area: a2a, assertions: [proxy_fronted_card], source: "agent_endpoints/a2a_endpoints.py get_agent_card", rationale: "/.well-known/agent-card.json serves the proxy url + supportedInterfaces and the LiteLLM virtual-key bearer scheme, not the upstream"} +- {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} +- {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} +- {id: other.a2a.version.serves_pinned_1_0, module: other, tier: P1, area: a2a, assertions: [serves_pinned_1_0], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 1.0 returns the nested 1.0 message shape (result.message with ROLE_AGENT)"} From 48572c9516abc0585c3133aa02254e4b06706f01 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 20 Jul 2026 19:20:34 -0700 Subject: [PATCH 02/84] 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 03/84] 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 04/84] 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 3a55dda7ec5c78f87e3289298c6221fe9e1d06e3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Jul 2026 22:22:28 -0700 Subject: [PATCH 05/84] refactor(ui): migrate memory table onto shared DataTable Move the admin dashboard Memory table off the hand-rolled antd onto the shared DataTable and cell library, matching the pattern already used by Teams, Virtual Keys, and Guardrails. MemoryView keeps the data (server useQuery, mutations) and owns the detail drawer, edit modal, and delete modal; it now renders a thin MemoryTable consumer plus a getMemoryTableColumns columns file. The server pagination moves the full PaginationState up to the parent so the shared footer's rows-per-page selector works, the key-prefix search runs through the shared toolbar and resets the page on change, and per-row view/edit/delete collapse into a single overflow menu. Sorting stays off since the backend returns updated_at DESC. The old page-reset effect is gone (the page now resets inside the search handler), so its react-hooks/set-state-in-effect suppression is pruned. The detail drawer moves into its own MemoryDetailDrawer component to keep the parent under the complexity budget. --- ui/litellm-dashboard/eslint-suppressions.json | 5 - .../memory/_components/MemoryDetailDrawer.tsx | 114 ++++++ .../memory/_components/MemoryTable.test.tsx | 144 ++++++++ .../memory/_components/MemoryTable.tsx | 94 +++++ .../memory/_components/MemoryTableColumns.tsx | 150 ++++++++ .../memory/_components/MemoryView.test.tsx | 45 +++ .../memory/_components/MemoryView.tsx | 333 ++++-------------- 7 files changed, 606 insertions(+), 279 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryDetailDrawer.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTableColumns.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index f935af8907d..b8e9668b21a 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -639,11 +639,6 @@ "count": 2 } }, - "src/app/(dashboard)/memory/_components/MemoryView.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx": { "no-restricted-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryDetailDrawer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryDetailDrawer.tsx new file mode 100644 index 00000000000..970e088ec00 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryDetailDrawer.tsx @@ -0,0 +1,114 @@ +"use client"; + +import { Drawer, Space, Typography } from "antd"; +import React from "react"; + +import { MemoryRow } from "@/components/networking"; + +const { Text, Paragraph } = Typography; + +interface MemoryDetailDrawerProps { + row: MemoryRow | null; + onClose: () => void; +} + +function formatTimestamp(ts?: string): string { + if (!ts) return "—"; + try { + const d = new Date(ts); + return d.toLocaleString(); + } catch { + return ts; + } +} + +export function MemoryDetailDrawer({ row, onClose }: MemoryDetailDrawerProps) { + return ( + + {row.key} + + ) : ( + "Memory" + ) + } + width={720} + destroyOnClose + > + {row && ( + + +
+ + Memory ID + + + {row.memory_id} + +
+
+ + User ID + + {row.user_id ?? "-"} +
+
+ + Team ID + + {row.team_id ?? "-"} +
+
+
+ Value + + {row.value} + +
+ {row.metadata !== undefined && row.metadata !== null && ( +
+ Metadata + + {JSON.stringify(row.metadata, null, 2)} + +
+ )} + ·} wrap size="small" style={{ color: "rgba(0,0,0,0.45)" }}> + + Created {formatTimestamp(row.created_at)} + {row.created_by ? ` by ${row.created_by}` : ""} + + + Updated {formatTimestamp(row.updated_at)} + {row.updated_by ? ` by ${row.updated_by}` : ""} + + +
+ )} +
+ ); +} + +export default MemoryDetailDrawer; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx new file mode 100644 index 00000000000..0dcbd6796a8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx @@ -0,0 +1,144 @@ +import { PaginationState } from "@tanstack/react-table"; +import { render, screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import { MemoryRow } from "@/components/networking"; + +import { MemoryTable } from "./MemoryTable"; + +const makeMemory = (overrides: Partial = {}): MemoryRow => ({ + memory_id: "mem-1", + key: "user:profile", + value: "The user prefers concise answers.", + metadata: null, + user_id: "user-42", + team_id: "team-7", + updated_at: "2024-05-01T12:00:00Z", + ...overrides, +}); + +const baseProps = { + data: [makeMemory()], + isLoading: false, + rowCount: 1, + pagination: { pageIndex: 0, pageSize: 50 } as PaginationState, + onPaginationChange: vi.fn(), + searchValue: "", + onSearchChange: vi.fn(), + isRefreshing: false, + onRefresh: vi.fn(), + hasActiveSearch: false, + onViewClick: vi.fn(), + onEditClick: vi.fn(), + onDeleteClick: vi.fn(), +}; + +describe("MemoryTable", () => { + it("renders every column header", () => { + render(); + for (const header of ["ID", "Name", "Preview", "User ID", "Team ID", "Updated"]) { + expect(screen.getByText(header)).toBeInTheDocument(); + } + }); + + it("opens the detail view when the ID identity cell is clicked", async () => { + const user = userEvent.setup(); + const onViewClick = vi.fn(); + const row = makeMemory({ memory_id: "mem-click" }); + render(); + + await user.click(screen.getByText("mem-click")); + + expect(onViewClick).toHaveBeenCalledTimes(1); + expect(onViewClick).toHaveBeenCalledWith(row); + }); + + it("routes each overflow-menu action to its callback with the row", async () => { + const user = userEvent.setup(); + const onViewClick = vi.fn(); + const onEditClick = vi.fn(); + const onDeleteClick = vi.fn(); + const row = makeMemory({ memory_id: "mem-9" }); + render( + , + ); + + await user.click(screen.getByTestId("memory-actions-mem-9")); + await user.click(await screen.findByTestId("memory-action-edit")); + expect(onEditClick).toHaveBeenCalledWith(row); + expect(onViewClick).not.toHaveBeenCalled(); + expect(onDeleteClick).not.toHaveBeenCalled(); + + await user.click(screen.getByTestId("memory-actions-mem-9")); + await user.click(await screen.findByTestId("memory-action-delete")); + expect(onDeleteClick).toHaveBeenCalledWith(row); + + await user.click(screen.getByTestId("memory-actions-mem-9")); + await user.click(await screen.findByTestId("memory-action-view")); + expect(onViewClick).toHaveBeenCalledWith(row); + }); + + it("shows the empty-only copy when there is no data and no active search", () => { + render(); + expect(screen.getByText("No memories stored yet")).toBeInTheDocument(); + expect(screen.queryByText("No matching memories")).not.toBeInTheDocument(); + }); + + it("shows the filtered-empty copy when a search is active", () => { + render(); + expect(screen.getByText("No matching memories")).toBeInTheDocument(); + expect(screen.queryByText("No memories stored yet")).not.toBeInTheDocument(); + }); + + it("renders loading skeleton rows instead of the empty state while loading", () => { + render(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); + expect(screen.queryByText("No memories stored yet")).not.toBeInTheDocument(); + }); + + it("drives the pagination footer from the server rowCount, not the page's row length", () => { + render(); + const range = screen.getByTestId("pagination-range"); + expect(range).toHaveTextContent("Showing 1-50 of 120"); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 3"); + expect(screen.getByTestId("pagination-next")).toBeEnabled(); + }); + + it("advances the page through the server pagination handler", async () => { + const user = userEvent.setup(); + const onPaginationChange = vi.fn(); + render(); + + await user.click(screen.getByTestId("pagination-next")); + + expect(onPaginationChange).toHaveBeenCalled(); + }); + + it("forwards toolbar search input and refresh to their callbacks", async () => { + const user = userEvent.setup(); + const onSearchChange = vi.fn(); + const onRefresh = vi.fn(); + render(); + + await user.type(screen.getByTestId("datatable-search"), "u"); + expect(onSearchChange).toHaveBeenCalledWith("u"); + + await user.click(screen.getByTestId("datatable-refresh")); + expect(onRefresh).toHaveBeenCalledTimes(1); + }); + + it("renders secondary id and date cells for the row", () => { + render(); + const table = screen.getByRole("table"); + expect(within(table).getByText("user-42")).toBeInTheDocument(); + expect(within(table).getByText("team-7")).toBeInTheDocument(); + expect(within(table).getByText("user:profile")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx new file mode 100644 index 00000000000..50dd04ee14c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx @@ -0,0 +1,94 @@ +"use client"; + +import { OnChangeFn, PaginationState } from "@tanstack/react-table"; +import { Database } from "lucide-react"; +import React, { useMemo } from "react"; + +import { MemoryRow } from "@/components/networking"; +import { DataTable, DataTableToolbar } from "@/components/shared/DataTable"; + +import { getMemoryTableColumns } from "./MemoryTableColumns"; + +interface MemoryTableProps { + data: MemoryRow[]; + isLoading: boolean; + rowCount: number; + pagination: PaginationState; + onPaginationChange: OnChangeFn; + searchValue: string; + onSearchChange: (value: string) => void; + isRefreshing: boolean; + onRefresh: () => void; + hasActiveSearch: boolean; + onViewClick: (row: MemoryRow) => void; + onEditClick: (row: MemoryRow) => void; + onDeleteClick: (row: MemoryRow) => void; +} + +function MemoryEmptyState({ hasActiveSearch }: { hasActiveSearch: boolean }) { + return ( +
+
+ +
+
+ {hasActiveSearch ? "No matching memories" : "No memories stored yet"} +
+
+ {hasActiveSearch + ? "No memories have keys starting with your search." + : "Memories your agents store under /v1/memory will appear here."} +
+
+ ); +} + +export function MemoryTable({ + data, + isLoading, + rowCount, + pagination, + onPaginationChange, + searchValue, + onSearchChange, + isRefreshing, + onRefresh, + hasActiveSearch, + onViewClick, + onEditClick, + onDeleteClick, +}: MemoryTableProps) { + const columns = useMemo(() => { + const columnDeps = { onViewClick, onEditClick, onDeleteClick }; + return getMemoryTableColumns(columnDeps); + }, [onViewClick, onEditClick, onDeleteClick]); + + return ( + row.memory_id} + paginationMode="server" + pagination={pagination} + onPaginationChange={onPaginationChange} + rowCount={rowCount} + isLoading={isLoading} + loadingMessage="Loading memories…" + noDataMessage={} + size="compact" + toolbar={(table) => ( + + )} + /> + ); +} + +export default MemoryTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTableColumns.tsx new file mode 100644 index 00000000000..6b2a6b08704 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTableColumns.tsx @@ -0,0 +1,150 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { Eye, MoreHorizontal, Pencil, Trash2 } from "lucide-react"; + +import { MemoryRow } from "@/components/networking"; +import { DateCell, IdCell, IdentityCell } from "@/components/shared/table_cells"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; + +interface MemoryRowActionsProps { + row: MemoryRow; + onViewClick: (row: MemoryRow) => void; + onEditClick: (row: MemoryRow) => void; + onDeleteClick: (row: MemoryRow) => void; +} + +function MemoryRowActions({ row, onViewClick, onEditClick, onDeleteClick }: MemoryRowActionsProps) { + return ( + + + + + + onViewClick(row)}> + + View + + onEditClick(row)}> + + Edit + + + onDeleteClick(row)}> + + Delete + + + + ); +} + +export interface MemoryTableColumnsDeps { + onViewClick: (row: MemoryRow) => void; + onEditClick: (row: MemoryRow) => void; + onDeleteClick: (row: MemoryRow) => void; +} + +export const getMemoryTableColumns = ({ + onViewClick, + onEditClick, + onDeleteClick, +}: MemoryTableColumnsDeps): ColumnDef[] => [ + { + id: "memory_id", + accessorKey: "memory_id", + meta: { title: "ID" }, + header: "ID", + size: 180, + enableSorting: false, + cell: ({ row }) => ( + onViewClick(row.original)} + /> + ), + }, + { + id: "key", + accessorKey: "key", + meta: { title: "Name" }, + header: "Name", + size: 200, + enableSorting: false, + cell: ({ row }) => ( + + {row.original.key} + + ), + }, + { + id: "value", + accessorKey: "value", + meta: { title: "Preview" }, + header: "Preview", + enableSorting: false, + cell: ({ row }) => ( + + {row.original.value || "-"} + + ), + }, + { + id: "user_id", + accessorKey: "user_id", + meta: { title: "User ID" }, + header: "User ID", + size: 160, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "team_id", + accessorKey: "team_id", + meta: { title: "Team ID" }, + header: "Team ID", + size: 160, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "updated_at", + accessorKey: "updated_at", + meta: { title: "Updated" }, + header: "Updated", + size: 170, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + }, +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx new file mode 100644 index 00000000000..f415c99225a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx @@ -0,0 +1,45 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render } from "@testing-library/react"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; + +import { MemoryRow } from "@/components/networking"; + +import { MemoryView } from "./MemoryView"; + +interface CapturedTableProps { + isLoading: boolean; + rowCount: number; + data: MemoryRow[]; + hasActiveSearch: boolean; +} + +const captured = vi.hoisted(() => ({ current: null as CapturedTableProps | null })); + +vi.mock("./MemoryTable", () => ({ + MemoryTable: function MemoryTableMock(props: CapturedTableProps) { + captured.current = props; + return
; + }, +})); + +const renderView = (accessToken: string | null) => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render( + + + , + ); +}; + +describe("MemoryView", () => { + it("keeps the table out of the skeleton state when the token is null (disabled query)", () => { + renderView(null); + + expect(captured.current).not.toBeNull(); + expect(captured.current?.isLoading).toBe(false); + expect(captured.current?.data).toEqual([]); + expect(captured.current?.rowCount).toBe(0); + expect(captured.current?.hasActiveSearch).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx index 4ee784f4664..fcb15978f47 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx @@ -1,21 +1,19 @@ "use client"; -import React, { useMemo, useState } from "react"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; -import { Button, Card, Drawer, Empty, Input, Space, Table, Typography, message } from "antd"; -import type { ColumnsType } from "antd/es/table"; -import { - DeleteOutlined, - EditOutlined, - EyeOutlined, - PlusOutlined, - ReloadOutlined, - SearchOutlined, -} from "@ant-design/icons"; +import type { PaginationState } from "@tanstack/react-table"; +import { PlusOutlined } from "@ant-design/icons"; +import { Button, Space, Typography, message } from "antd"; +import React, { useCallback, useMemo, useState } from "react"; + import { MemoryRow, createMemory, deleteMemory, fetchMemoryList, updateMemory } from "@/components/networking"; -import { DateCell, IdCell } from "@/components/shared/table_cells"; -import { MemoryEditModal } from "./MemoryEditModal"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; + +import { MemoryDetailDrawer } from "./MemoryDetailDrawer"; +import { MemoryEditModal } from "./MemoryEditModal"; +import { MemoryTable } from "./MemoryTable"; const { Text, Paragraph, Title } = Typography; @@ -25,38 +23,16 @@ interface MemoryViewProps { userRole: string | null; } -function previewValue(value: string, max = 120): string { - if (!value) return ""; - const trimmed = value.trim(); - if (trimmed.length <= max) return trimmed; - return `${trimmed.slice(0, max)}…`; -} - -function formatTimestamp(ts?: string): string { - if (!ts) return "—"; - try { - const d = new Date(ts); - return d.toLocaleString(); - } catch { - return ts; - } -} - -const PAGE_SIZE = 50; +const DEFAULT_PAGE_SIZE = 50; export const MemoryView: React.FC = ({ accessToken }) => { const [searchInput, setSearchInput] = useState(""); - const [appliedSearch, setAppliedSearch] = useState(""); + const [debouncedSearch] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); + const [pagination, setPagination] = useState({ pageIndex: 0, pageSize: DEFAULT_PAGE_SIZE }); const [detailRow, setDetailRow] = useState(null); const [editRow, setEditRow] = useState(null); const [deleteRow, setDeleteRow] = useState(null); const [isCreateOpen, setIsCreateOpen] = useState(false); - const [currentPage, setCurrentPage] = useState(1); - - // Reset to page 1 whenever the filter changes. - React.useEffect(() => { - setCurrentPage(1); - }, [appliedSearch]); const queryClient = useQueryClient(); // React Query key prefix for all memory-list variants (paged + filtered). @@ -65,15 +41,15 @@ export const MemoryView: React.FC = ({ accessToken }) => { const MEMORY_LIST_KEY = "memoryList" as const; const { data, isLoading, isFetching } = useQuery({ - queryKey: [MEMORY_LIST_KEY, appliedSearch, currentPage], + queryKey: [MEMORY_LIST_KEY, debouncedSearch, pagination.pageIndex, pagination.pageSize], queryFn: () => { if (!accessToken) throw new Error("Access token required"); // Prefix search matches the Redis-style mental model (namespace scan): // typing "user:" finds "user:profile", "user:prefs", etc. return fetchMemoryList(accessToken, { - keyPrefix: appliedSearch || undefined, - page: currentPage, - pageSize: PAGE_SIZE, + keyPrefix: debouncedSearch || undefined, + page: pagination.pageIndex + 1, + pageSize: pagination.pageSize, }); }, enabled: !!accessToken, @@ -88,7 +64,10 @@ export const MemoryView: React.FC = ({ accessToken }) => { // refetches from scratch (pagination + filter-aware). // - on error: surface the message via antd `message.error`. - const invalidateList = () => queryClient.invalidateQueries({ queryKey: [MEMORY_LIST_KEY] }); + const invalidateList = useCallback( + () => queryClient.invalidateQueries({ queryKey: [MEMORY_LIST_KEY] }), + [queryClient], + ); const createMutation = useMutation({ mutationFn: (args: { key: string; value: string; metadata: unknown }) => { @@ -133,9 +112,14 @@ export const MemoryView: React.FC = ({ accessToken }) => { }, }); - const handleDelete = (row: MemoryRow) => { - setDeleteRow(row); - }; + const handleSearchChange = useCallback((value: string) => { + setSearchInput(value); + setPagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); + + const handleView = useCallback((row: MemoryRow) => setDetailRow(row), []); + const handleEdit = useCallback((row: MemoryRow) => setEditRow(row), []); + const handleDelete = useCallback((row: MemoryRow) => setDeleteRow(row), []); const confirmDelete = async () => { if (!deleteRow) return; @@ -192,242 +176,43 @@ export const MemoryView: React.FC = ({ accessToken }) => { } }; - const columns: ColumnsType = [ - { - title: "ID", - dataIndex: "memory_id", - key: "memory_id", - width: 140, - render: (_: unknown, r: MemoryRow) => setDetailRow(r)} />, - }, - { - title: "Name", - dataIndex: "key", - key: "key", - width: 200, - render: (k: string) => {k}, - // No client-side sorter: pagination is server-side, so a client sort - // would only reorder the current page and mislead users into thinking - // the whole list is sorted. Backend returns rows ordered by - // `updated_at DESC`; use the prefix filter for discovery by name. - }, - { - title: "Preview", - dataIndex: "value", - key: "value", - render: (v: string) => ( - - {previewValue(v)} - - ), - }, - { - title: "User ID", - dataIndex: "user_id", - key: "user_id", - width: 160, - render: (uid?: string | null) => , - }, - { - title: "Team ID", - dataIndex: "team_id", - key: "team_id", - width: 160, - render: (tid?: string | null) => , - }, - { - title: "Updated", - dataIndex: "updated_at", - key: "updated_at", - width: 180, - render: (ts?: string) => , - // No sorter — backend already returns rows in `updated_at DESC` order, - // and a client-side sorter on a paginated view would only affect the - // current page. - }, - { - title: "", - key: "actions", - width: 140, - render: (_: unknown, r: MemoryRow) => ( - -
- - - - } - value={searchInput} - onChange={(e) => setSearchInput(e.target.value)} - onPressEnter={() => setAppliedSearch(searchInput.trim())} - onClear={() => { - setSearchInput(""); - setAppliedSearch(""); - }} - style={{ width: 280 }} - /> - - - - - - -
`${range[0]}–${range[1]} of ${n}`, - onChange: (page) => setCurrentPage(page), - }} - locale={{ - emptyText: ( - - ), - }} - /> - + {/* Detail drawer */} - setDetailRow(null)} - title={ - detailRow ? ( - - {detailRow.key} - - ) : ( - "Memory" - ) - } - width={720} - destroyOnClose - > - {detailRow && ( - - -
- - Memory ID - - - {detailRow.memory_id} - -
-
- - User ID - - {detailRow.user_id ?? "-"} -
-
- - Team ID - - {detailRow.team_id ?? "-"} -
-
-
- Value - - {detailRow.value} - -
- {detailRow.metadata !== undefined && detailRow.metadata !== null && ( -
- Metadata - - {JSON.stringify(detailRow.metadata, null, 2)} - -
- )} - ·} wrap size="small" style={{ color: "rgba(0,0,0,0.45)" }}> - - Created {formatTimestamp(detailRow.created_at)} - {detailRow.created_by ? ` by ${detailRow.created_by}` : ""} - - - Updated {formatTimestamp(detailRow.updated_at)} - {detailRow.updated_by ? ` by ${detailRow.updated_by}` : ""} - - -
- )} -
+ setDetailRow(null)} /> {/* Create / edit modal */} Date: Mon, 20 Jul 2026 22:23:54 -0700 Subject: [PATCH 06/84] refactor(ui): migrate audit logs table onto shared DataTable Move the Audit Logs table off the hand-rolled antd Table/Pagination onto the shared DataTable and cell library, matching the other migrated admin tables (Teams, Virtual Keys, Guardrails) The single audit_logs.tsx is split into three PascalCase files: AuditLogsPanel owns the data (server useQuery, pagination and filter state, the row-detail drawer, and the enterprise preview gate), AuditLogsTable is a thin DataTable consumer, and AuditLogsTableColumns exposes getAuditLogsTableColumns. The AuditLogEntry type moves out of the request-logs columns.tsx into the audit columns file, and AuditLogDrawer stays in the parent unchanged Server pagination is wired through paginationMode="server" with the shared footer replacing the standalone antd Pagination, keeping keepPreviousData semantics so page flips keep rows visible and only the initial load shows the skeleton. The six filters (Object ID, Changed By, Team ID, Key Hash, Action, Table) move into a DataTableFilterDrawer plus toolbar with active-filter chips, each resetting the page to the first. The Object ID cell is the clickable identity cell that opens the drawer; there is no whole-row navigation, no selection, and no per-row actions since the table is read-only The enterprise query is now also gated on premiumUser so the preview path no longer fires a doomed request for non-premium users --- .../AuditLogDrawer/AuditLogDrawer.tsx | 2 +- .../components/view_logs/AuditLogsPanel.tsx | 138 ++++++++ .../view_logs/AuditLogsTable.test.tsx | 146 ++++++++ .../components/view_logs/AuditLogsTable.tsx | 208 ++++++++++++ .../view_logs/AuditLogsTableColumns.tsx | 102 ++++++ .../src/components/view_logs/audit_logs.tsx | 315 ------------------ .../src/components/view_logs/columns.tsx | 12 - .../src/components/view_logs/index.tsx | 4 +- 8 files changed, 597 insertions(+), 330 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx create mode 100644 ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx create mode 100644 ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx create mode 100644 ui/litellm-dashboard/src/components/view_logs/AuditLogsTableColumns.tsx delete mode 100644 ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx index aa690787f66..81759c80ab6 100644 --- a/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogDrawer/AuditLogDrawer.tsx @@ -2,7 +2,7 @@ import { Drawer, Tag, Typography } from "antd"; import { CloseOutlined, CopyOutlined, CheckOutlined } from "@ant-design/icons"; import { useState, useCallback } from "react"; import moment from "moment"; -import { AuditLogEntry } from "../columns"; +import { AuditLogEntry } from "../AuditLogsTableColumns"; import DefaultProxyAdminTag from "../../common_components/DefaultProxyAdminTag"; const { Text } = Typography; diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx new file mode 100644 index 00000000000..81bd4a19f76 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx @@ -0,0 +1,138 @@ +import { useCallback, useState } from "react"; +import { useQuery, keepPreviousData } from "@tanstack/react-query"; +import { ColumnFiltersState, OnChangeFn, PaginationState } from "@tanstack/react-table"; +import { resolveLogoSrc } from "@/lib/assetPaths"; +import { uiAuditLogsCall } from "../networking"; +import { AuditLogEntry } from "./AuditLogsTableColumns"; +import { AuditLogsTable } from "./AuditLogsTable"; +import { AuditLogDrawer } from "./AuditLogDrawer/AuditLogDrawer"; + +interface AuditLogsProps { + accessToken: string | null; + token: string | null; + userRole: string | null; + userID: string | null; + isActive: boolean; + premiumUser: boolean; +} + +const asset_logos_folder = "/ui/assets/"; +const auditLogsPreviewImg = `${asset_logos_folder}audit-logs-preview.png`; + +const PAGE_SIZE = 50; + +interface AuditLogsResponse { + audit_logs: AuditLogEntry[]; + total: number; + page: number; + page_size: number; + total_pages: number; +} + +export default function AuditLogsPanel({ + userID, + userRole, + token, + accessToken, + isActive, + premiumUser, +}: AuditLogsProps) { + const [pagination, setPagination] = useState({ pageIndex: 0, pageSize: PAGE_SIZE }); + const [columnFilters, setColumnFilters] = useState([]); + const [selectedLog, setSelectedLog] = useState(null); + const [drawerOpen, setDrawerOpen] = useState(false); + + const getFilterValue = (columnId: string): string | undefined => { + const entry = columnFilters.find((filter) => filter.id === columnId); + return typeof entry?.value === "string" && entry.value.trim() ? entry.value.trim() : undefined; + }; + + const canQueryAuditLogs = !!accessToken && !!token && !!userRole && !!userID && isActive && premiumUser; + + const query = useQuery({ + queryKey: ["audit_logs", pagination.pageIndex, pagination.pageSize, columnFilters], + queryFn: async () => { + if (!accessToken) { + return { audit_logs: [], total: 0, page: 1, page_size: pagination.pageSize, total_pages: 0 }; + } + return uiAuditLogsCall({ + accessToken, + page: pagination.pageIndex + 1, + page_size: pagination.pageSize, + params: { + object_id: getFilterValue("object_id"), + changed_by: getFilterValue("changed_by"), + object_key_hash: getFilterValue("key_hash"), + object_team_id: getFilterValue("team_id"), + action: getFilterValue("action"), + table_name: getFilterValue("table_name"), + sort_by: "updated_at", + sort_order: "desc", + }, + }); + }, + enabled: canQueryAuditLogs, + placeholderData: keepPreviousData, + }); + + const handleColumnFiltersChange = useCallback>((updaterOrValue) => { + setColumnFilters(updaterOrValue); + setPagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); + + const handleViewLog = useCallback((log: AuditLogEntry) => { + setSelectedLog(log); + setDrawerOpen(true); + }, []); + + if (!premiumUser) { + return ( +
+

✨ Enterprise Feature.

+

+ This is a LiteLLM Enterprise feature, and requires a valid key to use. +

+

+ Here's a preview of what Audit Logs offer: +

+ Audit Logs Preview { + (e.target as HTMLImageElement).style.display = "none"; + }} + /> +
+ ); + } + + return ( + <> +
+

Audit Logs

+
+ + query.refetch()} + onViewLog={handleViewLog} + /> + + setDrawerOpen(false)} log={selectedLog} /> + + ); +} diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx new file mode 100644 index 00000000000..dbb0a39e2ee --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx @@ -0,0 +1,146 @@ +import type { ColumnFiltersState, PaginationState } from "@tanstack/react-table"; +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import { AuditLogsTable } from "./AuditLogsTable"; +import type { AuditLogEntry } from "./AuditLogsTableColumns"; + +const ROWS: AuditLogEntry[] = [ + { + id: "log-1", + updated_at: "2026-07-20T12:00:00Z", + changed_by: "default_user_id", + changed_by_api_key: "sk-hash-abc", + action: "created", + table_name: "LiteLLM_TeamTable", + object_id: "team-obj-123", + before_value: {}, + updated_values: { foo: "bar" }, + }, + { + id: "log-2", + updated_at: "2026-07-20T11:00:00Z", + changed_by: "user-42", + changed_by_api_key: "sk-hash-def", + action: "deleted", + table_name: "LiteLLM_UserTable", + object_id: "user-obj-456", + before_value: { a: 1 }, + updated_values: {}, + }, +]; + +const FIRST_PAGE: PaginationState = { pageIndex: 0, pageSize: 50 }; + +function renderTable(overrides: Partial> = {}) { + const props: React.ComponentProps = { + data: ROWS, + rowCount: ROWS.length, + isLoading: false, + isRefreshing: false, + pagination: FIRST_PAGE, + onPaginationChange: vi.fn(), + columnFilters: [], + onColumnFiltersChange: vi.fn(), + onRefresh: vi.fn(), + onViewLog: vi.fn(), + ...overrides, + }; + render(); + return props; +} + +describe("AuditLogsTable", () => { + it("renders each audit column with the migrated shared cells", () => { + renderTable(); + + // Action -> StatusBadge with a capitalized label + expect(screen.getByText("Created")).toBeInTheDocument(); + expect(screen.getByText("Deleted")).toBeInTheDocument(); + // Table name -> display mapping + expect(screen.getByText("Teams")).toBeInTheDocument(); + expect(screen.getByText("Users")).toBeInTheDocument(); + // Changed By -> DefaultProxyAdminTag (default_user_id becomes a labeled tag; other ids stay raw) + expect(screen.getByText("Default Proxy Admin")).toBeInTheDocument(); + expect(screen.getByText("user-42")).toBeInTheDocument(); + // Object ID + API key hash + expect(screen.getByText("team-obj-123")).toBeInTheDocument(); + expect(screen.getByText("sk-hash-abc")).toBeInTheDocument(); + }); + + it("opens the detail drawer from the Object ID identity cell with the full row", async () => { + const user = userEvent.setup(); + const props = renderTable(); + + await user.click(screen.getByText("team-obj-123")); + + expect(props.onViewLog).toHaveBeenCalledTimes(1); + expect(props.onViewLog).toHaveBeenCalledWith(ROWS[0]); + }); + + it("drives the shared footer from the server rowCount and reports page changes", async () => { + const user = userEvent.setup(); + const onPaginationChange = vi.fn(); + renderTable({ rowCount: 120, onPaginationChange }); + + // ceil(120 / 50) = 3 pages, proving rowCount (not data length) feeds the footer + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 3"); + + await user.click(screen.getByTestId("pagination-next")); + expect(onPaginationChange).toHaveBeenCalledTimes(1); + }); + + it("shows skeleton rows while loading and no data rows", () => { + renderTable({ isLoading: true, data: [] }); + + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); + expect(screen.queryByText("No audit logs yet")).toBeNull(); + }); + + it("uses a distinct empty state for unfiltered vs filtered-empty results", () => { + const { unmount } = render( + , + ); + expect(screen.getByText("No audit logs yet")).toBeInTheDocument(); + unmount(); + + renderTable({ data: [], rowCount: 0, columnFilters: [{ id: "action", value: "created" }] }); + expect(screen.getByText("No matching audit logs")).toBeInTheDocument(); + }); + + it("renders active filter chips with human-readable labels", () => { + const filters: ColumnFiltersState = [{ id: "action", value: "created" }]; + renderTable({ columnFilters: filters }); + + const chip = screen.getByTestId("filter-chip-action"); + expect(chip).toHaveTextContent("Action:"); + expect(chip).toHaveTextContent("Created"); + }); + + it("commits a text filter through the filter drawer and reports it to the parent", async () => { + const user = userEvent.setup(); + const onColumnFiltersChange = vi.fn(); + renderTable({ onColumnFiltersChange }); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + await user.type(await screen.findByPlaceholderText("Enter object ID…"), "obj-9"); + await user.click(screen.getByTestId("filter-drawer-apply")); + + expect(onColumnFiltersChange).toHaveBeenCalledTimes(1); + const arg = onColumnFiltersChange.mock.calls[0][0]; + const committed = typeof arg === "function" ? arg([]) : arg; + expect(committed).toEqual([{ id: "object_id", value: "obj-9" }]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx new file mode 100644 index 00000000000..bcdce12fce8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx @@ -0,0 +1,208 @@ +"use client"; + +import { ColumnFiltersState, OnChangeFn, PaginationState } from "@tanstack/react-table"; +import { ScrollText } from "lucide-react"; +import { useMemo, useState } from "react"; + +import { + DataTable, + DataTableFilterDrawer, + DataTableFilterField, + DataTableToolbar, +} from "@/components/shared/DataTable"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; + +import { AUDIT_TABLE_NAME_DISPLAY, AuditLogEntry, getAuditLogsTableColumns } from "./AuditLogsTableColumns"; + +interface AuditLogsTableProps { + data: AuditLogEntry[]; + rowCount: number; + isLoading: boolean; + isRefreshing: boolean; + pagination: PaginationState; + onPaginationChange: OnChangeFn; + columnFilters: ColumnFiltersState; + onColumnFiltersChange: OnChangeFn; + onRefresh: () => void; + onViewLog: (log: AuditLogEntry) => void; +} + +const ALL_VALUE = "all"; + +const ACTION_OPTIONS = [ + { label: "Created", value: "created" }, + { label: "Updated", value: "updated" }, + { label: "Deleted", value: "deleted" }, + { label: "Rotated", value: "rotated" }, +] as const; + +const TABLE_OPTIONS = [ + { label: "Keys", value: "LiteLLM_VerificationToken" }, + { label: "Teams", value: "LiteLLM_TeamTable" }, + { label: "Users", value: "LiteLLM_UserTable" }, + { label: "Organizations", value: "LiteLLM_OrganizationTable" }, + { label: "Models", value: "LiteLLM_ProxyModelTable" }, +] as const; + +const FILTER_LABELS: Record = { + object_id: "Object ID", + changed_by: "Changed By", + team_id: "Team ID", + key_hash: "Key Hash", + action: "Action", + table_name: "Table", +}; + +const formatFilterValue = (columnId: string, value: unknown): string => { + const raw = String(value); + if (columnId === "action") { + return ACTION_OPTIONS.find((option) => option.value === raw)?.label ?? raw; + } + if (columnId === "table_name") { + return AUDIT_TABLE_NAME_DISPLAY[raw] ?? raw; + } + return raw; +}; + +function AuditLogsEmptyState({ filtered }: { filtered: boolean }) { + return ( +
+
+ +
+
+ {filtered ? "No matching audit logs" : "No audit logs yet"} +
+
+ {filtered + ? "No audit log entries match your filters." + : "Administrative changes to keys, teams, users, and models will appear here."} +
+
+ ); +} + +export function AuditLogsTable({ + data, + rowCount, + isLoading, + isRefreshing, + pagination, + onPaginationChange, + columnFilters, + onColumnFiltersChange, + onRefresh, + onViewLog, +}: AuditLogsTableProps) { + const [filtersOpen, setFiltersOpen] = useState(false); + const columns = useMemo(() => getAuditLogsTableColumns({ onViewLog }), [onViewLog]); + + return ( + row.id} + paginationMode="server" + pagination={pagination} + onPaginationChange={onPaginationChange} + rowCount={rowCount} + filterMode="server" + columnFilters={columnFilters} + onColumnFiltersChange={onColumnFiltersChange} + isLoading={isLoading} + loadingMessage="Loading audit logs…" + noDataMessage={ 0} />} + size="compact" + toolbar={(table) => ( + <> + setFiltersOpen(true)} + filterLabels={FILTER_LABELS} + formatFilterValue={formatFilterValue} + showViewOptions={false} + /> + + {({ get, set }) => ( + <> + + set("object_id", event.target.value)} + placeholder="Enter object ID…" + /> + + + set("changed_by", event.target.value)} + placeholder="Enter user ID…" + /> + + + set("team_id", event.target.value)} + placeholder="Enter team ID…" + /> + + + set("key_hash", event.target.value)} + placeholder="Enter key hash…" + /> + + + + + + + + + )} + + + )} + /> + ); +} diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTableColumns.tsx new file mode 100644 index 00000000000..6910ca1c2f7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTableColumns.tsx @@ -0,0 +1,102 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; + +import { DateCell, IdCell, IdentityCell, StatusBadge, type StatusTone } from "@/components/shared/table_cells"; + +import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; + +export type AuditLogEntry = { + id: string; + updated_at: string; + changed_by: string; + changed_by_api_key: string; + action: string; + table_name: string; + object_id: string; + before_value: Record; + updated_values: Record; +}; + +export const AUDIT_TABLE_NAME_DISPLAY: Record = { + LiteLLM_VerificationToken: "Keys", + LiteLLM_TeamTable: "Teams", + LiteLLM_UserTable: "Users", + LiteLLM_OrganizationTable: "Organizations", + LiteLLM_ProxyModelTable: "Models", +}; + +const ACTION_TONE: Record = { + created: "success", + updated: "info", + deleted: "error", + rotated: "warning", +}; + +const capitalize = (value: string): string => (value ? value.charAt(0).toUpperCase() + value.slice(1) : value); + +interface AuditLogsTableColumnsDeps { + onViewLog: (log: AuditLogEntry) => void; +} + +export const getAuditLogsTableColumns = ({ onViewLog }: AuditLogsTableColumnsDeps): ColumnDef[] => [ + { + id: "updated_at", + accessorKey: "updated_at", + header: "Timestamp", + size: 200, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "action", + accessorKey: "action", + header: "Action", + size: 110, + enableSorting: false, + cell: ({ row }) => ( + + ), + }, + { + id: "table_name", + accessorKey: "table_name", + header: "Table", + size: 130, + enableSorting: false, + cell: ({ row }) => ( + {AUDIT_TABLE_NAME_DISPLAY[row.original.table_name] ?? row.original.table_name} + ), + }, + { + id: "object_id", + accessorKey: "object_id", + header: "Object ID", + minSize: 220, + enableSorting: false, + cell: ({ row }) => ( + onViewLog(row.original)} + /> + ), + }, + { + id: "changed_by", + accessorKey: "changed_by", + header: "Changed By", + size: 200, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "changed_by_api_key", + accessorKey: "changed_by_api_key", + header: "API Key (Hash)", + size: 160, + enableSorting: false, + cell: ({ row }) => , + }, +]; diff --git a/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx b/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx deleted file mode 100644 index d811d3b9402..00000000000 --- a/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx +++ /dev/null @@ -1,315 +0,0 @@ -import { useState } from "react"; -import { useQuery, keepPreviousData } from "@tanstack/react-query"; -import { Table, Tag, Input, Select, Button, Pagination, Spin } from "antd"; -import { ReloadOutlined, LoadingOutlined } from "@ant-design/icons"; -import type { ColumnsType } from "antd/es/table"; -import { resolveLogoSrc } from "@/lib/assetPaths"; -import { DateCell, IdCell } from "@/components/shared/table_cells"; -import { uiAuditLogsCall } from "../networking"; -import { AuditLogEntry } from "./columns"; -import { AuditLogDrawer } from "./AuditLogDrawer/AuditLogDrawer"; -import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; - -const { Search } = Input; - -interface AuditLogsProps { - accessToken: string | null; - token: string | null; - userRole: string | null; - userID: string | null; - isActive: boolean; - premiumUser: boolean; -} - -const asset_logos_folder = "/ui/assets/"; -export const auditLogsPreviewImg = `${asset_logos_folder}audit-logs-preview.png`; - -const TABLE_NAME_DISPLAY: Record = { - LiteLLM_VerificationToken: "Keys", - LiteLLM_TeamTable: "Teams", - LiteLLM_UserTable: "Users", - LiteLLM_OrganizationTable: "Organizations", - LiteLLM_ProxyModelTable: "Models", -}; - -const ACTION_COLOR: Record = { - created: "green", - updated: "blue", - deleted: "red", - rotated: "orange", -}; - -const PAGE_SIZE = 50; - -export default function AuditLogs({ userID, userRole, token, accessToken, isActive, premiumUser }: AuditLogsProps) { - const [page, setPage] = useState(1); - - // Filter state - const [objectId, setObjectId] = useState(""); - const [changedBy, setChangedBy] = useState(""); - const [keyHash, setKeyHash] = useState(""); - const [teamId, setTeamId] = useState(""); - const [action, setAction] = useState(undefined); - const [tableName, setTableName] = useState(undefined); - - // Drawer state - const [selectedLog, setSelectedLog] = useState(null); - const [drawerOpen, setDrawerOpen] = useState(false); - - const query = useQuery({ - queryKey: ["audit_logs", page, PAGE_SIZE, objectId, changedBy, keyHash, teamId, action, tableName], - queryFn: async () => { - if (!accessToken || !token || !userRole || !userID) { - return { audit_logs: [], total: 0, page: 1, page_size: PAGE_SIZE, total_pages: 0 }; - } - return uiAuditLogsCall({ - accessToken, - page, - page_size: PAGE_SIZE, - params: { - object_id: objectId || undefined, - changed_by: changedBy || undefined, - object_key_hash: keyHash || undefined, - object_team_id: teamId || undefined, - action: action || undefined, - table_name: tableName || undefined, - sort_by: "updated_at", - sort_order: "desc", - }, - }); - }, - enabled: !!accessToken && !!token && !!userRole && !!userID && isActive, - placeholderData: keepPreviousData, - }); - - const resetPage = () => setPage(1); - - const handleRowClick = (log: AuditLogEntry) => { - setSelectedLog(log); - setDrawerOpen(true); - }; - - const columns: ColumnsType = [ - { - title: "Timestamp", - dataIndex: "updated_at", - key: "updated_at", - width: 200, - render: (val: string) => , - }, - { - title: "Action", - dataIndex: "action", - key: "action", - width: 100, - render: (val: string) => ( - - {val} - - ), - }, - { - title: "Table", - dataIndex: "table_name", - key: "table_name", - width: 130, - render: (val: string) => TABLE_NAME_DISPLAY[val] ?? val, - }, - { - title: "Object ID", - dataIndex: "object_id", - key: "object_id", - render: (val: string) => , - }, - { - title: "Changed By", - dataIndex: "changed_by", - key: "changed_by", - width: 200, - render: (val: string) => , - }, - { - title: "API Key (Hash)", - dataIndex: "changed_by_api_key", - key: "changed_by_api_key", - width: 140, - render: (val: string) => , - }, - ]; - - if (!premiumUser) { - return ( -
-

✨ Enterprise Feature.

-

- This is a LiteLLM Enterprise feature, and requires a valid key to use. -

-

- Here's a preview of what Audit Logs offer: -

- Audit Logs Preview { - (e.target as HTMLImageElement).style.display = "none"; - }} - /> -
- ); - } - - const auditLogs: AuditLogEntry[] = query.data?.audit_logs ?? []; - const total: number = query.data?.total ?? 0; - - return ( - <> -
- {/* Header */} -
-
-

Audit Logs

-
- - {/* Filters + pagination on same row */} -
- { - setObjectId(val); - resetPage(); - }} - onChange={(e) => { - if (!e.target.value) { - setObjectId(""); - resetPage(); - } - }} - /> - { - setChangedBy(val); - resetPage(); - }} - onChange={(e) => { - if (!e.target.value) { - setChangedBy(""); - resetPage(); - } - }} - /> - { - setTeamId(val); - resetPage(); - }} - onChange={(e) => { - if (!e.target.value) { - setTeamId(""); - resetPage(); - } - }} - /> - { - setKeyHash(val); - resetPage(); - }} - onChange={(e) => { - if (!e.target.value) { - setKeyHash(""); - resetPage(); - } - }} - /> - { - setTableName(val); - resetPage(); - }} - /> - - {/* Pagination + refresh pushed to the right */} -
-
-
-
- - {/* Table — pagination handled in header */} - - columns={columns} - dataSource={auditLogs} - rowKey="id" - loading={{ - spinning: query.isLoading, - indicator: } size="small" />, - }} - size="small" - pagination={false} - onRow={(record) => ({ - onClick: () => handleRowClick(record), - style: { cursor: "pointer" }, - })} - /> -
- - setDrawerOpen(false)} log={selectedLog} /> - - ); -} diff --git a/ui/litellm-dashboard/src/components/view_logs/columns.tsx b/ui/litellm-dashboard/src/components/view_logs/columns.tsx index 0ce5e8e4717..d3ba90d4e2d 100644 --- a/ui/litellm-dashboard/src/components/view_logs/columns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/columns.tsx @@ -539,15 +539,3 @@ const CollapsibleJsonCell = ({ jsonData }: { jsonData: any }) => { ); }; - -export type AuditLogEntry = { - id: string; - updated_at: string; - changed_by: string; - changed_by_api_key: string; - action: string; - table_name: string; - object_id: string; - before_value: Record; - updated_values: Record; -}; diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index ee08712e56b..aa8077a02e9 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -8,7 +8,7 @@ import { KeyResponse } from "../key_team_helpers/key_list"; import FilterComponent from "../molecules/filter"; import { keyInfoV1Call } from "../networking"; import KeyInfoView from "../templates/key_info_view"; -import AuditLogs from "./audit_logs"; +import AuditLogsPanel from "./AuditLogsPanel"; import { createColumns, LogEntry, type LogsSortField } from "./columns"; import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "./constants"; import { getLogFilterOptions } from "./filter_options"; @@ -296,7 +296,7 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p )} - Date: Mon, 20 Jul 2026 22:24:28 -0700 Subject: [PATCH 07/84] refactor(ui): migrate organizations table onto shared DataTable The organizations admin table was a hand-rolled tremor/antd table in a single snake_case file. This moves it onto the shared DataTable and cell library the other migrated tables use, splitting it into a data-owning OrganizationsPanel, a thin OrganizationsTable consumer, and a getOrganizationsTableColumns module The models column no longer uses a per-row accordion whose expand state lived in the parent; it renders the shared ModelsCell with truncation and a "+N more" tooltip, matching every other table with a models column. Row actions (Edit, Delete) move into a per-row overflow menu gated to proxy admins, while the detail view, create modal, and delete modal stay in the panel. The server-side org id / org alias search stays wired to the useOrganizations hook, and the table gains an initial-load skeleton plus a search-aware empty state. The dead sort_by / sort_order filter fields, the misnamed "Info" column that only ever showed a member count, and an unused refresh affordance are dropped; the default created_at descending sort is preserved --- ui/litellm-dashboard/eslint-suppressions.json | 5 - .../OrganizationFilters.test.tsx | 2 - .../organizations/OrganizationFilters.tsx | 2 - .../_components/OrganizationsPanel.test.tsx | 57 ++ .../_components/OrganizationsPanel.tsx | 299 ++++++++++ .../_components/OrganizationsTable.test.tsx | 188 ++++++ .../_components/OrganizationsTable.tsx | 75 +++ .../_components/OrganizationsTableColumns.tsx | 186 ++++++ .../_components/organizations.test.tsx | 39 -- .../_components/organizations.tsx | 535 ------------------ .../app/(dashboard)/organizations/page.tsx | 4 +- 11 files changed, 807 insertions(+), 585 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.test.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index f935af8907d..8e81dc55f31 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -697,11 +697,6 @@ "count": 1 } }, - "src/app/(dashboard)/organizations/_components/organizations.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx": { "no-restricted-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/OrganizationFilters.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/OrganizationFilters.test.tsx index 814625ff6be..37eeaf4c2af 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/OrganizationFilters.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/OrganizationFilters.test.tsx @@ -7,8 +7,6 @@ describe("OrganizationFilters", () => { const defaultFilters: FilterState = { org_id: "", org_alias: "", - sort_by: "", - sort_order: "asc", }; it("should render", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/OrganizationFilters.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/OrganizationFilters.tsx index 5643a4bc51a..6ad2f00fdb0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/OrganizationFilters.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/OrganizationFilters.tsx @@ -14,8 +14,6 @@ interface OrganizationFiltersProps { type FilterState = { org_id: string; org_alias: string; - sort_by: string; - sort_order: "asc" | "desc"; }; const OrganizationFilters = ({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx new file mode 100644 index 00000000000..d381e5e65ca --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx @@ -0,0 +1,57 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen } from "@testing-library/react"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; + +vi.mock("@/components/vector_store_management/VectorStoreSelector", () => ({ + __esModule: true, + default: () => null, +})); +vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ + __esModule: true, + default: () => null, +})); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ + accessToken: null, + userId: null, + userRole: null, + }), +})); +vi.mock("./OrganizationsTable", () => ({ + __esModule: true, + default: (props: { isLoading: boolean }) => ( +
isLoading:{String(props.isLoading)}
+ ), +})); + +import OrganizationsPanel from "./OrganizationsPanel"; + +const renderWithQueryClient = (ui: React.ReactElement) => { + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + return render({ui}); +}; + +describe("OrganizationsPanel", () => { + it("gates non-premium users behind the enterprise notice", () => { + renderWithQueryClient(); + + expect(screen.getByText(/LiteLLM Enterprise feature/i)).toBeInTheDocument(); + expect(screen.queryByText("+ Create New Organization")).not.toBeInTheDocument(); + }); + + it("shows the create button for a premium admin", () => { + renderWithQueryClient(); + + expect(screen.getByText("+ Create New Organization")).toBeInTheDocument(); + }); + + it("resolves the loading skeleton to false when the query is disabled (no token)", () => { + renderWithQueryClient(); + + // A disabled React Query keeps isPending true forever; feeding isLoading avoids a stuck skeleton. + expect(screen.getByTestId("organizations-table")).toHaveTextContent("isLoading:false"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.tsx new file mode 100644 index 00000000000..9f7e029a1d4 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.tsx @@ -0,0 +1,299 @@ +import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { useUserModels } from "@/app/(dashboard)/hooks/models/useModels"; +import OrganizationFilters, { FilterState } from "@/app/(dashboard)/organizations/OrganizationFilters"; +import { InfoCircleOutlined } from "@ant-design/icons"; +import { Form, Input, Modal, Select as Select2, Tooltip } from "antd"; +import { useQueryClient } from "@tanstack/react-query"; +import React, { useState } from "react"; +import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; +import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; +import { ModelSelect } from "@/components/ModelSelect/ModelSelect"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { organizationCreateCall, organizationDeleteCall } from "@/components/networking"; +import OrganizationInfoView from "@/components/organization/organization_view"; +import NumericalInput from "@/components/shared/numerical_input"; +import { Button } from "@/components/ui/button"; +import VectorStoreSelector from "@/components/vector_store_management/VectorStoreSelector"; + +import OrganizationsTable from "./OrganizationsTable"; + +interface OrganizationsPanelProps { + userRole: string; + accessToken: string | null; + premiumUser: boolean; +} + +const OrganizationsPanel: React.FC = ({ userRole, accessToken, premiumUser }) => { + const [selectedOrgId, setSelectedOrgId] = useState(null); + const [editOrg, setEditOrg] = useState(false); + const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); + const [orgToDelete, setOrgToDelete] = useState(null); + const [isDeleting, setIsDeleting] = useState(false); + const [isOrgModalVisible, setIsOrgModalVisible] = useState(false); + const [form] = Form.useForm(); + const [showFilters, setShowFilters] = useState(false); + const [filters, setFilters] = useState({ org_id: "", org_alias: "" }); + + const queryClient = useQueryClient(); + const { data: organizations = [], isLoading } = useOrganizations({ + org_id: filters.org_id, + org_alias: filters.org_alias, + }); + const { data: userModels = [] } = useUserModels(); + + const searchActive = Boolean(filters.org_id || filters.org_alias); + + const refetchOrganizations = () => queryClient.invalidateQueries({ queryKey: organizationKeys.lists() }); + + const handleFilterChange = (key: keyof FilterState, value: string) => { + setFilters((previousFilters) => ({ ...previousFilters, [key]: value })); + }; + + const handleFilterReset = () => { + setFilters({ org_id: "", org_alias: "" }); + }; + + const handleDelete = (orgId: string | null) => { + if (!orgId) return; + + setOrgToDelete(orgId); + setIsDeleteModalOpen(true); + }; + + const confirmDelete = async () => { + if (!orgToDelete || !accessToken) return; + + try { + setIsDeleting(true); + await organizationDeleteCall(accessToken, orgToDelete); + NotificationsManager.success("Organization deleted successfully"); + + setIsDeleteModalOpen(false); + setOrgToDelete(null); + await refetchOrganizations(); + } catch (error) { + console.error("Error deleting organization:", error); + } finally { + setIsDeleting(false); + } + }; + + const cancelDelete = () => { + setIsDeleteModalOpen(false); + setOrgToDelete(null); + }; + + const handleCreate = async (values: any) => { + try { + if (!accessToken) return; + + // Transform allowed_vector_store_ids and allowed_mcp_servers_and_groups into object_permission + if ( + (values.allowed_vector_store_ids && values.allowed_vector_store_ids.length > 0) || + (values.allowed_mcp_servers_and_groups && + (values.allowed_mcp_servers_and_groups.servers?.length > 0 || + values.allowed_mcp_servers_and_groups.accessGroups?.length > 0)) + ) { + values.object_permission = {}; + if (values.allowed_vector_store_ids && values.allowed_vector_store_ids.length > 0) { + values.object_permission.vector_stores = values.allowed_vector_store_ids; + delete values.allowed_vector_store_ids; + } + if (values.allowed_mcp_servers_and_groups) { + if (values.allowed_mcp_servers_and_groups.servers?.length > 0) { + values.object_permission.mcp_servers = values.allowed_mcp_servers_and_groups.servers; + } + if (values.allowed_mcp_servers_and_groups.accessGroups?.length > 0) { + values.object_permission.mcp_access_groups = values.allowed_mcp_servers_and_groups.accessGroups; + } + delete values.allowed_mcp_servers_and_groups; + } + } + + await organizationCreateCall(accessToken, values); + NotificationsManager.success("Organization created successfully"); + setIsOrgModalVisible(false); + form.resetFields(); + await refetchOrganizations(); + } catch (error) { + console.error("Error creating organization:", error); + } + }; + + const handleCancel = () => { + setIsOrgModalVisible(false); + form.resetFields(); + }; + + if (!premiumUser) { + return ( +
+

+ This is a LiteLLM Enterprise feature, and requires a valid key to use. Get a trial key{" "} + + here + + . +

+
+ ); + } + + return ( +
+ {(userRole === "Admin" || userRole === "Org Admin") && ( + + )} + + {selectedOrgId ? ( + { + setSelectedOrgId(null); + setEditOrg(false); + }} + accessToken={accessToken} + is_org_admin={true} + is_proxy_admin={userRole === "Admin"} + userModels={userModels} + editOrg={editOrg} + /> + ) : ( + <> +

Click on an organization ID to view its details.

+ + { + setSelectedOrgId(organizationId); + setEditOrg(true); + }} + onDeleteClick={handleDelete} + /> + + )} + + +
+ + + + + form.setFieldValue("models", values)} + context="organization" + /> + + + + + + + + daily + weekly + monthly + + + + + + + + + + + Allowed Vector Stores{" "} + + + + + } + name="allowed_vector_store_ids" + className="mt-4" + help="Select vector stores this organization can access. Leave empty for access to all vector stores" + > + form.setFieldValue("allowed_vector_store_ids", values)} + value={form.getFieldValue("allowed_vector_store_ids")} + accessToken={accessToken || ""} + placeholder="Select vector stores (optional)" + /> + + + + Allowed MCP Servers{" "} + + + + + } + name="allowed_mcp_servers_and_groups" + className="mt-4" + help="Select MCP servers and access groups this organization can access." + > + form.setFieldValue("allowed_mcp_servers_and_groups", values)} + value={form.getFieldValue("allowed_mcp_servers_and_groups")} + accessToken={accessToken || ""} + placeholder="Select MCP servers and access groups (optional)" + /> + + + + + + +
+ +
+ +
+ + +
+ ); +}; + +export default OrganizationsPanel; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx new file mode 100644 index 00000000000..a06c5c885e3 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx @@ -0,0 +1,188 @@ +import { render, screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; + +import { Organization } from "@/components/networking"; + +import OrganizationsTable from "./OrganizationsTable"; + +const makeOrganization = (overrides: Partial = {}): Organization => ({ + organization_id: "org-alpha", + organization_alias: "Alpha", + budget_id: "budget-1", + metadata: {}, + models: [], + spend: 0, + model_spend: {}, + created_at: "2023-01-01T00:00:00Z", + created_by: "someone", + updated_at: "2023-01-01T00:00:00Z", + updated_by: "someone", + litellm_budget_table: null, + teams: null, + users: null, + members: null, + ...overrides, +}); + +const baseProps = { + isLoading: false, + userRole: "Admin", + searchActive: false, + onOrganizationClick: vi.fn(), + onEditClick: vi.fn(), + onDeleteClick: vi.fn(), +}; + +describe("OrganizationsTable", () => { + it("renders every column header", () => { + render(); + for (const header of [ + "Organization ID", + "Organization Name", + "Created", + "Spend (USD)", + "Budget (USD)", + "Models", + "TPM / RPM Limits", + "Members", + ]) { + expect(screen.getByText(header)).toBeInTheDocument(); + } + }); + + it("opens the detail view when the organization ID cell is clicked", async () => { + const user = userEvent.setup(); + const onOrganizationClick = vi.fn(); + render( + , + ); + + await user.click(screen.getByText("org-123")); + + expect(onOrganizationClick).toHaveBeenCalledWith("org-123"); + }); + + it("edits and deletes an organization through the ⋯ actions menu (admin)", async () => { + const user = userEvent.setup(); + const onEditClick = vi.fn(); + const onDeleteClick = vi.fn(); + render( + , + ); + + await user.click(screen.getByTestId("organization-actions-org-9")); + await user.click(await screen.findByTestId("organization-action-edit")); + expect(onEditClick).toHaveBeenCalledWith("org-9"); + + await user.click(screen.getByTestId("organization-actions-org-9")); + await user.click(await screen.findByTestId("organization-action-delete")); + expect(onDeleteClick).toHaveBeenCalledWith("org-9"); + }); + + it("hides the row actions menu from non-admins", () => { + render( + , + ); + + expect(screen.queryByTestId("organization-actions-org-9")).not.toBeInTheDocument(); + }); + + it("sorts by created_at descending by default", () => { + render( + , + ); + + const rows = screen.getAllByRole("row"); + // rows[0] is the header row; the newest organization must lead the body. + expect(within(rows[1]).getByText("Newer")).toBeInTheDocument(); + expect(within(rows[2]).getByText("Older")).toBeInTheDocument(); + }); + + it("renders budget, limits, members, and models for a fully-populated organization", () => { + render( + , + ); + + expect(screen.getByText("$100.00")).toBeInTheDocument(); + expect(screen.getByText("TPM: 1000")).toBeInTheDocument(); + expect(screen.getByText("RPM: 60")).toBeInTheDocument(); + expect(screen.getByText("3 Members")).toBeInTheDocument(); + // Five models, three visible -> the shared ModelsCell collapses the rest. + expect(screen.getByText("+2 more")).toBeInTheDocument(); + }); + + it("shows Unlimited budget and All Proxy Models when unset", () => { + render( + , + ); + + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + // Budget shows a standalone "Unlimited"; the limits fall back inline. + expect(screen.getByText("Unlimited")).toBeInTheDocument(); + expect(screen.getByText("TPM: Unlimited")).toBeInTheDocument(); + expect(screen.getByText("RPM: Unlimited")).toBeInTheDocument(); + }); + + it("renders loading skeletons instead of rows while loading", () => { + render( + , + ); + + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); + expect(screen.queryByText("ShouldNotShow")).not.toBeInTheDocument(); + }); + + it("uses a search-aware empty state", () => { + const { rerender } = render(); + expect(screen.getByText("No organizations yet")).toBeInTheDocument(); + + rerender(); + expect(screen.getByText("No matching organizations")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx new file mode 100644 index 00000000000..8e68a57d2f7 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx @@ -0,0 +1,75 @@ +"use client"; + +import { SortingState } from "@tanstack/react-table"; +import { Building2, SearchX } from "lucide-react"; +import React, { useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; +import { Organization } from "@/components/networking"; + +import { getOrganizationsTableColumns } from "./OrganizationsTableColumns"; + +interface OrganizationsTableProps { + organizations: Organization[]; + isLoading: boolean; + userRole: string; + searchActive: boolean; + onOrganizationClick: (organizationId: string) => void; + onEditClick: (organizationId: string) => void; + onDeleteClick: (organizationId: string) => void; +} + +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +function EmptyState({ searchActive }: { searchActive: boolean }) { + const Icon = searchActive ? SearchX : Building2; + return ( +
+
+ +
+
+ {searchActive ? "No matching organizations" : "No organizations yet"} +
+
+ {searchActive + ? "No organizations match your search. Try a different name or ID." + : "Create an organization to group teams, models, and budgets."} +
+
+ ); +} + +const OrganizationsTable: React.FC = ({ + organizations, + isLoading, + userRole, + searchActive, + onOrganizationClick, + onEditClick, + onDeleteClick, +}) => { + const [sorting, setSorting] = useState(DEFAULT_SORTING); + + const columns = useMemo(() => { + const deps = { userRole, onOrganizationClick, onEditClick, onDeleteClick }; + return getOrganizationsTableColumns(deps); + }, [userRole, onOrganizationClick, onEditClick, onDeleteClick]); + + return ( + organization.organization_id || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + isLoading={isLoading} + loadingMessage="Loading organizations…" + noDataMessage={} + size="compact" + /> + ); +}; + +export default OrganizationsTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx new file mode 100644 index 00000000000..31f6a00916c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx @@ -0,0 +1,186 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { MoreHorizontal, Pencil, Trash2 } from "lucide-react"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdentityCell, ModelsCell, MoneyCell } from "@/components/shared/table_cells"; +import { Organization } from "@/components/networking"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; + +interface OrganizationBudget { + max_budget?: number | null; + tpm_limit?: number | null; + rpm_limit?: number | null; +} + +const getOrganizationBudget = (organization: Organization): OrganizationBudget => + (organization.litellm_budget_table ?? {}) as OrganizationBudget; + +function OrganizationLimitsCell({ organization }: { organization: Organization }) { + const { tpm_limit, rpm_limit } = getOrganizationBudget(organization); + return ( +
+ TPM: {tpm_limit ? tpm_limit : "Unlimited"} + RPM: {rpm_limit ? rpm_limit : "Unlimited"} +
+ ); +} + +interface OrganizationRowActionsProps { + organization: Organization; + onEditClick: (organizationId: string) => void; + onDeleteClick: (organizationId: string) => void; +} + +function OrganizationRowActions({ organization, onEditClick, onDeleteClick }: OrganizationRowActionsProps) { + return ( + + + + + + onEditClick(organization.organization_id)} + > + + Edit + + onDeleteClick(organization.organization_id)} + > + + Delete + + + + ); +} + +export interface OrganizationsTableColumnsDeps { + userRole: string; + onOrganizationClick: (organizationId: string) => void; + onEditClick: (organizationId: string) => void; + onDeleteClick: (organizationId: string) => void; +} + +export const getOrganizationsTableColumns = ({ + userRole, + onOrganizationClick, + onEditClick, + onDeleteClick, +}: OrganizationsTableColumnsDeps): ColumnDef[] => [ + { + id: "organization_id", + accessorKey: "organization_id", + meta: { title: "Organization ID" }, + header: ({ column }) => , + size: 220, + enableSorting: true, + cell: ({ row }) => ( + onOrganizationClick(row.original.organization_id)} + /> + ), + }, + { + id: "organization_alias", + accessorKey: "organization_alias", + meta: { title: "Organization Name" }, + header: ({ column }) => , + size: 200, + enableSorting: true, + cell: ({ row }) => { + const alias = row.original.organization_alias; + return ( + + {alias || "-"} + + ); + }, + }, + { + id: "created_at", + accessorKey: "created_at", + sortingFn: "datetime", + meta: { title: "Created" }, + header: ({ column }) => , + size: 130, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "spend", + accessorKey: "spend", + meta: { title: "Spend (USD)" }, + header: ({ column }) => , + size: 120, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "max_budget", + meta: { title: "Budget (USD)" }, + header: "Budget (USD)", + size: 120, + enableSorting: false, + cell: ({ row }) => ( + + ), + }, + { + id: "models", + meta: { title: "Models", skeleton: "chips" }, + header: "Models", + size: 260, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "limits", + meta: { title: "TPM / RPM Limits" }, + header: "TPM / RPM Limits", + size: 150, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "members", + meta: { title: "Members" }, + header: "Members", + size: 100, + enableSorting: false, + cell: ({ row }) => {row.original.members?.length ?? 0} Members, + }, + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => + userRole === "Admin" ? ( +
+ +
+ ) : null, + }, +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.test.tsx deleted file mode 100644 index 75a6d30ac2e..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.test.tsx +++ /dev/null @@ -1,39 +0,0 @@ -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render } from "@testing-library/react"; -import React from "react"; -import { describe, expect, it, vi } from "vitest"; - -vi.mock("@/components/vector_store_management/VectorStoreSelector", () => ({ - __esModule: true, - default: () => null, -})); -vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ - __esModule: true, - default: () => null, -})); -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: () => ({ - accessToken: null, - userId: null, - userRole: null, - }), -})); - -import OrganizationsTable from "./organizations"; - -const renderWithQueryClient = (ui: React.ReactElement) => { - const queryClient = new QueryClient({ - defaultOptions: { queries: { retry: false } }, - }); - return render({ui}); -}; - -describe("OrganizationsTable", () => { - it("should render the OrganizationsTable component", () => { - const { getByText } = renderWithQueryClient( - , - ); - - expect(getByText("+ Create New Organization")).toBeInTheDocument(); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx deleted file mode 100644 index 87d8010759d..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/organizations.tsx +++ /dev/null @@ -1,535 +0,0 @@ -import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; -import { useUserModels } from "@/app/(dashboard)/hooks/models/useModels"; -import OrganizationFilters, { FilterState } from "@/app/(dashboard)/organizations/OrganizationFilters"; -import { InfoCircleOutlined } from "@ant-design/icons"; -import { ChevronDownIcon, ChevronRightIcon, RefreshIcon } from "@heroicons/react/outline"; -import { - Badge, - Button, - Card, - Col, - Grid, - Icon, - Tab, - TabGroup, - Table, - TableBody, - TableCell, - TableHead, - TableHeaderCell, - TableRow, - TabList, - TabPanel, - TabPanels, - Text, - TextInput, -} from "@tremor/react"; -import { Form, Input, Modal, Select as Select2, Tooltip } from "antd"; -import { useQueryClient } from "@tanstack/react-query"; -import React, { useState } from "react"; -import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells"; -import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import TableIconActionButton from "@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; -import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; -import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; -import { ModelSelect } from "@/components/ModelSelect/ModelSelect"; -import NotificationsManager from "@/components/molecules/notifications_manager"; -import { - Organization, - organizationCreateCall, - organizationDeleteCall, - organizationListCall, -} from "@/components/networking"; -import OrganizationInfoView from "@/components/organization/organization_view"; -import NumericalInput from "@/components/shared/numerical_input"; -import VectorStoreSelector from "@/components/vector_store_management/VectorStoreSelector"; - -interface OrganizationsTableProps { - userRole: string; - accessToken: string | null; - lastRefreshed?: string; - handleRefreshClick?: () => void; - premiumUser: boolean; -} - -export const fetchOrganizations = async ( - accessToken: string, - setOrganizations: (organizations: Organization[]) => void, - org_id: string | null = null, - org_alias: string | null = null, -) => { - const organizations = await organizationListCall(accessToken, org_id, org_alias); - setOrganizations(organizations); -}; - -const OrganizationsTable: React.FC = ({ - userRole, - accessToken, - lastRefreshed, - handleRefreshClick, - premiumUser, -}) => { - const [selectedOrgId, setSelectedOrgId] = useState(null); - const [editOrg, setEditOrg] = useState(false); - const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); - const [orgToDelete, setOrgToDelete] = useState(null); - const [isDeleting, setIsDeleting] = useState(false); - const [isOrgModalVisible, setIsOrgModalVisible] = useState(false); - const [form] = Form.useForm(); - const [expandedAccordions, setExpandedAccordions] = useState>({}); - const [showFilters, setShowFilters] = useState(false); - const [filters, setFilters] = useState({ - org_id: "", - org_alias: "", - sort_by: "created_at", - sort_order: "desc", - }); - - const queryClient = useQueryClient(); - const { data: organizations = [] } = useOrganizations({ org_id: filters.org_id, org_alias: filters.org_alias }); - const { data: userModels = [] } = useUserModels(); - - const refetchOrganizations = () => queryClient.invalidateQueries({ queryKey: organizationKeys.lists() }); - - const handleFilterChange = (key: keyof FilterState, value: string) => { - setFilters((previousFilters) => ({ ...previousFilters, [key]: value })); - }; - - const handleFilterReset = () => { - setFilters({ - org_id: "", - org_alias: "", - sort_by: "created_at", - sort_order: "desc", - }); - }; - - const handleDelete = (orgId: string | null) => { - if (!orgId) return; - - setOrgToDelete(orgId); - setIsDeleteModalOpen(true); - }; - - const confirmDelete = async () => { - if (!orgToDelete || !accessToken) return; - - try { - setIsDeleting(true); - await organizationDeleteCall(accessToken, orgToDelete); - NotificationsManager.success("Organization deleted successfully"); - - setIsDeleteModalOpen(false); - setOrgToDelete(null); - await refetchOrganizations(); - } catch (error) { - console.error("Error deleting organization:", error); - } finally { - setIsDeleting(false); - } - }; - - const cancelDelete = () => { - setIsDeleteModalOpen(false); - setOrgToDelete(null); - }; - - const handleCreate = async (values: any) => { - try { - if (!accessToken) return; - - // Transform allowed_vector_store_ids and allowed_mcp_servers_and_groups into object_permission - if ( - (values.allowed_vector_store_ids && values.allowed_vector_store_ids.length > 0) || - (values.allowed_mcp_servers_and_groups && - (values.allowed_mcp_servers_and_groups.servers?.length > 0 || - values.allowed_mcp_servers_and_groups.accessGroups?.length > 0)) - ) { - values.object_permission = {}; - if (values.allowed_vector_store_ids && values.allowed_vector_store_ids.length > 0) { - values.object_permission.vector_stores = values.allowed_vector_store_ids; - delete values.allowed_vector_store_ids; - } - if (values.allowed_mcp_servers_and_groups) { - if (values.allowed_mcp_servers_and_groups.servers?.length > 0) { - values.object_permission.mcp_servers = values.allowed_mcp_servers_and_groups.servers; - } - if (values.allowed_mcp_servers_and_groups.accessGroups?.length > 0) { - values.object_permission.mcp_access_groups = values.allowed_mcp_servers_and_groups.accessGroups; - } - delete values.allowed_mcp_servers_and_groups; - } - } - - await organizationCreateCall(accessToken, values); - NotificationsManager.success("Organization created successfully"); - setIsOrgModalVisible(false); - form.resetFields(); - await refetchOrganizations(); - } catch (error) { - console.error("Error creating organization:", error); - } - }; - - const handleCancel = () => { - setIsOrgModalVisible(false); - form.resetFields(); - }; - - if (!premiumUser) { - return ( -
- - This is a LiteLLM Enterprise feature, and requires a valid key to use. Get a trial key{" "} - - here - - . - -
- ); - } - - return ( -
- -
- {(userRole === "Admin" || userRole === "Org Admin") && ( - - )} - {selectedOrgId ? ( - { - setSelectedOrgId(null); - setEditOrg(false); - }} - accessToken={accessToken} - is_org_admin={true} // You'll need to implement proper org admin check - is_proxy_admin={userRole === "Admin"} - userModels={userModels} - editOrg={editOrg} - /> - ) : ( - - -
- Your Organizations -
-
- {lastRefreshed && Last Refreshed: {lastRefreshed}} - -
-
- - - Click on “Organization ID” to view organization details. - -
- -
-
- -
-
-
- - - Organization ID - Organization Name - Created - Spend (USD) - Budget (USD) - Models - TPM / RPM Limits - Info - Actions - - - - - {organizations && organizations.length > 0 - ? organizations - .sort((a, b) => new Date(b.created_at).getTime() - new Date(a.created_at).getTime()) - .map((org: Organization) => ( - - - - - {org.organization_alias} - - - - - - - - - - 3 ? "px-0" : ""} - > -
- {Array.isArray(org.models) ? ( -
- {org.models.length === 0 ? ( - - All Proxy Models - - ) : ( - <> -
- {org.models.length > 3 && ( -
- { - setExpandedAccordions((prev) => ({ - ...prev, - [org.organization_id || ""]: - !prev[org.organization_id || ""], - })); - }} - /> -
- )} -
- {org.models.slice(0, 3).map((model, index) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} - {org.models.length > 3 && - !expandedAccordions[org.organization_id || ""] && ( - - - +{org.models.length - 3}{" "} - {org.models.length - 3 === 1 - ? "more model" - : "more models"} - - - )} - {expandedAccordions[org.organization_id || ""] && ( -
- {org.models.slice(3).map((model, index) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} -
- )} -
-
- - )} -
- ) : null} -
-
- - - TPM:{" "} - {org.litellm_budget_table?.tpm_limit - ? org.litellm_budget_table?.tpm_limit - : "Unlimited"} -
- RPM:{" "} - {org.litellm_budget_table?.rpm_limit - ? org.litellm_budget_table?.rpm_limit - : "Unlimited"} -
-
- - {org.members?.length || 0} Members - - - {userRole === "Admin" && ( - <> - { - setSelectedOrgId(org.organization_id); - setEditOrg(true); - }} - /> - handleDelete(org.organization_id)} - /> - - )} - -
- )) - : null} -
-
- - - - - - - )} - - - -
- - - - - form.setFieldValue("models", values)} - context="organization" - /> - - - - - - - - daily - weekly - monthly - - - - - - - - - - - Allowed Vector Stores{" "} - - - - - } - name="allowed_vector_store_ids" - className="mt-4" - help="Select vector stores this organization can access. Leave empty for access to all vector stores" - > - form.setFieldValue("allowed_vector_store_ids", values)} - value={form.getFieldValue("allowed_vector_store_ids")} - accessToken={accessToken || ""} - placeholder="Select vector stores (optional)" - /> - - - - Allowed MCP Servers{" "} - - - - - } - name="allowed_mcp_servers_and_groups" - className="mt-4" - help="Select MCP servers and access groups this organization can access." - > - form.setFieldValue("allowed_mcp_servers_and_groups", values)} - value={form.getFieldValue("allowed_mcp_servers_and_groups")} - accessToken={accessToken || ""} - placeholder="Select MCP servers and access groups (optional)" - /> - - - - - - -
- -
-
-
- - - - ); -}; - -export default OrganizationsTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/page.tsx index 649e54f63eb..a492a572580 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/page.tsx @@ -1,9 +1,9 @@ "use client"; -import OrganizationsTable from "./_components/organizations"; +import OrganizationsPanel from "./_components/OrganizationsPanel"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; export default function OrganizationsPage() { const { accessToken, userRole, premiumUser } = useAuthorized(); - return ; + return ; } From ff8d8797dd9512c1dbf4504805256c9678ed4078 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Jul 2026 23:26:27 -0700 Subject: [PATCH 08/84] test(ui): pin memory table page-size behavior on the last page Changing rows-per-page while on the last page recomputes the page index from the top visible row, so the table lands on the new last page instead of an out-of-range one. Pin that, since it depends on the parent holding the full PaginationState rather than just the page index. --- .../memory/_components/MemoryTable.test.tsx | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx index 0dcbd6796a8..f664c650cd4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx @@ -1,6 +1,7 @@ import { PaginationState } from "@tanstack/react-table"; import { render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; +import React, { useState } from "react"; import { describe, expect, it, vi } from "vitest"; import { MemoryRow } from "@/components/networking"; @@ -134,6 +135,31 @@ describe("MemoryTable", () => { expect(onRefresh).toHaveBeenCalledTimes(1); }); + it("keeps the page in range when the rows-per-page selector shrinks the page count", async () => { + const user = userEvent.setup(); + const rowCount = 120; + const seen: PaginationState[] = []; + + function Harness() { + const [pagination, setPagination] = useState({ pageIndex: 4, pageSize: 25 }); + seen.push(pagination); + return ( + + ); + } + + render(); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 5 of 5"); + + await user.click(screen.getByTestId("pagination-page-size")); + await user.click(await screen.findByRole("option", { name: "100" })); + + const final = seen[seen.length - 1]; + expect(final.pageSize).toBe(100); + expect(final.pageIndex).toBeLessThanOrEqual(Math.ceil(rowCount / final.pageSize) - 1); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 2 of 2"); + }); + it("renders secondary id and date cells for the row", () => { render(); const table = screen.getByRole("table"); From e411d637b350c5403ed55627e8d38a6fbf963c2c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 20 Jul 2026 15:32:36 -0700 Subject: [PATCH 09/84] feat(gemini): day-0 pricing for gemini-3.6-flash and gemini-3.5-flash-lite --- ...odel_prices_and_context_window_backup.json | 333 ++++++++++++++++++ model_prices_and_context_window.json | 333 ++++++++++++++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 72 ++++ 3 files changed, 738 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bb6243e50ed..d3917886060 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -17566,6 +17566,61 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.5-flash-lite": { + "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -18233,6 +18288,60 @@ }, "web_search_billing_unit": "per_query" }, + "vertex_ai/gemini-3.6-flash": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_flex": 7.5e-08, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 7.5e-06, + "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, + "output_cost_per_token_flex": 3.75e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 2.7e-06, + "output_cost_per_token_priority": 1.35e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "vertex_ai/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -19585,6 +19694,63 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.5-flash-lite": { + "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "rpm": 15, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 250000, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, "input_cost_per_audio_token": 1e-06, @@ -19691,6 +19857,63 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.6-flash": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_flex": 7.5e-08, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 7.5e-06, + "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, + "output_cost_per_token_flex": 3.75e-06, + "rpm": 2000, + "source": "https://ai.google.dev/pricing/gemini-3", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 800000, + "input_cost_per_token_priority": 2.7e-06, + "output_cost_per_token_priority": 1.35e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, @@ -19971,6 +20194,61 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.6-flash": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_flex": 7.5e-08, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 7.5e-06, + "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, + "output_cost_per_token_flex": 3.75e-06, + "source": "https://ai.google.dev/pricing/gemini-3", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 2.7e-06, + "output_cost_per_token_priority": 1.35e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, @@ -37232,6 +37510,61 @@ }, "web_search_billing_unit": "per_query" }, + "vertex_ai/gemini-3.5-flash-lite": { + "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "vertex_ai/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5efd61f9747..c9d871fc41d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -17644,6 +17644,61 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.5-flash-lite": { + "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -18311,6 +18366,60 @@ }, "web_search_billing_unit": "per_query" }, + "vertex_ai/gemini-3.6-flash": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_flex": 7.5e-08, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 7.5e-06, + "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, + "output_cost_per_token_flex": 3.75e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 2.7e-06, + "output_cost_per_token_priority": 1.35e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "vertex_ai/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -19663,6 +19772,63 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.5-flash-lite": { + "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "rpm": 15, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 250000, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, "input_cost_per_audio_token": 1e-06, @@ -19769,6 +19935,63 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.6-flash": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_flex": 7.5e-08, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 7.5e-06, + "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, + "output_cost_per_token_flex": 3.75e-06, + "rpm": 2000, + "source": "https://ai.google.dev/pricing/gemini-3", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 800000, + "input_cost_per_token_priority": 2.7e-06, + "output_cost_per_token_priority": 1.35e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, @@ -20049,6 +20272,61 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.6-flash": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_flex": 7.5e-08, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, + "input_cost_per_token_flex": 7.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 7.5e-06, + "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, + "output_cost_per_token_flex": 3.75e-06, + "source": "https://ai.google.dev/pricing/gemini-3", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 2.7e-06, + "output_cost_per_token_priority": 1.35e-05, + "cache_read_input_token_cost_priority": 2.7e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, @@ -37323,6 +37601,61 @@ }, "web_search_billing_unit": "per_query" }, + "vertex_ai/gemini-3.5-flash-lite": { + "cache_read_input_token_cost": 3e-08, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token": 3e-07, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "vertex_ai/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index b156faf3ea6..9ff67a82f40 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2237,3 +2237,75 @@ def test_token_type_cost_breakdown_applies_regional_uplift(): text_input_cost = 600 * model_info["input_cost_per_token"] * uplift assert text_output_cost + eu.reasoning_cost == pytest.approx(completion_cost) assert text_input_cost + eu.cache_read_cost == pytest.approx(prompt_cost) + + +GEMINI_DAY0_LAUNCH_PRICING = [ + ("gemini-3.6-flash", 1.5e-06, 7.5e-06, 1.5e-07), + ("gemini/gemini-3.6-flash", 1.5e-06, 7.5e-06, 1.5e-07), + ("vertex_ai/gemini-3.6-flash", 1.5e-06, 7.5e-06, 1.5e-07), + ("gemini-3.5-flash-lite", 3e-07, 2.5e-06, 3e-08), + ("gemini/gemini-3.5-flash-lite", 3e-07, 2.5e-06, 3e-08), + ("vertex_ai/gemini-3.5-flash-lite", 3e-07, 2.5e-06, 3e-08), +] + + +@pytest.mark.parametrize("model,input_cost,output_cost,cache_read_cost", GEMINI_DAY0_LAUNCH_PRICING) +def test_gemini_36_flash_and_35_flash_lite_launch_pricing(model, input_cost, output_cost, cache_read_cost): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model_cost_map = litellm.model_cost[model] + assert model_cost_map["input_cost_per_token"] == input_cost + assert model_cost_map["output_cost_per_token"] == output_cost + assert model_cost_map["output_cost_per_reasoning_token"] == output_cost + assert model_cost_map["cache_read_input_token_cost"] == cache_read_cost + assert model_cost_map["mode"] == "chat" + assert model_cost_map["supports_reasoning"] is True + assert model_cost_map["supports_function_calling"] is True + assert model_cost_map["max_input_tokens"] == 1048576 + + +def test_generic_cost_per_token_gemini_36_flash(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=200, + text_tokens=300, + ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=1000), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model="gemini-3.6-flash", + usage=usage, + custom_llm_provider="gemini", + ) + assert prompt_cost == pytest.approx(0.0015) + assert completion_cost == pytest.approx(0.00375) + + +def test_generic_cost_per_token_gemini_35_flash_lite(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=200, + text_tokens=300, + ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=1000), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model="gemini-3.5-flash-lite", + usage=usage, + custom_llm_provider="gemini", + ) + assert prompt_cost == pytest.approx(0.0003) + assert completion_cost == pytest.approx(0.00125) From e20d3d4eccfcd35a4208b79bada247519e8f2da2 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 21 Jul 2026 09:22:02 -0700 Subject: [PATCH 10/84] fix(ui): serve /ui/assets from the nginx image instead of SPA fallback (#34066) --- ui/nginx.conf | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ui/nginx.conf b/ui/nginx.conf index adc394a28aa..a41ee5bd5b4 100644 --- a/ui/nginx.conf +++ b/ui/nginx.conf @@ -49,6 +49,9 @@ http { expires 1y; add_header Cache-Control "public, immutable"; } + location ^~ /ui/assets/ { + alias /usr/share/nginx/html/assets/; + } location = /favicon.ico { try_files $uri =404; expires 1d; From 4647f859586678a5d0f4513d0e84b3bfdd317d96 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 21 Jul 2026 10:28:21 -0700 Subject: [PATCH 11/84] Merge pull request #34078 from BerriAI/litellm_/elated-thompson-4a0c84 refactor(ui): migrate access groups table to shared DataTable --- .../_components/AccessGroupsPage.test.tsx | 198 ++++++------- .../_components/AccessGroupsPage.tsx | 277 +++--------------- .../_components/AccessGroupsTable.tsx | 72 +++++ .../_components/AccessGroupsTableColumns.tsx | 182 ++++++++++++ 4 files changed, 375 insertions(+), 354 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsTable.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsTableColumns.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.test.tsx index 7c8aaa2b785..a1484ffb5c5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.test.tsx @@ -38,6 +38,7 @@ const mockAccessGroups: AccessGroupResponse[] = [ const mockUseAccessGroups = vi.fn(); const mockUseDeleteAccessGroup = vi.fn(); const mockMutate = vi.fn(); +const mockUseAuthorized = vi.fn(); vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ useAccessGroups: () => mockUseAccessGroups(), @@ -47,6 +48,10 @@ vi.mock("@/app/(dashboard)/hooks/accessGroups/useDeleteAccessGroup", () => ({ useDeleteAccessGroup: () => mockUseDeleteAccessGroup(), })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + vi.mock("./AccessGroupsDetailsPage", () => ({ AccessGroupDetail: ({ accessGroupId, onBack }: { accessGroupId: string; onBack: () => void }) => (
@@ -65,49 +70,42 @@ vi.mock("./AccessGroupsModal/AccessGroupCreateModal", () => ({ ) : null, })); -vi.mock("@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton", () => ({ - default: ({ variant, tooltipText, onClick }: { variant: string; tooltipText: string; onClick: () => void }) => ( - - ), -})); +const makeGroups = (count: number): AccessGroupResponse[] => + Array.from({ length: count }, (_, index) => { + const suffix = String(index + 1).padStart(2, "0"); + return { + ...mockAccessGroups[0], + access_group_id: `ag-${suffix}`, + access_group_name: `Group ${suffix}`, + description: `Group ${suffix} description`, + }; + }); + +const openRowMenu = async (user: ReturnType, groupId: string) => { + await user.click(screen.getByTestId(`access-group-actions-${groupId}`)); + return screen.findByTestId("access-group-action-delete"); +}; describe("AccessGroupsPage", () => { beforeEach(() => { vi.clearAllMocks(); - mockUseAccessGroups.mockReturnValue({ - data: mockAccessGroups, - isLoading: false, - }); - mockUseDeleteAccessGroup.mockReturnValue({ - mutate: mockMutate, - isPending: false, - }); + mockUseAccessGroups.mockReturnValue({ data: mockAccessGroups, isLoading: false }); + mockUseDeleteAccessGroup.mockReturnValue({ mutate: mockMutate, isPending: false }); + mockUseAuthorized.mockReturnValue({ userRole: "Admin", accessToken: "sk-test" }); }); - it("should render", () => { - renderWithProviders(); - expect(screen.getByRole("heading", { name: "Access Groups" })).toBeInTheDocument(); - }); - - it("should display page title and subtitle", () => { + it("renders the page title and subtitle", () => { renderWithProviders(); expect(screen.getByRole("heading", { name: "Access Groups" })).toBeInTheDocument(); expect(screen.getByText("Manage resource permissions for your organization")).toBeInTheDocument(); }); - it("should display Create Access Group button", () => { + it("shows the Create Access Group button for an admin", () => { renderWithProviders(); expect(screen.getByRole("button", { name: /create access group/i })).toBeInTheDocument(); }); - it("should display search input with placeholder", () => { - renderWithProviders(); - expect(screen.getByPlaceholderText("Search groups by name, ID, or description...")).toBeInTheDocument(); - }); - - it("should display access groups in table", () => { + it("renders every access group row", () => { renderWithProviders(); expect(screen.getByText("ag-1")).toBeInTheDocument(); expect(screen.getByText("Admin Group")).toBeInTheDocument(); @@ -115,57 +113,70 @@ describe("AccessGroupsPage", () => { expect(screen.getByText("Read Only")).toBeInTheDocument(); }); - it("should display resource counts for each group", () => { + it("renders resource counts for each group", () => { renderWithProviders(); - const table = screen.getByRole("table"); - expect(table).toHaveTextContent("2"); - expect(table).toHaveTextContent("1"); + // ag-1 has 2 models, 1 mcp server, 1 agent. + const adminRow = screen.getByText("ag-1").closest("tr") as HTMLElement; + expect(within(adminRow).getByTitle("2 Models")).toHaveTextContent("2"); + expect(within(adminRow).getByTitle("1 MCP Servers")).toHaveTextContent("1"); + expect(within(adminRow).getByTitle("1 Agents")).toHaveTextContent("1"); }); - it("should filter groups by search text matching name", async () => { + it("shows the expected column headers", () => { + renderWithProviders(); + expect(screen.getByRole("columnheader", { name: /^ID$/i })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: /Name/i })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: /Resources/i })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: /Created/i })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: /Updated/i })).toBeInTheDocument(); + }); + + it("filters by name", async () => { const user = userEvent.setup(); renderWithProviders(); - const searchInput = screen.getByPlaceholderText("Search groups by name, ID, or description..."); - await user.type(searchInput, "Admin"); + await user.type(screen.getByPlaceholderText("Search groups by name, ID, or description..."), "Admin"); expect(screen.getByText("Admin Group")).toBeInTheDocument(); expect(screen.queryByText("Read Only")).not.toBeInTheDocument(); }); - it("should filter groups by search text matching ID", async () => { + it("filters by ID", async () => { const user = userEvent.setup(); renderWithProviders(); - const searchInput = screen.getByPlaceholderText("Search groups by name, ID, or description..."); - await user.type(searchInput, "ag-2"); + await user.type(screen.getByPlaceholderText("Search groups by name, ID, or description..."), "ag-2"); expect(screen.getByText("Read Only")).toBeInTheDocument(); expect(screen.queryByText("Admin Group")).not.toBeInTheDocument(); }); - it("should filter groups by search text matching description", async () => { + it("filters by description", async () => { const user = userEvent.setup(); renderWithProviders(); - const searchInput = screen.getByPlaceholderText("Search groups by name, ID, or description..."); - await user.type(searchInput, "read-only"); + await user.type(screen.getByPlaceholderText("Search groups by name, ID, or description..."), "read-only"); expect(screen.getByText("Read Only")).toBeInTheDocument(); expect(screen.queryByText("Admin Group")).not.toBeInTheDocument(); }); - it("should reset to first page when search text changes", async () => { + it("shows the filtered empty state when nothing matches", async () => { const user = userEvent.setup(); renderWithProviders(); - const searchInput = screen.getByPlaceholderText("Search groups by name, ID, or description..."); - await user.type(searchInput, "Admin"); - const pagination = screen.getByText(/groups/); - expect(pagination).toHaveTextContent("1 groups"); + await user.type(screen.getByPlaceholderText("Search groups by name, ID, or description..."), "no-such-group"); + expect(screen.getByText("No matching access groups")).toBeInTheDocument(); + expect(screen.queryByText("Admin Group")).not.toBeInTheDocument(); }); - it("should open create modal when Create Access Group button is clicked", async () => { - const user = userEvent.setup(); + it("shows the empty state when there are no groups", () => { + mockUseAccessGroups.mockReturnValue({ data: [], isLoading: false }); renderWithProviders(); - await user.click(screen.getByRole("button", { name: /create access group/i })); - expect(screen.getByTestId("create-access-group-modal")).toBeInTheDocument(); + expect(screen.getByText("No access groups yet")).toBeInTheDocument(); }); - it("should close create modal when cancel is clicked", async () => { + it("renders loading skeletons on the initial load", () => { + mockUseAccessGroups.mockReturnValue({ data: undefined, isLoading: true }); + renderWithProviders(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); + expect(screen.queryByText("Admin Group")).not.toBeInTheDocument(); + }); + + it("opens and closes the create modal", async () => { const user = userEvent.setup(); renderWithProviders(); await user.click(screen.getByRole("button", { name: /create access group/i })); @@ -174,33 +185,22 @@ describe("AccessGroupsPage", () => { expect(screen.queryByTestId("create-access-group-modal")).not.toBeInTheDocument(); }); - it("should navigate to detail view when group ID is clicked", async () => { + it("opens the detail view when the ID cell is clicked and returns via Back", async () => { const user = userEvent.setup(); renderWithProviders(); await user.click(screen.getByText("ag-1")); expect(screen.getByTestId("access-group-detail")).toBeInTheDocument(); expect(screen.getByText("Detail for ag-1")).toBeInTheDocument(); - }); - - it("should return to list view when Back is clicked from detail", async () => { - const user = userEvent.setup(); - renderWithProviders(); - await user.click(screen.getByText("ag-1")); - expect(screen.getByTestId("access-group-detail")).toBeInTheDocument(); await user.click(screen.getByRole("button", { name: "Back" })); expect(screen.queryByTestId("access-group-detail")).not.toBeInTheDocument(); expect(screen.getByText("Admin Group")).toBeInTheDocument(); }); - it("should open delete modal when delete action is clicked", async () => { + it("opens the delete modal from the row actions menu", async () => { const user = userEvent.setup(); renderWithProviders(); - const deleteButtons = screen.getAllByRole("button", { - name: "Delete access group", - }); - await user.click(deleteButtons[0]); + await user.click(await openRowMenu(user, "ag-1")); const dialog = screen.getByRole("dialog", { name: "Delete Access Group" }); - expect(dialog).toBeInTheDocument(); expect( within(dialog).getByText("Are you sure you want to delete this access group? This action cannot be undone."), ).toBeInTheDocument(); @@ -209,71 +209,49 @@ describe("AccessGroupsPage", () => { expect(within(dialog).getByText("Admin Group")).toBeInTheDocument(); }); - it("should close delete modal when cancel is clicked", async () => { + it("closes the delete modal on cancel without deleting", async () => { const user = userEvent.setup(); renderWithProviders(); - const deleteButtons = screen.getAllByRole("button", { - name: "Delete access group", - }); - await user.click(deleteButtons[0]); + await user.click(await openRowMenu(user, "ag-1")); const dialog = screen.getByRole("dialog", { name: "Delete Access Group" }); await user.click(within(dialog).getByRole("button", { name: "Cancel" })); expect(screen.queryByRole("dialog", { name: "Delete Access Group" })).not.toBeInTheDocument(); + expect(mockMutate).not.toHaveBeenCalled(); }); - it("should call delete mutation when delete is confirmed", async () => { + it("calls the delete mutation with the group ID when confirmed", async () => { const user = userEvent.setup(); mockMutate.mockImplementation((_id: string, opts?: { onSuccess?: () => void }) => { opts?.onSuccess?.(); }); renderWithProviders(); - const deleteButtons = screen.getAllByRole("button", { - name: "Delete access group", - }); - await user.click(deleteButtons[0]); + await user.click(await openRowMenu(user, "ag-1")); const dialog = screen.getByRole("dialog", { name: "Delete Access Group" }); - const deleteConfirmButton = within(dialog).getByRole("button", { name: /delete/i }); - await user.click(deleteConfirmButton); + await user.click(within(dialog).getByRole("button", { name: /delete/i })); expect(mockMutate).toHaveBeenCalledWith("ag-1", expect.any(Object)); }); - it("should display pagination with total count", () => { - renderWithProviders(); - expect(screen.getByText("2 groups")).toBeInTheDocument(); - }); - - it("should show table headers for ID, Name, Resources, and Actions", () => { - renderWithProviders(); - expect(screen.getByRole("columnheader", { name: /ID/i })).toBeInTheDocument(); - expect(screen.getByRole("columnheader", { name: /Name/i })).toBeInTheDocument(); - expect(screen.getByRole("columnheader", { name: /Resources/i })).toBeInTheDocument(); - expect(screen.getByRole("columnheader", { name: /Actions/i })).toBeInTheDocument(); - }); - - it("should display loading state when data is loading", () => { - mockUseAccessGroups.mockReturnValue({ - data: undefined, - isLoading: true, - }); - renderWithProviders(); - const table = screen.getByRole("table"); - expect(table).toBeInTheDocument(); - }); - - it("should display empty state when no groups match search", async () => { + it("still shows matches when searching from a later page", async () => { const user = userEvent.setup(); + mockUseAccessGroups.mockReturnValue({ data: makeGroups(25), isLoading: false }); renderWithProviders(); - const searchInput = screen.getByPlaceholderText("Search groups by name, ID, or description..."); - await user.type(searchInput, "nonexistent-group-xyz"); - expect(screen.getByRole("table")).toBeInTheDocument(); + + await user.click(screen.getByTestId("pagination-next")); + expect(screen.getByText("ag-11")).toBeInTheDocument(); + expect(screen.queryByText("ag-01")).not.toBeInTheDocument(); + + // The only match lives on page 1, so the page index must reset or the table reads as empty. + await user.type(screen.getByPlaceholderText("Search groups by name, ID, or description..."), "ag-01"); + expect(await screen.findByText("ag-01")).toBeInTheDocument(); + expect(screen.queryByText("No matching access groups")).not.toBeInTheDocument(); }); - it("should display empty data when useAccessGroups returns empty array", () => { - mockUseAccessGroups.mockReturnValue({ - data: [], - isLoading: false, - }); + it("hides the Create button and row actions for a non-admin", () => { + mockUseAuthorized.mockReturnValue({ userRole: "Admin Viewer", accessToken: "sk-test" }); renderWithProviders(); - expect(screen.getByRole("table")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /create access group/i })).not.toBeInTheDocument(); + expect(screen.queryByTestId("access-group-actions-ag-1")).not.toBeInTheDocument(); + // The read-only view still lists the groups. + expect(screen.getByText("Admin Group")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx index dbbf4e35900..0de6596f57c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx @@ -1,38 +1,17 @@ import { AccessGroupResponse, useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups"; import { useDeleteAccessGroup } from "@/app/(dashboard)/hooks/accessGroups/useDeleteAccessGroup"; import { PlusOutlined } from "@ant-design/icons"; -import { - ColumnDef, - flexRender, - getCoreRowModel, - getSortedRowModel, - Row, - SortingState, - useReactTable, -} from "@tanstack/react-table"; -import { Button, Card, Flex, Input, Layout, Pagination, Space, Table, Tag, theme, Tooltip, Typography } from "antd"; -import { BotIcon, LayersIcon, SearchIcon, ServerIcon } from "lucide-react"; -import { useEffect, useMemo, useState } from "react"; +import { Button, Flex, Input, Layout, Space, theme, Typography } from "antd"; +import { SearchIcon } from "lucide-react"; +import { useMemo, useState } from "react"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import TableIconActionButton from "@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; -import { - SortState, - TableHeaderSortDropdown, -} from "@/components/common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; -import { DateCell, IdCell } from "@/components/shared/table_cells"; import { AccessGroupDetail } from "./AccessGroupsDetailsPage"; import { AccessGroupCreateModal } from "./AccessGroupsModal/AccessGroupCreateModal"; +import { AccessGroupsTable } from "./AccessGroupsTable"; import { AccessGroup } from "./types"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { isProxyAdminRole } from "@/utils/roles"; -declare module "@tanstack/react-table" { - // eslint-disable-next-line @typescript-eslint/no-unused-vars - interface ColumnMeta { - responsive?: string[]; - } -} - const { Title, Text } = Typography; const { Content } = Layout; @@ -52,55 +31,6 @@ function mapResponseToAccessGroup(r: AccessGroupResponse): AccessGroup { updatedBy: r.updated_by ?? "", }; } -function buildAntdColumns( - table: ReturnType>, - rowLookup: Map>, - onSortingChange: (s: SortingState) => void, -) { - const headers = table.getHeaderGroups()[0]?.headers ?? []; - - return headers.map((header) => { - const canSort = header.column.getCanSort(); - const isSorted = header.column.getIsSorted(); - const meta = header.column.columnDef.meta as { responsive?: string[] } | undefined; - - const col: Record = { - title: ( -
- {header.isPlaceholder ? null : flexRender(header.column.columnDef.header, header.getContext())} - {canSort && ( - { - if (newState === false) { - onSortingChange([]); - } else { - onSortingChange([{ id: header.column.id, desc: newState === "desc" }]); - } - }} - columnId={header.column.id} - /> - )} -
- ), - key: header.id, - width: header.column.columnDef.size, - render: (_: unknown, record: AccessGroup) => { - const row = rowLookup.get(record.id); - if (!row) return null; - const cell = row.getVisibleCells().find((c) => c.column.id === header.id); - if (!cell) return null; - return flexRender(cell.column.columnDef.cell, cell.getContext()); - }, - }; - - if (meta?.responsive) { - col.responsive = meta.responsive; - } - - return col; - }); -} export function AccessGroupsPage() { const { token } = theme.useToken(); @@ -113,151 +43,19 @@ export function AccessGroupsPage() { const [selectedGroupId, setSelectedGroupId] = useState(null); const [isCreateModalVisible, setIsCreateModalVisible] = useState(false); const [searchText, setSearchText] = useState(""); - const [currentPage, setCurrentPage] = useState(1); - const [sorting, setSorting] = useState([]); const [groupToDelete, setGroupToDelete] = useState(null); const deleteMutation = useDeleteAccessGroup(); - const pageSize = 10; - useEffect(() => { - setCurrentPage(1); - }, [searchText]); - - // ---------- filtered data ---------- - const filteredGroups = useMemo( - () => - groups.filter( - (group) => - group.name.toLowerCase().includes(searchText.toLowerCase()) || - group.id.toLowerCase().includes(searchText.toLowerCase()) || - group.description.toLowerCase().includes(searchText.toLowerCase()), - ), - [groups, searchText], - ); - - // ---------- TanStack column definitions ---------- - const columnDefs = useMemo[]>( - () => [ - { - id: "id", - accessorKey: "id", - header: () => ID, - enableSorting: false, - size: 170, - cell: ({ row }) => , - }, - { - id: "name", - accessorKey: "name", - header: () => Name, - enableSorting: true, - cell: ({ getValue }) => getValue() as string, - }, - { - id: "resources", - header: () => Resources, - enableSorting: false, - cell: ({ row }) => { - const record = row.original; - const modelIds = record.modelIds ?? []; - const mcpServerIds = record.mcpServerIds ?? []; - const agentIds = record.agentIds ?? []; - return ( - - - - - - {modelIds?.length} - - - - - - - - {mcpServerIds?.length} - - - - - - - - {agentIds?.length} - - - - - ); - }, - }, - { - id: "createdAt", - accessorKey: "createdAt", - header: () => Created, - enableSorting: true, - sortingFn: "datetime", - cell: ({ getValue }) => , - meta: { responsive: ["lg"] }, - }, - { - id: "updatedAt", - accessorKey: "updatedAt", - header: () => Updated, - enableSorting: false, - cell: ({ getValue }) => , - meta: { responsive: ["xl"] }, - }, - ...(canModify - ? [ - { - id: "actions", - header: () => Actions, - enableSorting: false, - cell: ({ row }: { row: Row }) => ( - - setGroupToDelete(row.original)} - /> - - ), - }, - ] - : []), - ], - // setSelectedGroup is stable (useState setter) - // eslint-disable-next-line react-hooks/exhaustive-deps - [canModify], - ); - - // ---------- TanStack table instance ---------- - const table = useReactTable({ - data: filteredGroups, - columns: columnDefs, - state: { sorting }, - onSortingChange: setSorting, - getCoreRowModel: getCoreRowModel(), - getSortedRowModel: getSortedRowModel(), - getRowId: (row) => row.id, - }); - - // All sorted rows from TanStack - const sortedRows = table.getRowModel().rows; - - // Paginated slice - const paginatedRows = sortedRows.slice((currentPage - 1) * pageSize, currentPage * pageSize); - - // Map for O(1) lookup by record id in antd render() - const rowLookup = useMemo(() => new Map(paginatedRows.map((row) => [row.original.id, row])), [paginatedRows]); - - // Convert TanStack headers → antd columns - const antdColumns = buildAntdColumns(table, rowLookup, setSorting); - - // antd dataSource (just the originals for the current page) - const dataSource = paginatedRows.map((row) => row.original); + const filteredGroups = useMemo(() => { + const query = searchText.trim().toLowerCase(); + if (!query) return groups; + return groups.filter( + (group) => + group.name.toLowerCase().includes(query) || + group.id.toLowerCase().includes(query) || + group.description.toLowerCase().includes(query), + ); + }, [groups, searchText]); if (selectedGroupId) { return setSelectedGroupId(null)} />; @@ -279,34 +77,25 @@ export function AccessGroupsPage() { )} - - - } - placeholder="Search groups by name, ID, or description..." - style={{ maxWidth: 400 }} - value={searchText} - onChange={(e) => setSearchText(e.target.value)} - allowClear - /> - setCurrentPage(page)} - size="small" - showTotal={(total) => `${total} groups`} - showSizeChanger={false} - /> - - - + + } + placeholder="Search groups by name, ID, or description..." + style={{ maxWidth: 400 }} + value={searchText} + onChange={(e) => setSearchText(e.target.value)} + allowClear + /> + + + 0} + canModify={canModify} + onGroupClick={setSelectedGroupId} + onDeleteClick={setGroupToDelete} + /> setIsCreateModalVisible(false)} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsTable.tsx new file mode 100644 index 00000000000..10d1735d3e7 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsTable.tsx @@ -0,0 +1,72 @@ +"use client"; + +import { SortingState } from "@tanstack/react-table"; +import { Layers } from "lucide-react"; +import { useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; + +import { getAccessGroupsTableColumns } from "./AccessGroupsTableColumns"; +import { AccessGroup } from "./types"; + +interface AccessGroupsTableProps { + groups: AccessGroup[]; + isLoading: boolean; + isFiltered: boolean; + canModify: boolean; + onGroupClick: (id: string) => void; + onDeleteClick: (group: AccessGroup) => void; +} + +const PAGE_SIZE_OPTIONS = [10, 25, 50]; + +function EmptyState({ isFiltered }: { isFiltered: boolean }) { + return ( +
+
+ +
+
+ {isFiltered ? "No matching access groups" : "No access groups yet"} +
+
+ {isFiltered + ? "Try a different search term." + : "Create an access group to manage resource permissions for your organization."} +
+
+ ); +} + +export function AccessGroupsTable({ + groups, + isLoading, + isFiltered, + canModify, + onGroupClick, + onDeleteClick, +}: AccessGroupsTableProps) { + const [sorting, setSorting] = useState([]); + + const columns = useMemo(() => { + const deps = { canModify, onGroupClick, onDeleteClick }; + return getAccessGroupsTableColumns(deps); + }, [canModify, onGroupClick, onDeleteClick]); + + return ( + group.id || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + paginationMode="client" + pageSizeOptions={PAGE_SIZE_OPTIONS} + isLoading={isLoading} + loadingMessage="Loading access groups…" + noDataMessage={} + size="compact" + /> + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsTableColumns.tsx new file mode 100644 index 00000000000..ae65f161b1e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsTableColumns.tsx @@ -0,0 +1,182 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { Bot, Layers, MoreHorizontal, Server, Trash2 } from "lucide-react"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdentityCell } from "@/components/shared/table_cells"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; + +import { AccessGroup } from "./types"; + +interface ResourceTone { + icon: typeof Layers; + className: string; +} + +const RESOURCE_TONES: Record<"models" | "mcpServers" | "agents", ResourceTone> = { + models: { icon: Layers, className: "bg-blue-50 text-blue-700 ring-blue-600/20" }, + mcpServers: { icon: Server, className: "bg-cyan-50 text-cyan-700 ring-cyan-600/20" }, + agents: { icon: Bot, className: "bg-purple-50 text-purple-700 ring-purple-600/20" }, +}; + +function ResourcesCell({ group }: { group: AccessGroup }) { + const items = [ + { key: "models" as const, label: "Models", count: group.modelIds.length }, + { key: "mcpServers" as const, label: "MCP Servers", count: group.mcpServerIds.length }, + { key: "agents" as const, label: "Agents", count: group.agentIds.length }, + ]; + + return ( +
+ {items.map((item) => { + const tone = RESOURCE_TONES[item.key]; + const Icon = tone.icon; + return ( + + + {item.count} + + ); + })} +
+ ); +} + +function AccessGroupRowActions({ + group, + onDeleteClick, +}: { + group: AccessGroup; + onDeleteClick: (group: AccessGroup) => void; +}) { + return ( + + + + + + onDeleteClick(group)} + > + + Delete access group + + + + ); +} + +interface AccessGroupsTableColumnsDeps { + canModify: boolean; + onGroupClick: (id: string) => void; + onDeleteClick: (group: AccessGroup) => void; +} + +export const getAccessGroupsTableColumns = ({ + canModify, + onGroupClick, + onDeleteClick, +}: AccessGroupsTableColumnsDeps): ColumnDef[] => { + const columns: ColumnDef[] = [ + { + id: "id", + accessorKey: "id", + meta: { title: "ID" }, + header: "ID", + size: 200, + enableSorting: false, + cell: ({ row }) => ( + onGroupClick(row.original.id)} + /> + ), + }, + { + id: "name", + accessorKey: "name", + meta: { title: "Name" }, + header: ({ column }) => , + size: 220, + enableSorting: true, + cell: ({ row }) => { + const name = row.original.name; + return ( + + {name || "-"} + + ); + }, + }, + { + id: "resources", + meta: { title: "Resources" }, + header: "Resources", + size: 220, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "createdAt", + accessorKey: "createdAt", + meta: { title: "Created" }, + header: ({ column }) => , + size: 150, + enableSorting: true, + sortingFn: "datetime", + cell: ({ row }) => , + }, + { + id: "updatedAt", + accessorKey: "updatedAt", + meta: { title: "Updated" }, + header: "Updated", + size: 150, + enableSorting: false, + cell: ({ row }) => , + }, + ]; + + if (!canModify) { + return columns; + } + + return [ + ...columns, + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + }, + ]; +}; From 049c6836d205d024e3beb7ab881ec62c289cd142 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 21 Jul 2026 10:28:24 -0700 Subject: [PATCH 12/84] fix(model_armor): sanitize error details by default (#33908) * fix(model_armor): sanitize error details by default Generated with AI Co-Authored-By: Claude Code * fix(model_armor): sanitize handler-raised HTTP errors and redact scanned content in guardrail logging The async HTTP handler raises MaskedHTTPStatusError on any non-2xx via raise_for_status, so the non-200 branch in make_model_armor_request never ran against a live API and the raw upstream body reached callers and logs. Catch the raised error and build the sanitized detail from the response status Replace the empty-dict guardrail logging payload with field-level redaction of the keys that echo scanned content (text, sanitizedText, findings) so guardrail traces keep filter states and block reasons while scanned content stays out Restore the upstream status code in the sanitized error detail, read guardrail metadata from the same key the hooks write, and keep guardrail_status within its typed literal values * fix(model_armor): bound redactor recursion depth and allowlist it in the recursion detector _redact_scanned_content walks provider JSON bounded by _REDACT_MAX_DEPTH=20 and fails closed by returning the redaction sentinel at the cap * fix(model_armor): honor fail_on_error for upstream API failures API failures now raise a dedicated ModelArmorAPIError so hooks can tell them apart from content-block HTTPExceptions; fail_on_error=False lets the request proceed on a Model Armor outage again while fail-closed configs get the same sanitized 400 as before Also addresses review notes: sanitize_error_detail constructor annotation matches the nullable config field, redaction is owned by the metadata write sites so _process_response no longer re-applies it, and the request and response debug log branches move into helpers * test(model_armor): cover fail_on_error routing on during-call, post-call, streaming, and file-scan paths * chore: remove accidentally committed pytest cache files * fix(model_armor): keep sanitize_error_detail coerced across in-memory config reloads update_in_memory_litellm_params assigns raw LitellmParams fields, so a hot reloaded config carrying an explicit null would silently disable sanitization; re-apply the only-explicit-False-opts-out coercion after the update * fix(model_armor): redact matched malicious URIs and reuse the shared recursion depth constant maliciousUriMatchedItems echoes the caller-supplied URL including path and query, so it joins the scanned-content key set; the redactor depth cap now comes from DEFAULT_MAX_RECURSE_DEPTH in litellm constants instead of a local literal * fix(model_armor): keep API failures out of the intervention trace status Fail-closed upstream failures re-raise ModelArmorAPIError instead of converting to HTTPException(400), so the shared guardrail logging keeps recording them as guardrail_failed_to_respond while content blocks stay guardrail_intervened. Callers see the same 500 shape as before this PR, with the sanitized message * chore(model_armor): drop explanatory comment per repository comment policy --------- Co-authored-by: eugene-yao-zocdoc --- .../guardrail_hooks/model_armor/__init__.py | 1 + .../model_armor/model_armor.py | 202 +++++-- litellm/types/guardrails.py | 7 + .../guardrails/guardrail_hooks/model_armor.py | 7 + .../code_coverage_tests/recursive_detector.py | 1 + .../guardrail_hooks/test_model_armor.py | 557 +++++++++++++++++- 6 files changed, 706 insertions(+), 69 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py index 5e62ab96f0c..d91ddffa0c1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py @@ -27,6 +27,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" mask_response_content=litellm_params.mask_response_content, fail_on_error=litellm_params.fail_on_error, skip_unscannable_attachments=litellm_params.skip_unscannable_attachments, + sanitize_error_detail=litellm_params.sanitize_error_detail, ) litellm.logging_callback_manager.add_litellm_callback(_model_armor_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 32a3cebfca0..31535a5b569 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -11,6 +11,7 @@ from typing import ( Union, ) +import httpx from fastapi import HTTPException if TYPE_CHECKING: @@ -35,7 +36,8 @@ from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import ( MODEL_ARMOR_MAX_FILE_SIZE_BYTES, plan_file_scans, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ( CallTypes, @@ -50,6 +52,33 @@ from litellm.types.utils import ( GUARDRAIL_NAME = "model_armor" +class ModelArmorAPIError(Exception): + """Model Armor API failure (non-2xx), distinct from a content-block decision so + hooks can honor fail_on_error. The detail is already sanitized per configuration.""" + + def __init__(self, detail: str): + super().__init__(detail) + self.detail = detail + + +_SCANNED_CONTENT_KEYS = frozenset({"text", "sanitizedText", "findings", "maliciousUriMatchedItems"}) + +RedactablePayload = Union[dict, list, str, int, float, bool, None] + + +def _redact_scanned_content(payload: RedactablePayload, depth: int = 0) -> RedactablePayload: + if depth >= DEFAULT_MAX_RECURSE_DEPTH: + return "[REDACTED]" + if isinstance(payload, dict): + return { + key: "[REDACTED]" if key in _SCANNED_CONTENT_KEYS else _redact_scanned_content(value, depth + 1) + for key, value in payload.items() + } + if isinstance(payload, list): + return [_redact_scanned_content(item, depth + 1) for item in payload] + return payload + + class ModelArmorGuardrail(CustomGuardrail, VertexBase): """ Google Cloud Model Armor Guardrail integration for LiteLLM. @@ -76,6 +105,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): location: Optional[str] = None, credentials: Optional[Any] = None, api_endpoint: Optional[str] = None, + sanitize_error_detail: "bool | None" = True, **kwargs, ): # Set supported event hooks if not already provided @@ -98,6 +128,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): self.location = location or "us-central1" self.credentials = credentials self.api_endpoint = api_endpoint + self.sanitize_error_detail = sanitize_error_detail is not False # Store optional params self.optional_params = kwargs @@ -141,6 +172,67 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): verbose_proxy_logger.debug("Model Armor: Skipping non-ModelResponse type: %s", type(response).__name__) return "" + def _build_api_error_detail(self, status_code: int, response_text: str) -> str: + if self.sanitize_error_detail: + return f"Model Armor API error (upstream {status_code})" + return f"Model Armor API error (upstream {status_code}): {response_text}" + + def _build_block_error_detail(self, message: str, armor_response: RedactablePayload) -> dict: + if self.sanitize_error_detail: + return {"error": message} + return {"error": message, "model_armor_response": armor_response} + + def _build_logging_response(self, armor_response: RedactablePayload) -> RedactablePayload: + if self.sanitize_error_detail: + return _redact_scanned_content(armor_response) + return armor_response + + def _raise_if_fail_closed(self, e: ModelArmorAPIError) -> None: + if self.optional_params.get("fail_on_error", True): + raise e from None + + def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: + super().update_in_memory_litellm_params(litellm_params) + self.sanitize_error_detail = self.sanitize_error_detail is not False + + def _log_request_debug( + self, + url: str, + body: dict, + file_bytes: "bytes | None", + file_type: "str | None", + ) -> None: + # Never log byteData: it is the full base64 of the scanned document. Log only its + # type and size so debug deployments cannot leak the contents the guardrail inspects. + if file_bytes is not None and file_type is not None: + verbose_proxy_logger.debug( + "Model Armor file request - URL: %s, byteDataType: %s, bytes: %d", + url, + file_type, + len(file_bytes), + ) + elif self.sanitize_error_detail: + verbose_proxy_logger.debug("Model Armor request - URL: %s", url) + else: + verbose_proxy_logger.debug( + "Model Armor request - URL: %s, Body: %s", + url, + body, + ) + + def _log_response_debug(self, status_code: int, response_text: str) -> None: + if self.sanitize_error_detail: + verbose_proxy_logger.debug( + "Model Armor response - Status: %s", + status_code, + ) + else: + verbose_proxy_logger.debug( + "Model Armor response - Status: %s, Body: %s", + status_code, + response_text, + ) + async def make_model_armor_request( self, content: Optional[str] = None, @@ -185,48 +277,37 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): "Authorization": f"Bearer {access_token}", } - # Never log byteData: it is the full base64 of the scanned document. Log only its - # type and size so debug deployments cannot leak the contents the guardrail inspects. - if file_bytes is not None and file_type is not None: - verbose_proxy_logger.debug( - "Model Armor file request - URL: %s, byteDataType: %s, bytes: %d", - url, - file_type, - len(file_bytes), - ) - else: - verbose_proxy_logger.debug( - "Model Armor request - URL: %s, Body: %s", - url, - body, - ) + self._log_request_debug(url=url, body=body, file_bytes=file_bytes, file_type=file_type) # Make request if self.async_handler is None: raise ValueError("Async handler not initialized") - response = await self.async_handler.post( - url=url, - json=body, - headers=headers, - ) + try: + response = await self.async_handler.post( + url=url, + json=body, + headers=headers, + ) + except httpx.HTTPStatusError as e: + detail = self._build_api_error_detail(e.response.status_code, e.response.text) + verbose_proxy_logger.error( + "Model Armor API error - Status: %s, Detail: %s", + e.response.status_code, + detail, + ) + raise ModelArmorAPIError(detail) from None - verbose_proxy_logger.debug( - "Model Armor response - Status: %s, Body: %s", - response.status_code, - response.text, - ) + self._log_response_debug(status_code=response.status_code, response_text=response.text) if response.status_code != 200: + detail = self._build_api_error_detail(response.status_code, response.text) verbose_proxy_logger.error( - "Model Armor API error - Status: %s, Response: %s", + "Model Armor API error - Status: %s, Detail: %s", response.status_code, - response.text, - ) - raise HTTPException( - status_code=400, - detail=f"Model Armor API error (upstream {response.status_code}): {response.text}", + detail, ) + raise ModelArmorAPIError(detail) json_response = response.json() if hasattr(json_response, "__await__"): @@ -351,9 +432,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): Override to store only the Model Armor API response, not the entire data dict. This prevents circular references in logging. """ - # Retrieve the Model Armor response & status stored on the per-request `metadata` object. metadata = request_data.get("metadata", {}) if isinstance(request_data, dict) else {} - guardrail_response = metadata.get("_model_armor_response", {}) # Determine status – default to "success" but prefer the explicit value if present. @@ -444,6 +523,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): file_bytes=attachment.file_bytes, file_type=attachment.byte_data_type, ) + except ModelArmorAPIError as e: + self._raise_if_fail_closed(e) + continue except HTTPException: raise except Exception as e: @@ -459,7 +541,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # otherwise a PII-only (SDP deidentify) document would pass through unscrubbed. blocked = self._should_block_content(armor_response, allow_sanitization=False) metadata["_model_armor_response"] = self._append_armor_response( - metadata.get("_model_armor_response"), armor_response + metadata.get("_model_armor_response"), + self._build_logging_response(armor_response), ) if blocked or metadata.get("_model_armor_status") == "blocked": metadata["_model_armor_status"] = "blocked" @@ -469,10 +552,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if blocked: raise HTTPException( status_code=400, - detail={ - "error": "Content blocked by Model Armor", - "model_armor_response": armor_response, - }, + detail=self._build_block_error_detail("Content blocked by Model Armor", armor_response), ) @log_guardrail_information @@ -530,7 +610,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): metadata = data.setdefault("metadata", {}) # ensures metadata exists and is unique per request # Accumulate so a prior file scan on the same request is not overwritten by this text scan. metadata["_model_armor_response"] = self._append_armor_response( - metadata.get("_model_armor_response"), armor_response + metadata.get("_model_armor_response"), + self._build_logging_response(armor_response), ) # Pre-compute guardrail status for downstream logging. A blocked response will eventually raise # an HTTPException, however in scenarios where the caller decides to ignore the exception (e.g. @@ -548,10 +629,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if blocked: raise HTTPException( status_code=400, - detail={ - "error": "Content blocked by Model Armor", - "model_armor_response": armor_response, - }, + detail=self._build_block_error_detail("Content blocked by Model Armor", armor_response), ) # If mask_request_content is enabled, update messages with sanitized content @@ -565,6 +643,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): data["messages"] = set_last_user_message(messages, sanitized_content) + except ModelArmorAPIError as e: + self._raise_if_fail_closed(e) except HTTPException: raise except Exception as e: @@ -625,7 +705,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): metadata = data.setdefault("metadata", {}) # Accumulate so a prior file scan on the same request is not overwritten by this text scan. metadata["_model_armor_response"] = self._append_armor_response( - metadata.get("_model_armor_response"), armor_response + metadata.get("_model_armor_response"), + self._build_logging_response(armor_response), ) if blocked or metadata.get("_model_armor_status") == "blocked": metadata["_model_armor_status"] = "blocked" @@ -640,10 +721,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if blocked: raise HTTPException( status_code=400, - detail={ - "error": "Content blocked by Model Armor", - "model_armor_response": armor_response, - }, + detail=self._build_block_error_detail("Content blocked by Model Armor", armor_response), ) # If mask_request_content is enabled, update messages with sanitized content @@ -656,6 +734,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): data["messages"] = set_last_user_message(messages, sanitized_content) + except ModelArmorAPIError as e: + self._raise_if_fail_closed(e) except HTTPException: raise except Exception as e: @@ -698,7 +778,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Attach Model Armor response & status to this request's metadata to prevent race conditions if isinstance(armor_response, dict): model_armor_logged_object = { - "model_armor_response": armor_response, + "model_armor_response": self._build_logging_response(armor_response), "model_armor_status": ( "blocked" if self._should_block_content( @@ -729,10 +809,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content): raise HTTPException( status_code=400, - detail={ - "error": "Response blocked by Model Armor", - "model_armor_response": armor_response, - }, + detail=self._build_block_error_detail("Response blocked by Model Armor", armor_response), ) # If mask_response_content is enabled, update response with sanitized content @@ -746,6 +823,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if choice.message.content: choice.message.content = sanitized_content + except ModelArmorAPIError as e: + self._raise_if_fail_closed(e) except HTTPException: raise except Exception as e: @@ -790,7 +869,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # Attach Model Armor response & status to this request's metadata to avoid race conditions if isinstance(request_data, dict): metadata = request_data.setdefault("metadata", {}) - metadata["_model_armor_response"] = armor_response + metadata["_model_armor_response"] = self._build_logging_response(armor_response) metadata["_model_armor_status"] = ( "blocked" if self._should_block_content(armor_response) else "success" ) @@ -809,10 +888,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if self._should_block_content(armor_response): raise HTTPException( status_code=400, - detail={ - "error": "Streaming response blocked by Model Armor", - "model_armor_response": armor_response, - }, + detail=self._build_block_error_detail( + "Streaming response blocked by Model Armor", + armor_response, + ), ) # Apply sanitization if enabled @@ -831,6 +910,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): yield chunk return + except ModelArmorAPIError as e: + if self.optional_params.get("fail_on_error", True): + error_obj = {"message": e.detail, "code": "500"} + yield f"data: {json.dumps({'error': error_obj})}\n\n" + return except HTTPException as e: # Yield error as SSE event so create_response() detects it and # returns a proper JSON error response with the correct status code. diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 3605ab95d1b..47d93fc2d7a 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -826,6 +826,13 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up "while fail_on_error still governs real Model Armor API errors. Default False blocks them." ), ) + sanitize_error_detail: Optional[bool] = Field( + default=True, + description=( + "For guardrail='model_armor': omit the raw Model Armor response from " + "caller-facing errors and logs by default. Set False to restore verbose output." + ), + ) additional_provider_specific_params: Optional[Dict[str, Any]] = Field( default=None, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py b/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py index 628ac0442de..d5e601ce8ea 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py @@ -20,6 +20,13 @@ class ModelArmorGuardrailConfigModel(GuardrailConfigModel): default=True, description="Whether to fail the request if Model Armor encounters an error", ) + sanitize_error_detail: Optional[bool] = Field( + default=True, + description=( + "Omit the raw Model Armor response from caller-facing errors and logs " + "by default. Set False to restore verbose output." + ), + ) @staticmethod def ui_friendly_name() -> str: diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index e08d703d21f..fa81efde5db 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -55,6 +55,7 @@ IGNORE_FUNCTIONS = [ "_freeze_for_dedupe", # OTEL: max depth set (default 16, _FREEZE_MAX_DEPTH); fails closed by returning repr(value) at the cap. "apply_json_merge_patch", # max depth set (_MAX_MERGE_DEPTH=64); fails closed by raising ValueError at the cap. "_filter_mcp_argument_value", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by blocking the MCP call at the cap. + "_redact_scanned_content", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by returning "[REDACTED]" at the cap. ] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 07c40aa763d..4021f922877 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -10,14 +10,19 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) +import httpx from fastapi import HTTPException import litellm import litellm.types.utils from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache +from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail +from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ( + ModelArmorAPIError, +) from litellm.types.guardrails import GuardrailEventHooks @@ -403,8 +408,9 @@ async def test_model_armor_api_error_handling(): "metadata": {"guardrails": ["model-armor-test"]}, } - # Should raise HTTPException for API error - with pytest.raises(HTTPException) as exc_info: + # An API failure propagates as ModelArmorAPIError, not a content-block + # HTTPException, so guardrail trace status stays guardrail_failed_to_respond + with pytest.raises(ModelArmorAPIError) as exc_info: await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, cache=mock_cache, @@ -412,9 +418,8 @@ async def test_model_armor_api_error_handling(): call_type="completion", ) - assert exc_info.value.status_code == 400 - assert "Model Armor API error" in str(exc_info.value.detail) - assert "upstream 500" in str(exc_info.value.detail) + assert exc_info.value.detail == "Model Armor API error (upstream 500)" + assert "Internal Server Error" not in str(exc_info.value.detail) @pytest.mark.asyncio @@ -622,7 +627,7 @@ async def test_model_armor_streaming_block_yields_sse_error(): @pytest.mark.asyncio -async def test_model_armor_api_failure_returns_400(): +async def test_model_armor_api_failure_raises_sanitized_error(): """Test that Model Armor API failures raise HTTP 400, not the upstream status code.""" guardrail = ModelArmorGuardrail( template_id="test-template", @@ -643,15 +648,544 @@ async def test_model_armor_api_failure_returns_400(): with patch.object( guardrail.async_handler, "post", AsyncMock(return_value=mock_response) ): - with pytest.raises(HTTPException) as exc_info: + with pytest.raises(ModelArmorAPIError) as exc_info: await guardrail.make_model_armor_request( content="test content", source="user_prompt", ) - # Should be 400, NOT the upstream 500 - assert exc_info.value.status_code == 400 - assert "upstream 500" in str(exc_info.value.detail) + assert exc_info.value.detail == "Model Armor API error (upstream 500)" + assert "Internal Server Error" not in str(exc_info.value.detail) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sanitize", [True, False]) +async def test_model_armor_error_output_sanitization(sanitize: bool): + marker = "SYNTHETIC_MODEL_ARMOR_MARKER" + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + sanitize_error_detail=sanitize, + ) + guardrail._ensure_access_token_async = AsyncMock( + return_value=("test-token", "test-project") + ) + + error_response = AsyncMock(status_code=500, text=marker) + with patch.object( + guardrail.async_handler, "post", AsyncMock(return_value=error_response) + ), patch.object(verbose_proxy_logger, "debug") as debug_log, patch.object( + verbose_proxy_logger, "error" + ) as error_log, pytest.raises(ModelArmorAPIError) as exc_info: + await guardrail.make_model_armor_request(content=marker) + + direct_log = f"{debug_log.call_args_list} {error_log.call_args_list}" + if sanitize: + assert marker not in str(exc_info.value.detail) + assert marker not in direct_log + else: + assert marker in str(exc_info.value.detail) + assert marker in direct_log + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_on_error", [True, False]) +async def test_model_armor_api_error_honors_fail_open(fail_on_error: bool): + """An upstream API failure (raised by the real handler as MaskedHTTPStatusError) + must block with a sanitized 400 when fail_on_error is true and let the request + proceed when the operator configured fail-open.""" + marker = "SYNTHETIC_FAIL_OPEN_MARKER" + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + fail_on_error=fail_on_error, + ) + guardrail._ensure_access_token_async = AsyncMock( + return_value=("test-token", "test-project") + ) + guardrail.should_run_guardrail = Mock(return_value=True) + + request = httpx.Request("POST", "https://modelarmor.example.test/v1") + upstream = httpx.Response(503, content=marker.encode(), request=request) + original = httpx.HTTPStatusError("Service Unavailable", request=request, response=upstream) + masked = MaskedHTTPStatusError(original, message=marker, text=marker) + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "synthetic input"}], + "metadata": {}, + } + + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=masked)): + if fail_on_error: + with pytest.raises(ModelArmorAPIError) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + assert exc_info.value.detail == "Model Armor API error (upstream 503)" + assert marker not in str(exc_info.value.detail) + else: + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + assert result is request_data + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_on_error", [True, False]) +async def test_model_armor_api_error_fail_open_moderation_and_post_call(fail_on_error: bool): + """The during-call and post-call hooks route API failures through fail_on_error + exactly like pre-call: sanitized 400 when failing closed, pass-through when open.""" + api_error = ModelArmorAPIError("Model Armor API error (upstream 503)") + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + fail_on_error=fail_on_error, + ) + guardrail.make_model_armor_request = AsyncMock(side_effect=api_error) + guardrail.should_run_guardrail = Mock(return_value=True) + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "synthetic input"}], + "metadata": {}, + } + mock_llm_response = litellm.ModelResponse() + mock_llm_response.choices = [ + litellm.Choices(message=litellm.Message(content="model output")) + ] + + if fail_on_error: + with pytest.raises(ModelArmorAPIError) as mod_exc: + await guardrail.async_moderation_hook( + data=dict(request_data), + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + assert mod_exc.value.detail == "Model Armor API error (upstream 503)" + + with pytest.raises(ModelArmorAPIError) as post_exc: + await guardrail.async_post_call_success_hook( + data=dict(request_data), + user_api_key_dict=UserAPIKeyAuth(), + response=mock_llm_response, + ) + assert post_exc.value.detail == "Model Armor API error (upstream 503)" + else: + moderated = await guardrail.async_moderation_hook( + data=dict(request_data), + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + assert moderated is not None + + result = await guardrail.async_post_call_success_hook( + data=dict(request_data), + user_api_key_dict=UserAPIKeyAuth(), + response=mock_llm_response, + ) + assert result is mock_llm_response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_on_error", [True, False]) +async def test_model_armor_api_error_fail_open_streaming(fail_on_error: bool): + """A streaming-path API failure yields a sanitized SSE error frame when failing + closed and passes the original chunks through when the operator opted into fail-open.""" + api_error = ModelArmorAPIError("Model Armor API error (upstream 503)") + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + fail_on_error=fail_on_error, + ) + guardrail.make_model_armor_request = AsyncMock(side_effect=api_error) + guardrail.should_run_guardrail = Mock(return_value=True) + + async def mock_stream(): + yield litellm.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content="streamed output") + ) + ] + ) + + chunks = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=mock_stream(), + request_data={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "synthetic input"}], + "metadata": {}, + }, + ): + chunks.append(chunk) + + if fail_on_error: + assert len(chunks) == 1 + assert isinstance(chunks[0], str) + assert "Model Armor API error (upstream 503)" in chunks[0] + assert '"code": "500"' in chunks[0] + else: + assert len(chunks) == 1 + assert isinstance(chunks[0], litellm.ModelResponseStream) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_on_error", [True, False]) +async def test_model_armor_api_error_fail_open_file_scan(fail_on_error: bool): + """A file-scan API failure blocks with the sanitized detail when failing closed + and skips the attachment when the operator opted into fail-open.""" + api_error = ModelArmorAPIError("Model Armor API error (upstream 503)") + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + fail_on_error=fail_on_error, + ) + guardrail.make_model_armor_request = AsyncMock(side_effect=api_error) + + pdf_b64 = base64.b64encode(b"%PDF-1.4 synthetic").decode() + messages = [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_data": f"data:application/pdf;base64,{pdf_b64}", + "filename": "synthetic.pdf", + "format": "application/pdf", + }, + } + ], + } + ] + data = {"metadata": {}} + + if fail_on_error: + with pytest.raises(ModelArmorAPIError) as exc_info: + await guardrail._scan_request_files(messages=messages, data=data) + assert exc_info.value.detail == "Model Armor API error (upstream 503)" + else: + assert await guardrail._scan_request_files(messages=messages, data=data) is None + + +def test_model_armor_hot_reload_null_stays_sanitized(): + """update_in_memory_litellm_params assigns raw fields; an explicit null in a + hot-reloaded config must not disable sanitization.""" + from litellm.types.guardrails import LitellmParams + + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + ) + guardrail.update_in_memory_litellm_params( + LitellmParams(guardrail="model_armor", mode="pre_call", sanitize_error_detail=None) + ) + assert guardrail.sanitize_error_detail is True + + guardrail.update_in_memory_litellm_params( + LitellmParams(guardrail="model_armor", mode="pre_call", sanitize_error_detail=False) + ) + assert guardrail.sanitize_error_detail is False + + +def test_model_armor_redactor_depth_cap_fails_closed(): + """Past the recursion cap the redactor must return the redaction sentinel, + never raw content, and must not raise RecursionError.""" + from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH + from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ( + _redact_scanned_content, + ) + + marker = "SYNTHETIC_DEEP_MARKER" + payload: dict = {"safe_key": marker, "items": [{"safe_key": marker}]} + for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 5): + payload = {"nested": payload} + + redacted = _redact_scanned_content(payload) + assert marker not in str(redacted) + + shallow = _redact_scanned_content({"filterResults": [{"text": marker, "matchState": "MATCH_FOUND"}]}) + assert shallow == {"filterResults": [{"text": "[REDACTED]", "matchState": "MATCH_FOUND"}]} + + uri_payload = _redact_scanned_content( + { + "maliciousUriFilterResult": { + "matchState": "MATCH_FOUND", + "maliciousUriMatchedItems": [{"uri": f"https://evil.example/{marker}"}], + } + } + ) + assert uri_payload == { + "maliciousUriFilterResult": { + "matchState": "MATCH_FOUND", + "maliciousUriMatchedItems": "[REDACTED]", + } + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sanitize", [True, False]) +async def test_model_armor_handler_raised_http_error_sanitized(sanitize: bool): + """The real AsyncHTTPHandler raises on non-2xx via raise_for_status, so a non-200 + never returns a response object. The raised MaskedHTTPStatusError carries the raw + upstream body in its message; the guardrail must convert it to a sanitized + HTTPException instead of letting it bubble raw to callers and logs.""" + marker = "SYNTHETIC_MODEL_ARMOR_MARKER" + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + sanitize_error_detail=sanitize, + ) + guardrail._ensure_access_token_async = AsyncMock( + return_value=("test-token", "test-project") + ) + + request = httpx.Request("POST", "https://modelarmor.example.test/v1") + upstream = httpx.Response(403, content=marker.encode(), request=request) + original = httpx.HTTPStatusError("Forbidden", request=request, response=upstream) + masked = MaskedHTTPStatusError(original, message=marker, text=marker) + + with patch.object( + guardrail.async_handler, "post", AsyncMock(side_effect=masked) + ), patch.object(verbose_proxy_logger, "debug") as debug_log, patch.object( + verbose_proxy_logger, "error" + ) as error_log, pytest.raises(ModelArmorAPIError) as exc_info: + await guardrail.make_model_armor_request(content=marker) + + direct_log = f"{debug_log.call_args_list} {error_log.call_args_list}" + assert "403" in str(exc_info.value.detail) + if sanitize: + assert marker not in str(exc_info.value.detail) + assert marker not in direct_log + else: + assert marker in str(exc_info.value.detail) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sanitize", [True, False]) +async def test_model_armor_post_call_logging_redacts_scanned_content(sanitize: bool): + marker = "SYNTHETIC_POST_CALL_MARKER" + armor_response = { + "sanitizationResult": { + "filterMatchState": "NO_MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "deidentifyResult": { + "matchState": "MATCH_FOUND", + "data": {"text": marker}, + } + } + } + }, + } + } + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + mask_response_content=True, + sanitize_error_detail=sanitize, + ) + guardrail.make_model_armor_request = AsyncMock(return_value=armor_response) + guardrail.should_run_guardrail = Mock(return_value=True) + + mock_llm_response = litellm.ModelResponse() + mock_llm_response.choices = [ + litellm.Choices(message=litellm.Message(content="model output")) + ] + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "synthetic input"}], + "metadata": {}, + "litellm_logging_obj": MagicMock(), + } + + with patch( + "litellm.proxy.common_utils.callback_utils.add_guardrail_response_to_standard_logging_object" + ) as add_logging: + await guardrail.async_post_call_success_hook( + data=request_data, + user_api_key_dict=UserAPIKeyAuth(), + response=mock_llm_response, + ) + + logged = add_logging.call_args.kwargs["guardrail_response"] + assert logged["guardrail_status"] == "success" + logged_armor_response = logged["guardrail_response"]["model_armor_response"] + if sanitize: + assert marker not in str(logged_armor_response) + assert ( + logged_armor_response["sanitizationResult"]["filterResults"]["sdp"][ + "sdpFilterResult" + ]["deidentifyResult"]["matchState"] + == "MATCH_FOUND" + ) + else: + assert logged_armor_response == armor_response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sanitize", [True, False]) +async def test_model_armor_streaming_logging_redacts_scanned_content(sanitize: bool): + marker = "SYNTHETIC_STREAMING_MARKER" + armor_response = { + "sanitizationResult": { + "filterMatchState": "NO_MATCH_FOUND", + "sanitizedText": marker, + } + } + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + sanitize_error_detail=sanitize, + ) + guardrail.make_model_armor_request = AsyncMock(return_value=armor_response) + guardrail.should_run_guardrail = Mock(return_value=True) + + async def mock_stream(): + yield litellm.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content="streamed output") + ) + ] + ) + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "synthetic input"}], + "metadata": {}, + } + + async for _ in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=mock_stream(), + request_data=request_data, + ): + pass + + logged_response = request_data["metadata"]["_model_armor_response"] + if sanitize: + assert logged_response == { + "sanitizationResult": { + "filterMatchState": "NO_MATCH_FOUND", + "sanitizedText": "[REDACTED]", + } + } + assert marker not in str(logged_response) + else: + assert logged_response == armor_response + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sanitize", [True, False]) +async def test_model_armor_match_found_sanitizes_caller_and_logging(sanitize: bool): + marker = "SYNTHETIC_MATCH_FOUND_MARKER" + armor_response = { + "sanitizationResult": { + "filterResults": { + "sdp": { + "sdpFilterResult": { + "inspectResult": { + "matchState": "MATCH_FOUND", + "findings": [{"marker": marker}], + } + } + } + } + } + } + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + event_hook=[GuardrailEventHooks.pre_mcp_call], + sanitize_error_detail=sanitize, + ) + guardrail.make_model_armor_request = AsyncMock(return_value=armor_response) + guardrail.should_run_guardrail = Mock(return_value=True) + request_data = { + "messages": [{"role": "user", "content": "synthetic input"}], + "metadata": {}, + } + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type=litellm.types.utils.CallTypes.call_mcp_tool.value, + ) + + detail = exc_info.value.detail + logged_response = request_data["metadata"]["_model_armor_response"] + if sanitize: + assert detail == {"error": "Content blocked by Model Armor"} + assert logged_response == { + "sanitizationResult": { + "filterResults": { + "sdp": { + "sdpFilterResult": { + "inspectResult": { + "matchState": "MATCH_FOUND", + "findings": "[REDACTED]", + } + } + } + } + } + } + assert marker not in str(detail) + assert marker not in str(logged_response) + else: + assert detail["model_armor_response"] == armor_response + assert logged_response == armor_response + assert marker in str(detail) + assert marker in str(logged_response) + + +def test_model_armor_sanitize_error_detail_config_wiring(): + from litellm.proxy.guardrails.guardrail_hooks.model_armor import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + config = {"guardrail_name": "model-armor-test"} + params = { + "guardrail": "model_armor", + "mode": "pre_mcp_call", + "template_id": "test-template", + "project_id": "test-project", + } + opted_out = initialize_guardrail( + LitellmParams(**params, sanitize_error_detail=False), config + ) + explicit_null = initialize_guardrail( + LitellmParams(**params, sanitize_error_detail=None), config + ) + default = initialize_guardrail(LitellmParams(**params), config) + + assert opted_out.sanitize_error_detail is False + assert explicit_null.sanitize_error_detail is True + assert default.sanitize_error_detail is True def test_model_armor_ui_friendly_name(): @@ -1394,7 +1928,10 @@ async def test_model_armor_guardrail_status_intervened_vs_failed(): ) info = request_data["metadata"]["standard_logging_guardrail_information"] + assert info[0]["guardrail_name"] == guardrail.guardrail_name assert info[0]["guardrail_status"] == "guardrail_intervened" + assert "model_armor_response" not in info[0]["guardrail_response"] + assert "sanitizationResult" not in info[0]["guardrail_response"] # 2: if an API error - guardrail status should be guardrail_failed_to_respond" guardrail2 = ModelArmorGuardrail( From 212a9213c4997a4957dfb9337d3f7a94ca138fba Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 21 Jul 2026 10:28:49 -0700 Subject: [PATCH 13/84] refactor(ui): migrate agents table onto the shared DataTable (#34089) * refactor(ui): migrate agents table onto the shared DataTable Replace the hand-rolled tremor table inside AgentsPanel with the shared DataTable, splitting the surface into a data-owning panel, a thin AgentsTable consumer, and a getAgentsTableColumns definition composed from the shared cell library. Row delete moves from an inline icon button into the per-row overflow menu, and the health-check toggle moves into the table toolbar since it controls which rows the server returns. The loading skeleton is now initial-load-only, so refetches keep the current rows on screen. Drops the last @tremor/react import from AgentsPanel, so its grandfathered eslint suppressions are pruned from the baseline. * fix(ui): keep agents ordering and token changes correct in the migrated table Sorting by created_at went through a raw accessor, and TanStack places undefined ahead of real values, so an agent with no created_at jumped to the top of the newest-first list. The pre-migration sort coerced a missing date to epoch 0 and sorted it last; restore that by sorting on a derived timestamp. Reload the list when the access token changes rather than leaving the previous token's rows on screen: show the skeleton for the new token, drop the rows if that load fails, and ignore a superseded response so a slow earlier request cannot overwrite newer rows. Refetches triggered by delete or the health-check toggle still keep their rows. Tests also reset the networking mocks between cases so an unconsumed mockResolvedValueOnce queue cannot leak into the next test. --- ui/litellm-dashboard/eslint-suppressions.json | 8 - .../agents/_components/AgentsPanel.test.tsx | 221 +++++++++++++++--- .../agents/_components/AgentsPanel.tsx | 187 +++++---------- .../agents/_components/AgentsTable.test.tsx | 147 ++++++++++++ .../agents/_components/AgentsTable.tsx | 88 +++++++ .../agents/_components/AgentsTableColumns.tsx | 163 +++++++++++++ 6 files changed, 656 insertions(+), 158 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTableColumns.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 59b36255daa..d596897c4c9 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -12,14 +12,6 @@ "count": 1 } }, - "src/app/(dashboard)/agents/_components/AgentsPanel.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/app/(dashboard)/agents/_components/add_agent_form.tsx": { "no-nested-ternary": { "count": 3 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx index 48674f21883..441d300436a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.test.tsx @@ -1,12 +1,13 @@ import React from "react"; -import { render, screen, waitFor, act, fireEvent, within } from "@testing-library/react"; +import { act, render, screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import AgentsPanel from "./AgentsPanel"; import * as networking from "@/components/networking"; vi.mock("@/components/networking", () => ({ getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), - deleteAgentCall: vi.fn(), + deleteAgentCall: vi.fn().mockResolvedValue({}), })); vi.mock("./add_agent_form", () => ({ @@ -19,56 +20,54 @@ vi.mock("./agent_info", () => ({ describe("AgentsPanel", () => { beforeEach(() => { - vi.clearAllMocks(); + // mockReset (not mockClear) so an unconsumed *Once queue cannot leak into the next test + vi.mocked(networking.getAgentsList).mockReset().mockResolvedValue({ agents: [] }); + vi.mocked(networking.deleteAgentCall).mockReset().mockResolvedValue({}); }); - it("should render the Agents panel title", async () => { + it("should render the Agents panel title", () => { render(); expect(screen.getByText("Agents")).toBeInTheDocument(); }); - it("should show Add New Agent button for admin users", async () => { + it("should show Add New Agent button for admin users", () => { render(); - expect(screen.getByText("+ Add New Agent")).toBeInTheDocument(); + expect(screen.getByText("Add New Agent")).toBeInTheDocument(); }); - it("should show Add New Agent button for proxy_admin users", async () => { + it("should show Add New Agent button for proxy_admin users", () => { render(); - expect(screen.getByText("+ Add New Agent")).toBeInTheDocument(); + expect(screen.getByText("Add New Agent")).toBeInTheDocument(); }); - it("should not show Add New Agent button for internal_user role", async () => { + it("should not show Add New Agent button for internal_user role", () => { render(); - expect(screen.queryByText("+ Add New Agent")).not.toBeInTheDocument(); + expect(screen.queryByText("Add New Agent")).not.toBeInTheDocument(); }); - it("should not show Add New Agent button for internal_user_viewer role", async () => { + it("should not show Add New Agent button for internal_user_viewer role", () => { render(); - expect(screen.queryByText("+ Add New Agent")).not.toBeInTheDocument(); + expect(screen.queryByText("Add New Agent")).not.toBeInTheDocument(); }); - it("should show Actions column header for admin role", async () => { + it("should show the Actions column for admin role", async () => { render(); - await waitFor(() => { - expect(screen.getByRole("columnheader", { name: /actions/i })).toBeInTheDocument(); - }); + expect(await screen.findByRole("columnheader", { name: /actions/i })).toBeInTheDocument(); }); - it("should not show Actions column header for internal user role", async () => { + it("should not show the Actions column for internal user role", async () => { render(); await waitFor(() => { expect(screen.queryByRole("columnheader", { name: /actions/i })).not.toBeInTheDocument(); - // confirm table is rendered (not still loading) expect(screen.getByRole("table")).toBeInTheDocument(); }); }); - it("should render the Health Check toggle", async () => { - render(); + it("should render the Health Check toggle for admins and non-admins", () => { + const { unmount } = render(); expect(screen.getByText("Health Check")).toBeInTheDocument(); - }); + unmount(); - it("should render the Health Check toggle for non-admin users too", async () => { render(); expect(screen.getByText("Health Check")).toBeInTheDocument(); }); @@ -108,19 +107,187 @@ describe("AgentsPanel", () => { expect(within(keylessRow).getByText("Needs Setup")).toBeInTheDocument(); }); - it("should call getAgentsList with health_check=true when toggle is enabled", async () => { + it("should refetch with health_check=true when the toggle is enabled", async () => { + const user = userEvent.setup(); render(); await waitFor(() => { expect(networking.getAgentsList).toHaveBeenCalledWith("test-token", false); }); - const toggle = screen.getByRole("switch"); - await act(async () => { - fireEvent.click(toggle); - }); + await user.click(screen.getByRole("switch")); await waitFor(() => { expect(networking.getAgentsList).toHaveBeenCalledWith("test-token", true); }); }); + + it("should delete an agent through the ⋯ menu and confirm modal, then refetch", async () => { + const user = userEvent.setup(); + vi.mocked(networking.getAgentsList).mockResolvedValue({ + agents: [ + { + agent_id: "agent-9", + agent_name: "Doomed Agent", + litellm_params: { model: "gpt-4" }, + spend: 0, + keys: [], + }, + ], + }); + + render(); + + await user.click(await screen.findByTestId("agent-actions-agent-9")); + await user.click(await screen.findByTestId("agent-action-delete")); + + const modal = await screen.findByRole("dialog"); + await user.click(within(modal).getByRole("button", { name: /^delete$/i })); + + await waitFor(() => { + expect(networking.deleteAgentCall).toHaveBeenCalledWith("test-token", "agent-9"); + }); + // one initial load + one post-delete refetch + await waitFor(() => { + expect(vi.mocked(networking.getAgentsList).mock.calls.length).toBeGreaterThanOrEqual(2); + }); + }); + + it("should show a loading skeleton on initial load and clear it once agents arrive", async () => { + render(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); + await waitFor(() => { + expect(screen.queryByTestId("skeleton-row")).not.toBeInTheDocument(); + }); + }); + + it("should clear the loading state when there is no access token rather than skeleton forever", async () => { + render(); + await waitFor(() => { + expect(screen.queryByTestId("skeleton-row")).not.toBeInTheDocument(); + }); + expect(screen.getByText("No agents yet")).toBeInTheDocument(); + expect(networking.getAgentsList).not.toHaveBeenCalled(); + }); + + it("should not show rows fetched with a previous access token after the token changes", async () => { + const agentFor = (name: string) => ({ + agent_id: `id-${name}`, + agent_name: name, + litellm_params: { model: "gpt-4" }, + spend: 0, + keys: [], + }); + let resolveSecond: (value: { agents: ReturnType[] }) => void = () => {}; + vi.mocked(networking.getAgentsList) + .mockResolvedValueOnce({ agents: [agentFor("first-token-agent")] }) + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveSecond = resolve; + }), + ); + + const { rerender } = render(); + expect(await screen.findByText("first-token-agent")).toBeInTheDocument(); + + rerender(); + + // the previous token's rows must not linger while the new token loads + expect(screen.queryByText("first-token-agent")).not.toBeInTheDocument(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); + + await act(async () => { + resolveSecond({ agents: [agentFor("second-token-agent")] }); + }); + expect(await screen.findByText("second-token-agent")).toBeInTheDocument(); + }); + + it("should drop previous rows when the fetch for a new token fails", async () => { + vi.mocked(networking.getAgentsList) + .mockResolvedValueOnce({ + agents: [ + { agent_id: "stale", agent_name: "Stale Agent", litellm_params: { model: "gpt-4" }, spend: 0, keys: [] }, + ], + }) + .mockRejectedValueOnce(new Error("unauthorized")); + + const { rerender } = render(); + expect(await screen.findByText("Stale Agent")).toBeInTheDocument(); + + rerender(); + + await waitFor(() => { + expect(screen.getByText("No agents yet")).toBeInTheDocument(); + }); + expect(screen.queryByText("Stale Agent")).not.toBeInTheDocument(); + }); + + it("should ignore a superseded response so it cannot overwrite the current token's rows", async () => { + let resolveFirst: (value: { + agents: { agent_id: string; agent_name: string; litellm_params: { model: string }; spend: number; keys: [] }[]; + }) => void = () => {}; + vi.mocked(networking.getAgentsList) + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveFirst = resolve; + }), + ) + .mockResolvedValueOnce({ + agents: [ + { agent_id: "current", agent_name: "Current Agent", litellm_params: { model: "gpt-4" }, spend: 0, keys: [] }, + ], + }); + + const { rerender } = render(); + rerender(); + + expect(await screen.findByText("Current Agent")).toBeInTheDocument(); + + // the slow token-a response lands last and must be discarded + await act(async () => { + resolveFirst({ + agents: [ + { agent_id: "stale", agent_name: "Superseded Agent", litellm_params: { model: "gpt-4" }, spend: 0, keys: [] }, + ], + }); + }); + + expect(screen.queryByText("Superseded Agent")).not.toBeInTheDocument(); + expect(screen.getByText("Current Agent")).toBeInTheDocument(); + }); + + it("should keep rows visible during a health-check refetch instead of re-showing the skeleton", async () => { + const user = userEvent.setup(); + const agents = [ + { + agent_id: "agent-1", + agent_name: "Stable Agent", + litellm_params: { model: "gpt-4" }, + spend: 0, + keys: [], + }, + ]; + let resolveRefetch: (value: { agents: typeof agents }) => void = () => {}; + vi.mocked(networking.getAgentsList) + .mockResolvedValueOnce({ agents }) + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveRefetch = resolve; + }), + ); + + render(); + expect(await screen.findByText("Stable Agent")).toBeInTheDocument(); + + await user.click(screen.getByRole("switch")); + + expect(screen.getByText("Stable Agent")).toBeInTheDocument(); + expect(screen.queryByTestId("skeleton-row")).not.toBeInTheDocument(); + + await act(async () => { + resolveRefetch({ agents }); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx index 84634620426..a4a71530c84 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsPanel.tsx @@ -1,27 +1,15 @@ import React, { useState, useEffect } from "react"; -import { - Button, - Card, - Table, - TableBody, - TableCell, - TableHead, - TableHeaderCell, - TableRow, - Badge, - Text, -} from "@tremor/react"; -import { Modal, Alert, Tooltip, Skeleton, Switch } from "antd"; -import { CheckCircleOutlined } from "@ant-design/icons"; +import { Modal, Alert } from "antd"; +import { Plus } from "lucide-react"; import { getAgentsList, deleteAgentCall } from "@/components/networking"; import AddAgentForm from "./add_agent_form"; import { isAdminRole } from "@/utils/roles"; import AgentInfoView from "./agent_info"; +import AgentsTable from "./AgentsTable"; import NotificationsManager from "@/components/molecules/notifications_manager"; import { Agent } from "@/components/agents/types"; import { Team } from "@/components/key_team_helpers/key_list"; -import { DateCell, IdCell, MoneyCell, StatusBadge } from "@/components/shared/table_cells"; -import TableIconActionButton from "@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; +import { Button } from "@/components/ui/button"; interface AgentsPanelProps { accessToken: string | null; @@ -36,37 +24,66 @@ interface AgentsResponse { const AgentsPanel: React.FC = ({ accessToken, userRole, teams }) => { const [agentsList, setAgentsList] = useState([]); const [isAddModalVisible, setIsAddModalVisible] = useState(false); - const [isLoading, setIsLoading] = useState(false); + const [isLoading, setIsLoading] = useState(true); const [isDeleting, setIsDeleting] = useState(false); + const [isHealthCheckLoading, setIsHealthCheckLoading] = useState(false); const [agentToDelete, setAgentToDelete] = useState<{ id: string; name: string } | null>(null); const [selectedAgentId, setSelectedAgentId] = useState(null); const [healthCheckEnabled, setHealthCheckEnabled] = useState(false); const isAdmin = userRole ? isAdminRole(userRole) : false; - const fetchAgents = async (healthCheck?: boolean) => { + useEffect(() => { + let cancelled = false; + const loadForToken = async () => { + if (!accessToken) { + setAgentsList([]); + setIsLoading(false); + return; + } + setIsLoading(true); + try { + const response: AgentsResponse = await getAgentsList(accessToken, false); + if (!cancelled) { + setAgentsList(response.agents || []); + } + } catch (error) { + console.error("Error fetching agents:", error); + if (!cancelled) { + setAgentsList([]); + } + } finally { + if (!cancelled) { + setIsLoading(false); + } + } + }; + loadForToken(); + return () => { + cancelled = true; + }; + }, [accessToken]); + + const refetchAgents = async (healthCheck: boolean) => { if (!accessToken) { return; } - - setIsLoading(true); try { - const response: AgentsResponse = await getAgentsList(accessToken, healthCheck ?? healthCheckEnabled); + const response: AgentsResponse = await getAgentsList(accessToken, healthCheck); setAgentsList(response.agents || []); } catch (error) { console.error("Error fetching agents:", error); - } finally { - setIsLoading(false); } }; - useEffect(() => { - fetchAgents(); - }, [accessToken]); - - const handleHealthCheckToggle = (checked: boolean) => { + const handleHealthCheckToggle = async (checked: boolean) => { setHealthCheckEnabled(checked); - fetchAgents(checked); + setIsHealthCheckLoading(true); + try { + await refetchAgents(checked); + } finally { + setIsHealthCheckLoading(false); + } }; const handleAddAgent = () => { @@ -81,7 +98,7 @@ const AgentsPanel: React.FC = ({ accessToken, userRole, teams }; const handleSuccess = () => { - fetchAgents(); + refetchAgents(healthCheckEnabled); }; const handleDeleteClick = (agentId: string, agentName: string) => { @@ -95,7 +112,7 @@ const AgentsPanel: React.FC = ({ accessToken, userRole, teams try { await deleteAgentCall(accessToken, agentToDelete.id); NotificationsManager.success(`Agent "${agentToDelete.name}" deleted successfully`); - fetchAgents(); + await refetchAgents(healthCheckEnabled); } catch (error) { console.error("Error deleting agent:", error); NotificationsManager.fromBackend("Failed to delete agent"); @@ -109,14 +126,6 @@ const AgentsPanel: React.FC = ({ accessToken, userRole, teams setAgentToDelete(null); }; - const sortedAgents = [...agentsList].sort((a, b) => { - const dateA = a.created_at ? new Date(a.created_at).getTime() : 0; - const dateB = b.created_at ? new Date(b.created_at).getTime() : 0; - return dateB - dateA; - }); - - const columnCount = isAdmin ? 7 : 6; - return (
@@ -132,25 +141,14 @@ const AgentsPanel: React.FC = ({ accessToken, userRole, teams showIcon className="mb-3" /> -
- {isAdmin && ( + {isAdmin && ( +
- )} - -
- - Health Check - -
-
-
+
+ )}
{selectedAgentId ? ( @@ -161,73 +159,16 @@ const AgentsPanel: React.FC = ({ accessToken, userRole, teams isAdmin={isAdmin} /> ) : ( - - {isLoading ? ( - - ) : ( -
- - - Agent Name - Agent ID - Spend (USD) - Model - Created - Status - {isAdmin && Actions} - - - - {sortedAgents.length === 0 ? ( - - - - No agents found. Click "+ Add New Agent" to create one. - - - - ) : ( - sortedAgents.map((agent) => ( - - - {agent.agent_name} - - - setSelectedAgentId(id)} /> - - - - - - - {agent.litellm_params?.model || "N/A"} - - - - - - - {(agent.keys?.length ?? 0) > 0 ? ( - - ) : ( - - )} - - {isAdmin && ( - - handleDeleteClick(agent.agent_id, agent.agent_name)} - /> - - )} - - )) - )} - -
- )} -
+ setSelectedAgentId(id)} + onDeleteClick={handleDeleteClick} + /> )} = {}): Agent => ({ + agent_id: "agent-1", + agent_name: "Test Agent", + litellm_params: { model: "gpt-4" }, + spend: 0, + keys: [{ token: "hash-1", key_alias: "primary", key_name: "sk-...1" }], + created_at: "2023-01-01T00:00:00Z", + ...overrides, +}); + +describe("AgentsTable", () => { + it("renders every column header", () => { + render(); + for (const header of ["Agent Name", "Agent ID", "Spend (USD)", "Model", "Created", "Status"]) { + expect(screen.getByText(header)).toBeInTheDocument(); + } + }); + + it("renders the agent's model and opens the detail view when the ID cell is clicked", async () => { + const user = userEvent.setup(); + const onAgentClick = vi.fn(); + const agent = makeAgent({ agent_id: "agent-xyz", agent_name: "Router", litellm_params: { model: "claude-3-5" } }); + render(); + + expect(screen.getByText("claude-3-5")).toBeInTheDocument(); + + await user.click(screen.getByText("agent-xyz")); + expect(onAgentClick).toHaveBeenCalledWith("agent-xyz"); + }); + + it("marks agents Active when they have keys and Needs Setup when they have none", () => { + render( + , + ); + + const keyedRow = screen.getByText("Keyed Agent").closest("tr")!; + const keylessRow = screen.getByText("Keyless Agent").closest("tr")!; + expect(within(keyedRow).getByText("Active")).toBeInTheDocument(); + expect(within(keylessRow).getByText("Needs Setup")).toBeInTheDocument(); + }); + + it("deletes an agent through the ⋯ actions menu", async () => { + const user = userEvent.setup(); + const onDeleteClick = vi.fn(); + const agent = makeAgent({ agent_id: "agent-9", agent_name: "Doomed Agent" }); + render(); + + await user.click(screen.getByTestId("agent-actions-agent-9")); + await user.click(await screen.findByTestId("agent-action-delete")); + + expect(onDeleteClick).toHaveBeenCalledWith("agent-9", "Doomed Agent"); + }); + + it("hides the actions column entirely for non-admins", () => { + const agent = makeAgent({ agent_id: "agent-2" }); + render(); + + expect(screen.queryByTestId("agent-actions-agent-2")).not.toBeInTheDocument(); + expect(screen.queryByRole("columnheader", { name: /actions/i })).not.toBeInTheDocument(); + expect(screen.getByRole("table")).toBeInTheDocument(); + }); + + it("shows the actions column for admins", () => { + render(); + expect(screen.getByRole("columnheader", { name: /actions/i })).toBeInTheDocument(); + expect(screen.getByTestId("agent-actions-agent-3")).toBeInTheDocument(); + }); + + it("defaults to sorting by created_at descending (newest first)", () => { + render( + , + ); + + const bodyRows = screen.getAllByRole("row").slice(1); + expect(bodyRows[0].textContent).toContain("Beta Agent"); + expect(bodyRows[1].textContent).toContain("Alpha Agent"); + }); + + it("sorts agents with no created_at last, never ahead of dated ones", () => { + render( + , + ); + + const bodyRows = screen.getAllByRole("row").slice(1); + expect(bodyRows[0].textContent).toContain("Beta Agent"); + expect(bodyRows[1].textContent).toContain("Alpha Agent"); + expect(bodyRows[2].textContent).toContain("Undated Agent"); + }); + + it("shows a rich empty state when there are no agents", () => { + render(); + expect(screen.getByText("No agents yet")).toBeInTheDocument(); + expect(screen.queryByTestId("skeleton-row")).not.toBeInTheDocument(); + }); + + it("renders loading skeleton rows on initial load instead of the empty state", () => { + render(); + expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); + expect(screen.queryByText("No agents yet")).not.toBeInTheDocument(); + }); + + it("invokes the health-check toggle from the toolbar", async () => { + const user = userEvent.setup(); + const onHealthCheckToggle = vi.fn(); + render(); + + expect(screen.getByText("Health Check")).toBeInTheDocument(); + await user.click(screen.getByRole("switch")); + expect(onHealthCheckToggle).toHaveBeenCalledWith(true, expect.anything()); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx new file mode 100644 index 00000000000..824ae47f3e6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx @@ -0,0 +1,88 @@ +"use client"; + +import { SortingState } from "@tanstack/react-table"; +import { Tooltip, Switch } from "antd"; +import { CheckCircleOutlined } from "@ant-design/icons"; +import { Bot } from "lucide-react"; +import React, { useMemo, useState } from "react"; + +import { Agent } from "@/components/agents/types"; +import { DataTable } from "@/components/shared/DataTable"; + +import { getAgentsTableColumns } from "./AgentsTableColumns"; + +interface AgentsTableProps { + agents: Agent[]; + isLoading: boolean; + isAdmin: boolean; + healthCheckEnabled: boolean; + isHealthCheckLoading: boolean; + onHealthCheckToggle: (checked: boolean) => void; + onAgentClick: (agentId: string) => void; + onDeleteClick: (agentId: string, agentName: string) => void; +} + +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +function EmptyState() { + return ( +
+
+ +
+
No agents yet
+
Add an agent to make it available in your organization.
+
+ ); +} + +const AgentsTable: React.FC = ({ + agents, + isLoading, + isAdmin, + healthCheckEnabled, + isHealthCheckLoading, + onHealthCheckToggle, + onAgentClick, + onDeleteClick, +}) => { + const [sorting, setSorting] = useState(DEFAULT_SORTING); + + const columns = useMemo( + () => getAgentsTableColumns({ isAdmin, onAgentClick, onDeleteClick }), + [isAdmin, onAgentClick, onDeleteClick], + ); + + return ( + agent.agent_id || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + isLoading={isLoading} + loadingMessage="Loading agents…" + noDataMessage={} + size="compact" + toolbar={() => ( +
+ +
+ + Health Check + +
+
+
+ )} + /> + ); +}; + +export default AgentsTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTableColumns.tsx new file mode 100644 index 00000000000..a8fe3973a42 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTableColumns.tsx @@ -0,0 +1,163 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { MoreHorizontal, Trash2 } from "lucide-react"; + +import { Agent } from "@/components/agents/types"; +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdentityCell, MoneyCell, StatusBadge } from "@/components/shared/table_cells"; +import { Badge } from "@/components/ui/badge"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; + +interface AgentRowActionsProps { + agent: Agent; + onDeleteClick: (agentId: string, agentName: string) => void; +} + +function AgentRowActions({ agent, onDeleteClick }: AgentRowActionsProps) { + return ( + + + + + + onDeleteClick(agent.agent_id, agent.agent_name)} + > + + Delete + + + + ); +} + +interface AgentsTableColumnsDeps { + isAdmin: boolean; + onAgentClick: (agentId: string) => void; + onDeleteClick: (agentId: string, agentName: string) => void; +} + +export const getAgentsTableColumns = ({ + isAdmin, + onAgentClick, + onDeleteClick, +}: AgentsTableColumnsDeps): ColumnDef[] => [ + { + id: "agent_name", + accessorKey: "agent_name", + meta: { title: "Agent Name" }, + header: ({ column }) => , + size: 200, + enableSorting: true, + cell: ({ row }) => { + const name = row.original.agent_name; + return ( + + {name || "-"} + + ); + }, + }, + { + id: "agent_id", + accessorKey: "agent_id", + meta: { title: "Agent ID" }, + header: ({ column }) => , + size: 200, + enableSorting: true, + cell: ({ row }) => ( + onAgentClick(row.original.agent_id)} + /> + ), + }, + { + id: "spend", + accessorKey: "spend", + meta: { title: "Spend (USD)" }, + header: ({ column }) => , + size: 130, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "model", + meta: { title: "Model" }, + header: "Model", + size: 170, + enableSorting: false, + cell: ({ row }) => { + const model = row.original.litellm_params?.model; + if (!model) { + return N/A; + } + return ( + + + {model} + + + ); + }, + }, + { + id: "created_at", + accessorFn: (agent) => { + const timestamp = agent.created_at ? new Date(agent.created_at).getTime() : 0; + return Number.isNaN(timestamp) ? 0 : timestamp; + }, + meta: { title: "Created" }, + header: ({ column }) => , + size: 150, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "status", + meta: { title: "Status" }, + header: "Status", + size: 130, + enableSorting: false, + cell: ({ row }) => { + const hasKeys = (row.original.keys?.length ?? 0) > 0; + return hasKeys ? ( + + ) : ( + + ); + }, + }, + ...(isAdmin + ? [ + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + } satisfies ColumnDef, + ] + : []), +]; From 01d624e860a936f2664eb958f7942a9ebccb380e Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 21 Jul 2026 12:01:17 -0700 Subject: [PATCH 14/84] 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 15/84] 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 16/84] 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 17/84] 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 18/84] 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 1315ebd1f97c2c3bb56f278a45d1904d48657401 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 21 Jul 2026 20:14:21 +0000 Subject: [PATCH 19/84] test(e2e): guard 0.3.0-style semver protocolVersion registration Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/a2a/test_a2a_agent_e2e.py | 10 ++++++++++ tests/e2e/coverage_registry/other.yaml | 1 + 2 files changed, 11 insertions(+) diff --git a/tests/e2e/a2a/test_a2a_agent_e2e.py b/tests/e2e/a2a/test_a2a_agent_e2e.py index eb61ace238c..823f6b9c001 100644 --- a/tests/e2e/a2a/test_a2a_agent_e2e.py +++ b/tests/e2e/a2a/test_a2a_agent_e2e.py @@ -69,6 +69,16 @@ class TestA2AAgentLifecycle: assert fetched.agent_name == agent.agent_name assert fetched.agent_card_params.protocol_version == "0.3" + @pytest.mark.covers("other.a2a.register.semver_version_accepted") + def test_semver_protocol_version_registers_and_serves(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: + agent = _register(client, resources, "0.3.0") + assert agent.agent_card_params.protocol_version.startswith("0.3") + card = unwrap(client.agent_card(agent.agent_id, scoped_key)) + assert card.protocol_version.startswith("0.3") + result = unwrap(client.send_message(agent.agent_id, scoped_key, _ask("Say hi in one word"))).result + assert result is not None + assert result.text != "" + @pytest.mark.covers("other.a2a.discovery.proxy_fronted_card") def test_discovery_card_is_proxy_fronted(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None: agent = _register(client, resources, "0.3") diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index f4d0120e085..63a626caab4 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -30,6 +30,7 @@ - {id: other.key_mgmt.spend_reset.resets_to_value, module: other, tier: P1, area: auth, assertions: [resets_to_value], source: "key_management_endpoints.py:4841", rationale: "reset_spend resets accumulated spend"} - {id: other.a2a.register.persists, module: other, tier: P1, area: a2a, assertions: [persists], source: "agent_endpoints/endpoints.py:325-443", rationale: "POST /v1/agents registers an agent card; GET /v1/agents/{id} reads it back"} - {id: other.a2a.register.unsupported_version_rejected, module: other, tier: P1, area: a2a, assertions: [unsupported_version_rejected], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a protocolVersion outside SUPPORTED_A2A_PROTOCOL_VERSIONS is refused with 400"} +- {id: other.a2a.register.semver_version_accepted, module: other, tier: P1, area: a2a, assertions: [semver_version_accepted], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a patch-level semver like 0.3.0 (what the Google A2A SDK emits) registers and serves as the 0.3 family rather than 400ing; regression guard for the v1.92 report"} - {id: other.a2a.discovery.proxy_fronted_card, module: other, tier: P1, area: a2a, assertions: [proxy_fronted_card], source: "agent_endpoints/a2a_endpoints.py get_agent_card", rationale: "/.well-known/agent-card.json serves the proxy url + supportedInterfaces and the LiteLLM virtual-key bearer scheme, not the upstream"} - {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} - {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} 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 20/84] 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 21/84] 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 22/84] 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 23/84] 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 && ( - - )} - + + + row.id} + rowSelection={rowSelection} + onRowSelectionChange={setRowSelection} + /> + + ); +} + +describe("DataTable row selection", () => { + it("supports uncontrolled per-row toggle, select-all, and indeterminate", async () => { + const user = userEvent.setup(); + + render( + row.id} + toolbar={(table) => {table.getSelectedRowModel().rows.length}} + />, + ); + + expect(selectedCount()).toHaveTextContent("0"); + + await user.click(rowBox("m1")); + expect(selectedCount()).toHaveTextContent("1"); + expect(selectAll()).toHaveAttribute("aria-checked", "mixed"); + + await user.click(selectAll()); + expect(selectedCount()).toHaveTextContent("3"); + expect(selectAll()).toHaveAttribute("aria-checked", "true"); + + await user.click(selectAll()); + expect(selectedCount()).toHaveTextContent("0"); + }); + + it("keys controlled selection by getRowId so the parent can map back to entities", async () => { + const user = userEvent.setup(); + render(); + + await user.click(rowBox("m2")); + expect(screen.getByTestId("keys")).toHaveTextContent("m2"); + + await user.click(rowBox("m3")); + expect(screen.getByTestId("keys")).toHaveTextContent("m2,m3"); + }); + + it("lets the parent clear the selection, the pattern an external pager needs", async () => { + const user = userEvent.setup(); + render(); + + await user.click(selectAll()); + expect(screen.getByTestId("keys")).toHaveTextContent("m1,m2,m3"); + + await user.click(screen.getByTestId("clear")); + expect(screen.getByTestId("keys")).toBeEmptyDOMElement(); + expect(rowBox("m1")).toHaveAttribute("aria-checked", "false"); + }); + + it("respects an enableRowSelection predicate", async () => { + const user = userEvent.setup(); + + render( + row.id} + enableRowSelection={(row) => row.original.id !== "m2"} + toolbar={(table) => {table.getSelectedRowModel().rows.length}} + />, + ); + + expect(rowBox("m2")).toHaveAttribute("aria-disabled", "true"); + + await user.click(rowBox("m2")); + expect(selectedCount()).toHaveTextContent("0"); + + await user.click(rowBox("m1")); + expect(selectedCount()).toHaveTextContent("1"); + }); + + it("rejects controlled rowSelection without onRowSelectionChange", () => { + const errors = validateDataTableConfig({ data, columns, rowSelection: { m1: true } }); + + expect(errors).toContain( + "Controlled `rowSelection` requires `onRowSelectionChange`; without it selection changes are dropped.", + ); + }); + + it("does not complain when selection is left uncontrolled", () => { + expect(validateDataTableConfig({ data, columns })).toHaveLength(0); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSelectionColumn.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSelectionColumn.tsx new file mode 100644 index 00000000000..da32a01ab0e --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTableSelectionColumn.tsx @@ -0,0 +1,53 @@ +"use client"; + +import type { ColumnDef, Row, RowData, Table } from "@tanstack/react-table"; + +import { Checkbox } from "@/components/ui/checkbox"; + +interface SelectionColumnOptions { + rowAriaLabel?: (row: Row) => string; +} + +function SelectAllCheckbox({ table }: { table: Table }) { + const allSelected = table.getIsAllPageRowsSelected(); + const someSelected = table.getIsSomePageRowsSelected(); + + return ( + table.toggleAllPageRowsSelected(Boolean(checked))} + /> + ); +} + +function SelectRowCheckbox({ row, label }: { row: Row; label: string }) { + return ( + row.toggleSelected(Boolean(checked))} + /> + ); +} + +export function createSelectionColumn( + options: SelectionColumnOptions = {}, +): ColumnDef { + const { rowAriaLabel } = options; + + return { + id: "select", + size: 44, + enableSorting: false, + enableHiding: false, + enableResizing: false, + meta: { title: "Select", className: "w-11", headerClassName: "w-11" }, + header: ({ table }) => , + cell: ({ row }) => , + }; +} diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/index.ts b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts index 1ee1eed1258..62ddd1b0742 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/index.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts @@ -3,6 +3,7 @@ import "./columnMeta"; export { DataTable, DataTableConfigError, validateDataTableConfig } from "./DataTable"; export { DataTableFilterDrawer, DataTableFilterField, type FilterDraft } from "./DataTableFilterDrawer"; export { DataTablePagination, DEFAULT_PAGE_SIZE_OPTIONS } from "./DataTablePagination"; +export { createSelectionColumn } from "./DataTableSelectionColumn"; export { DataTableToolbar } from "./DataTableToolbar"; export { DataTableViewOptions } from "./DataTableViewOptions"; export { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts index 672ab512ef4..40f3a4df204 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts @@ -6,6 +6,7 @@ import type { PaginationState, Row, RowData, + RowSelectionState, SortingState, Table, VisibilityState, @@ -59,6 +60,10 @@ export interface DataTableProps { expanded?: ExpandedState; onExpandedChange?: OnChangeFn; + enableRowSelection?: boolean | ((row: Row) => boolean); + rowSelection?: RowSelectionState; + onRowSelectionChange?: OnChangeFn; + onRowClick?: (row: TData) => void; rowClassName?: (row: Row) => string; diff --git a/ui/litellm-dashboard/src/components/ui/checkbox.tsx b/ui/litellm-dashboard/src/components/ui/checkbox.tsx new file mode 100644 index 00000000000..93f419e79a7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ui/checkbox.tsx @@ -0,0 +1,28 @@ +"use client"; + +import { Checkbox as CheckboxPrimitive } from "@base-ui/react/checkbox"; + +import { cn } from "@/lib/cva.config"; +import { CheckIcon } from "lucide-react"; + +function Checkbox({ className, ...props }: CheckboxPrimitive.Root.Props) { + return ( + + + + + + ); +} + +export { Checkbox }; From fa025fc4748f94af928cf347a69392d7cdeafc1c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:09:15 -0700 Subject: [PATCH 48/84] chore(tests): replace a customer name and domain with neutral placeholders --- tests/e2e/coverage_registry/llm_conversational.yaml | 8 ++++---- ...ssages_mid_conversation_system_native_providers_e2e.py | 2 +- .../test_azure_anthropic_messages_transformation.py | 2 +- ..._vertex_ai_partner_models_anthropic_messages_config.py | 2 +- .../management_endpoints/scim/test_scim_v2_endpoints.py | 6 +++--- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 3b1aff80024..26280d35da0 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -46,10 +46,10 @@ - {id: llm.messages.anthropic.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Extended thinking via Messages API"} - {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Flagged Claude 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (#32578/#32831/#32882)", fail_before_fix: proven} - {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Claude <= 4.7 rejects role system inside messages; unflagged models must hoist reminders into top-level system or every Claude Code session 400s (#32831)", fail_before_fix: proven} -- {id: llm.messages.azure_foundry.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: azure_foundry, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/azure_ai/anthropic/messages_transformation.py", rationale: "Azure Foundry serves Claude on the native Anthropic contract, so flagged 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (Kraken Tech RCA gap)", fail_before_fix: proven} -- {id: llm.messages.azure_foundry.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: azure_foundry, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/azure_ai/anthropic/messages_transformation.py", rationale: "Azure Foundry Claude <= 4.7 rejects role system inside messages; unflagged models must hoist reminders into top-level system or every Claude Code session 400s (Kraken Tech RCA gap)", fail_before_fix: proven} -- {id: llm.messages.vertex.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: vertex, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py", rationale: "Vertex serves Claude on the native Anthropic contract, so flagged 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (Kraken Tech RCA gap)", fail_before_fix: proven} -- {id: llm.messages.vertex.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: vertex, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py", rationale: "Vertex Claude <= 4.7 rejects role system inside messages; unflagged models must hoist reminders into top-level system or every Claude Code session 400s (Kraken Tech RCA gap)", fail_before_fix: proven} +- {id: llm.messages.azure_foundry.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: azure_foundry, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/azure_ai/anthropic/messages_transformation.py", rationale: "Azure Foundry serves Claude on the native Anthropic contract, so flagged 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (customer RCA gap)", fail_before_fix: proven} +- {id: llm.messages.azure_foundry.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: azure_foundry, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/azure_ai/anthropic/messages_transformation.py", rationale: "Azure Foundry Claude <= 4.7 rejects role system inside messages; unflagged models must hoist reminders into top-level system or every Claude Code session 400s (customer RCA gap)", fail_before_fix: proven} +- {id: llm.messages.vertex.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: vertex, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py", rationale: "Vertex serves Claude on the native Anthropic contract, so flagged 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (customer RCA gap)", fail_before_fix: proven} +- {id: llm.messages.vertex.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: vertex, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py", rationale: "Vertex Claude <= 4.7 rejects role system inside messages; unflagged models must hoist reminders into top-level system or every Claude Code session 400s (customer RCA gap)", fail_before_fix: proven} - {id: llm.responses.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Core endpoint; OpenAI Responses native"} - {id: llm.responses.openai.basic.stream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: stream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Streaming via /v1/responses"} - {id: llm.responses.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "response_api_endpoints/endpoints.py:26", rationale: "Cost logged on responses"} diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py index 6d6830075bc..97d24e0564b 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_native_providers_e2e.py @@ -7,7 +7,7 @@ accepted in place on Claude 4.8+/5 (200) but rejected on Claude 4.7 and older ("role 'system' is not supported on this model", 400), and a *leading* system entry is rejected on every model ("messages.0: use the top-level 'system' parameter"). This mirrors Bedrock Invoke (PRs #32578/#32831/#32882); the same -model-gated hoist now runs for these two providers (Kraken Tech RCA gap #3). +model-gated hoist now runs for these two providers (customer RCA gap #3). Flagged models (``supports_mid_conversation_system`` in the cost map: Claude 4.8+ and the 5 family) must keep the reminder in ``messages`` so the top-level diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 25d24cfc3ac..1e1b98861b4 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -412,7 +412,7 @@ class TestAzureAnthropicMidConversationSystem: older Claude, and a *leading* system entry 400s on every model ("messages.0: use the top-level 'system' parameter"). These tests pin the model-aware hoist the config applies so Claude Code sessions neither collapse the prompt cache - on 4.8+ nor hard-fail on 4.7 and older (RCA: Kraken Tech high-spend).""" + on 4.8+ nor hard-fail on 4.7 and older (RCA: customer high-spend).""" def test_supported_model_keeps_mid_conversation_system_in_place(self, local_model_cost_map): messages = [ diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 2d09cc0ed32..292bddf1274 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -591,7 +591,7 @@ class TestVertexAnthropicMidConversationSystem: Claude, and a *leading* system entry 400s on every model ("messages.0: use the top-level 'system' parameter"). These tests pin the model-aware hoist so Claude Code sessions neither collapse the prompt cache on 4.8+ nor hard-fail - on 4.7 and older (RCA: Kraken Tech high-spend).""" + on 4.7 and older (RCA: customer high-spend).""" def test_supported_model_keeps_mid_conversation_system_in_place(self, local_model_cost_map): messages = [ diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index f27f1197090..3ff8e2a6886 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -347,8 +347,8 @@ async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeyp # Step 3: Create a user via SCIM scim_user = SCIMUser( schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], - userName="idontexist@krakentest.tech", - emails=[SCIMUserEmail(value="idontexist@krakentest.tech")], + userName="idontexist@example.com", + emails=[SCIMUserEmail(value="idontexist@example.com")], ) mock_prisma_client = mocker.MagicMock() @@ -364,7 +364,7 @@ async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeyp new_user_mock = mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2.new_user", - AsyncMock(return_value=NewUserRequest(user_id="idontexist@krakentest.tech")), + AsyncMock(return_value=NewUserRequest(user_id="idontexist@example.com")), ) mocker.patch( From 58ff0e32ba750c0e4605682a9984bc4720a0a6b1 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 21 Jul 2026 15:13:07 -0700 Subject: [PATCH 49/84] chore(deps): bump gitpython to 3.1.52 in uv.lock (#34168) --- uv.lock | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/uv.lock b/uv.lock index cee24aca330..0a90682a187 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-07-18T19:44:23.519632Z" +exclude-newer = "2026-07-18T21:57:43.13625Z" exclude-newer-span = "P3D" [manifest] @@ -2314,14 +2314,14 @@ wheels = [ [[package]] name = "gitpython" -version = "3.1.50" +version = "3.1.52" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "gitdb" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/33/f6/354ae6491228b5eb40e10d89c4d13c651fe1cf7556e35ebdded50cff57ce/gitpython-3.1.50.tar.gz", hash = "sha256:80da2d12504d52e1f998772dc5baf6e553f8d2fcfe1fcc226c9d9a2ee3372dcc", size = 219798, upload-time = "2026-05-06T04:01:26.571Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e5/fd/df0bafa4eb5ea2f51e1adee9f7a94c8e62c5d180e65117045dfca3439c8a/gitpython-3.1.52.tar.gz", hash = "sha256:de0a8ad86274c6e75ae8b37dd055ba68f19818c813108642263227b20775b48e", size = 223726, upload-time = "2026-07-16T03:15:59.599Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/20/7a/1c6e3562dfd8950adbb11ffbc65d21e7c89d01a6e4f137fa981056de25c5/gitpython-3.1.50-py3-none-any.whl", hash = "sha256:d352abe2908d07355014abdd21ddf798c2a961469239afec4962e9da884858f9", size = 212507, upload-time = "2026-05-06T04:01:23.799Z" }, + { url = "https://files.pythonhosted.org/packages/8d/90/04dff7c1e176bb1c3011ef1647393d368790da710d8dde1cdcfad301f45a/gitpython-3.1.52-py3-none-any.whl", hash = "sha256:79a36ee1f83523214a3f72d56cf1c4e490d577dc61af77e43dfe5862bd9da01a", size = 215366, upload-time = "2026-07-16T03:15:58.239Z" }, ] [[package]] From 72d458e416d1bf1d25f343b3d12a5304bbe56120 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 21 Jul 2026 15:17:58 -0700 Subject: [PATCH 50/84] feat(ui): add react-hook-form + zod form infrastructure (#34170) * feat(ui): add react-hook-form + zod form infrastructure Introduce the shared form layer the dashboard's antd forms will migrate onto, with no user-visible change yet. - pin react-hook-form, @hookform/resolvers, and zod (kept on 3.25.76 and imported via the zod/v4 entrypoint so openai's optional zod ^3 peer still resolves and npm ci stays clean) - vendor the base-vega Field family into components/shared/form as forwardRef components on the repo's cva.config, since base-vega ships no form primitive and its field source imports class-variance-authority and is React 19 style - add a FormField bridge that binds a react-hook-form Controller to the Field layer and wires label, description, and error ids into aria attributes - add pickDirty, which narrows a submitted body to the top-level keys the user actually touched so a partial update stops re-sending untouched fields pickDirty reads dirtiness at the top level because react-hook-form tracks it per leaf, so an edited array arrives as [true, false] and a cleared list as an empty array that still carries its default-length dirty markers; the falsy clear tokens (null, [], {}, 0, false) all survive. Tests cover the Field primitives, the FormField aria wiring against a live zod resolver, and pickDirty both as a unit and driven through a real react-hook-form instance. * test(ui): lock pickDirty behavior on a pure field-array reorder react-hook-form compares each array element to its default positionally by value, so useFieldArray move/swap and a reordered scalar array all mark the moved indices dirty and pickDirty sends the whole array; a swap of two equal elements is a value-level no-op and is correctly omitted. Covers the reorder case a review flagged as untested. --- ui/litellm-dashboard/package-lock.json | 34 ++- ui/litellm-dashboard/package.json | 5 +- .../components/shared/form/FormField.test.tsx | 180 ++++++++++++ .../src/components/shared/form/FormField.tsx | 75 +++++ .../src/components/shared/form/field.test.tsx | 125 +++++++++ .../src/components/shared/form/field.tsx | 223 +++++++++++++++ .../src/lib/forms/pickDirty.test.ts | 263 ++++++++++++++++++ .../src/lib/forms/pickDirty.ts | 24 ++ 8 files changed, 926 insertions(+), 3 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/shared/form/FormField.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/form/FormField.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/form/field.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/form/field.tsx create mode 100644 ui/litellm-dashboard/src/lib/forms/pickDirty.test.ts create mode 100644 ui/litellm-dashboard/src/lib/forms/pickDirty.ts diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 7a65b63b33c..49d289879d3 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -14,6 +14,7 @@ "@base-ui/react": "^1.6.0", "@headlessui/tailwindcss": "0.2.2", "@heroicons/react": "1.0.6", + "@hookform/resolvers": "5.4.0", "@tanstack/react-pacer": "0.22.1", "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", @@ -34,13 +35,15 @@ "react": "18.3.1", "react-copy-to-clipboard": "5.1.1", "react-dom": "18.3.1", + "react-hook-form": "7.82.0", "react-json-view-lite": "2.5.0", "react-markdown": "9.1.0", "react-syntax-highlighter": "15.6.6", "recharts": "3.9.2", "remark-gfm": "4.0.1", "tailwind-merge": "3.4.0", - "uuid": "14.0.0" + "uuid": "14.0.0", + "zod": "3.25.76" }, "devDependencies": { "@eslint/js": "9.39.2", @@ -1556,6 +1559,18 @@ "react": ">= 16" } }, + "node_modules/@hookform/resolvers": { + "version": "5.4.0", + "resolved": "https://registry.npmjs.org/@hookform/resolvers/-/resolvers-5.4.0.tgz", + "integrity": "sha512-EIsqr/t/qbinPIhGjMdtvutIN1Kk4uwbROE9/UQ93CAVGR7GkA7Y92+fX80OzXi/OB67jVFYwKGO1WzkxmkFZw==", + "license": "MIT", + "dependencies": { + "@standard-schema/utils": "^0.3.0" + }, + "peerDependencies": { + "react-hook-form": "^7.55.0" + } + }, "node_modules/@humanfs/core": { "version": "0.19.2", "resolved": "https://registry.npmjs.org/@humanfs/core/-/core-0.19.2.tgz", @@ -11780,6 +11795,22 @@ "react": "^18.3.1" } }, + "node_modules/react-hook-form": { + "version": "7.82.0", + "resolved": "https://registry.npmjs.org/react-hook-form/-/react-hook-form-7.82.0.tgz", + "integrity": "sha512-Zw/uFZ2dO+02GHlBn7JFGn8kZJ7LdM33B/0BXOovzFay+CMhf94JMw5BVu+F1tVkUKjNvBuaE3fz5BJhga10Tg==", + "license": "MIT", + "engines": { + "node": ">=18.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/react-hook-form" + }, + "peerDependencies": { + "react": "^16.8.0 || ^17 || ^18 || ^19" + } + }, "node_modules/react-is": { "version": "17.0.2", "resolved": "https://registry.npmjs.org/react-is/-/react-is-17.0.2.tgz", @@ -14156,7 +14187,6 @@ "version": "3.25.76", "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz", "integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==", - "devOptional": true, "license": "MIT", "funding": { "url": "https://github.com/sponsors/colinhacks" diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index b5e93d175bf..c29b53cd818 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -30,6 +30,7 @@ "@base-ui/react": "^1.6.0", "@headlessui/tailwindcss": "0.2.2", "@heroicons/react": "1.0.6", + "@hookform/resolvers": "5.4.0", "@tanstack/react-pacer": "0.22.1", "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", @@ -50,13 +51,15 @@ "react": "18.3.1", "react-copy-to-clipboard": "5.1.1", "react-dom": "18.3.1", + "react-hook-form": "7.82.0", "react-json-view-lite": "2.5.0", "react-markdown": "9.1.0", "react-syntax-highlighter": "15.6.6", "recharts": "3.9.2", "remark-gfm": "4.0.1", "tailwind-merge": "3.4.0", - "uuid": "14.0.0" + "uuid": "14.0.0", + "zod": "3.25.76" }, "devDependencies": { "@eslint/js": "9.39.2", diff --git a/ui/litellm-dashboard/src/components/shared/form/FormField.test.tsx b/ui/litellm-dashboard/src/components/shared/form/FormField.test.tsx new file mode 100644 index 00000000000..af9122e2bd5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/form/FormField.test.tsx @@ -0,0 +1,180 @@ +import { zodResolver } from "@hookform/resolvers/zod"; +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import * as React from "react"; +import { useForm } from "react-hook-form"; +import { describe, expect, it, vi } from "vitest"; +import { z } from "zod/v4"; + +import { Input } from "@/components/ui/input"; + +import { FormField } from "./FormField"; + +const schema = z.object({ + team_alias: z.string().min(1, "Please input a team name"), + owner: z.string(), +}); + +type FormInput = z.input; + +const TestForm = ({ + onSubmit, + defaultValues = { team_alias: "team-a", owner: "" }, + description, +}: { + onSubmit: (values: z.output) => void; + defaultValues?: FormInput; + description?: React.ReactNode; +}) => { + const form = useForm>({ + resolver: zodResolver(schema), + defaultValues, + }); + + return ( +
+ + {(field) => } + + +
+ ); +}; + +describe("FormField", () => { + it("associates the label with the control so it is reachable by its accessible name", () => { + render(); + + expect(screen.getByLabelText("Team Name")).toHaveValue("team-a"); + }); + + it("gives each field instance a unique control id", () => { + const Harness = () => { + const form = useForm({ defaultValues: { team_alias: "", owner: "" } }); + return ( + <> + + {(field) => } + + + {(field) => } + + + ); + }; + render(); + + expect(screen.getByLabelText("One").id).not.toBe(screen.getByLabelText("Two").id); + }); + + it("feeds edits back into form state and submits the parsed output", async () => { + const user = userEvent.setup(); + const onSubmit = vi.fn(); + render(); + + await user.clear(screen.getByLabelText("Team Name")); + await user.type(screen.getByLabelText("Team Name"), "team-b"); + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1)); + expect(onSubmit.mock.calls[0][0]).toEqual({ team_alias: "team-b", owner: "" }); + }); + + it("renders the zod message and blocks submit when validation fails", async () => { + const user = userEvent.setup(); + const onSubmit = vi.fn(); + render(); + + await user.clear(screen.getByLabelText("Team Name")); + await user.click(screen.getByRole("button", { name: "Save" })); + + expect(await screen.findByRole("alert")).toHaveTextContent("Please input a team name"); + expect(onSubmit).not.toHaveBeenCalled(); + }); + + it("marks the control invalid and points aria-describedby at the message", async () => { + const user = userEvent.setup(); + render(); + + await user.clear(screen.getByLabelText("Team Name")); + await user.click(screen.getByRole("button", { name: "Save" })); + + const control = await screen.findByLabelText("Team Name"); + await waitFor(() => expect(control).toHaveAttribute("aria-invalid", "true")); + expect(control.getAttribute("aria-describedby")).toBe(screen.getByRole("alert").id); + }); + + it("leaves a valid control free of aria-invalid", () => { + render(); + + expect(screen.getByLabelText("Team Name")).not.toHaveAttribute("aria-invalid"); + }); + + it("clears the message once the value becomes valid again", async () => { + const user = userEvent.setup(); + render(); + + await user.clear(screen.getByLabelText("Team Name")); + await user.click(screen.getByRole("button", { name: "Save" })); + expect(await screen.findByRole("alert")).toBeInTheDocument(); + + await user.type(screen.getByLabelText("Team Name"), "team-c"); + + await waitFor(() => expect(screen.queryByRole("alert")).not.toBeInTheDocument()); + }); + + it("describes the control by its description when there is no error", () => { + render(); + + const control = screen.getByLabelText("Team Name"); + const describedBy = control.getAttribute("aria-describedby"); + + expect(describedBy).not.toBeNull(); + expect(document.getElementById(describedBy!)).toHaveTextContent("Shown to team members"); + }); + + it("describes the control by both description and error while invalid", async () => { + const user = userEvent.setup(); + render(); + + await user.clear(screen.getByLabelText("Team Name")); + await user.click(screen.getByRole("button", { name: "Save" })); + await screen.findByRole("alert"); + + const ids = screen.getByLabelText("Team Name").getAttribute("aria-describedby")?.split(" ") ?? []; + + expect(ids).toHaveLength(2); + expect(ids).toContain(screen.getByRole("alert").id); + }); + + it("omits aria-describedby entirely when there is no description and no error", () => { + render(); + + expect(screen.getByLabelText("Team Name")).not.toHaveAttribute("aria-describedby"); + }); + + it("hands the control a value and onChange so non-native widgets can be wired", async () => { + const user = userEvent.setup(); + const seen: unknown[] = []; + const Harness = () => { + const form = useForm({ defaultValues: { team_alias: "team-a", owner: "" } }); + return ( + + {(field) => { + seen.push(field.value); + return ( + + ); + }} + + ); + }; + render(); + + await user.click(screen.getByRole("button", { name: "widget" })); + + await waitFor(() => expect(seen.at(-1)).toBe("from-widget")); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/form/FormField.tsx b/ui/litellm-dashboard/src/components/shared/form/FormField.tsx new file mode 100644 index 00000000000..3b9783333cc --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/form/FormField.tsx @@ -0,0 +1,75 @@ +"use client"; + +import * as React from "react"; +import { + Controller, + type Control, + type ControllerRenderProps, + type FieldPath, + type FieldValues, +} from "react-hook-form"; + +import { Field, FieldDescription, FieldError, FieldLabel } from "./field"; + +export type FormFieldControlProps< + TFieldValues extends FieldValues, + TName extends FieldPath, +> = ControllerRenderProps & { + id: string; + "aria-invalid": true | undefined; + "aria-describedby": string | undefined; +}; + +export interface FormFieldProps> { + control: Control; + name: TName; + label?: React.ReactNode; + description?: React.ReactNode; + orientation?: "vertical" | "horizontal" | "responsive"; + className?: string; + children: (control: FormFieldControlProps) => React.ReactNode; +} + +export const FormField = >({ + control, + name, + label, + description, + orientation, + className, + children, +}: FormFieldProps) => { + const reactId = React.useId(); + const controlId = `${reactId}-control`; + const descriptionId = `${reactId}-description`; + const errorId = `${reactId}-error`; + + return ( + { + const invalid = fieldState.error !== undefined; + const describedBy = + [description !== undefined ? descriptionId : undefined, invalid ? errorId : undefined] + .filter((id): id is string => id !== undefined) + .join(" ") || undefined; + const controlProps: FormFieldControlProps = { + ...field, + id: controlId, + "aria-invalid": invalid || undefined, + "aria-describedby": describedBy, + }; + + return ( + + {label !== undefined && {label}} + {children(controlProps)} + {description !== undefined && {description}} + + + ); + }} + /> + ); +}; diff --git a/ui/litellm-dashboard/src/components/shared/form/field.test.tsx b/ui/litellm-dashboard/src/components/shared/form/field.test.tsx new file mode 100644 index 00000000000..54b589ce2f4 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/form/field.test.tsx @@ -0,0 +1,125 @@ +import { render, screen } from "@testing-library/react"; +import * as React from "react"; +import { describe, expect, it } from "vitest"; + +import { + Field, + FieldContent, + FieldDescription, + FieldError, + FieldGroup, + FieldLabel, + FieldLegend, + FieldSeparator, + FieldSet, + FieldTitle, +} from "./field"; + +describe("FieldError", () => { + it("renders nothing when there are no errors and no children", () => { + const { container } = render(); + + expect(container).toBeEmptyDOMElement(); + }); + + it("renders nothing when every error entry is undefined", () => { + const { container } = render(); + + expect(container).toBeEmptyDOMElement(); + }); + + it("renders a single message as plain text, not a list", () => { + render(); + + expect(screen.getByRole("alert")).toHaveTextContent("Required"); + expect(screen.queryByRole("listitem")).not.toBeInTheDocument(); + }); + + it("collapses duplicate messages to a single entry", () => { + render(); + + expect(screen.getByRole("alert")).toHaveTextContent("Required"); + expect(screen.queryByRole("listitem")).not.toBeInTheDocument(); + }); + + it("renders distinct messages as a list", () => { + render(); + + const items = screen.getAllByRole("listitem"); + expect(items.map((item) => item.textContent)).toEqual(["Too short", "Must be lowercase"]); + }); + + it("prefers explicit children over the errors prop", () => { + render(from children); + + expect(screen.getByRole("alert")).toHaveTextContent("from children"); + expect(screen.getByRole("alert")).not.toHaveTextContent("from errors"); + }); + + it("exposes the message to assistive tech via role=alert", () => { + render(); + + expect(screen.getByRole("alert")).toBeInTheDocument(); + }); +}); + +describe("Field", () => { + it("marks itself invalid so descendants can style off it", () => { + render( + + child + , + ); + + expect(screen.getByRole("group")).toHaveAttribute("data-invalid", "true"); + }); + + it("defaults to vertical orientation", () => { + render(); + + expect(screen.getByRole("group")).toHaveAttribute("data-orientation", "vertical"); + }); + + it("honours an explicit orientation", () => { + render(); + + expect(screen.getByRole("group")).toHaveAttribute("data-orientation", "horizontal"); + }); +}); + +describe("field primitives forward refs to their DOM node", () => { + it.each([ + ["Field", Field, HTMLDivElement], + ["FieldContent", FieldContent, HTMLDivElement], + ["FieldDescription", FieldDescription, HTMLParagraphElement], + ["FieldGroup", FieldGroup, HTMLDivElement], + ["FieldLabel", FieldLabel, HTMLLabelElement], + ["FieldSeparator", FieldSeparator, HTMLDivElement], + ["FieldTitle", FieldTitle, HTMLDivElement], + ])("%s", (_name, Component, expected) => { + const ref = React.createRef(); + render(React.createElement(Component as React.ElementType, { ref })); + + expect(ref.current).toBeInstanceOf(expected); + }); + + it("FieldSet and FieldLegend", () => { + const fieldSet = React.createRef(); + const legend = React.createRef(); + render( +
+ Legend +
, + ); + + expect(fieldSet.current).toBeInstanceOf(HTMLFieldSetElement); + expect(legend.current).toBeInstanceOf(HTMLLegendElement); + }); + + it("FieldError", () => { + const ref = React.createRef(); + render(); + + expect(ref.current).toBeInstanceOf(HTMLDivElement); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/form/field.tsx b/ui/litellm-dashboard/src/components/shared/form/field.tsx new file mode 100644 index 00000000000..36ce691827c --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/form/field.tsx @@ -0,0 +1,223 @@ +"use client"; + +import * as React from "react"; +import { type VariantProps } from "cva"; + +import { Label } from "@/components/ui/label"; +import { Separator } from "@/components/ui/separator"; +import { cn, cva } from "@/lib/cva.config"; + +const FieldSet = React.forwardRef>( + ({ className, ...props }, ref) => ( +
[data-slot=checkbox-group]]:gap-3 has-[>[data-slot=radio-group]]:gap-3", + className, + )} + {...props} + /> + ), +); +FieldSet.displayName = "FieldSet"; + +const FieldLegend = React.forwardRef< + HTMLLegendElement, + React.ComponentPropsWithoutRef<"legend"> & { variant?: "legend" | "label" } +>(({ className, variant = "legend", ...props }, ref) => ( + +)); +FieldLegend.displayName = "FieldLegend"; + +const FieldGroup = React.forwardRef>( + ({ className, ...props }, ref) => ( +
+ ), +); +FieldGroup.displayName = "FieldGroup"; + +const fieldVariants = cva({ + base: "group/field flex w-full gap-3 data-[invalid=true]:text-destructive", + variants: { + orientation: { + vertical: "flex-col *:w-full [&>.sr-only]:w-auto", + horizontal: + "flex-row items-center has-[>[data-slot=field-content]]:items-start *:data-[slot=field-label]:flex-auto has-[>[data-slot=field-content]]:[&>[role=checkbox],[role=radio]]:mt-px", + responsive: + "flex-col *:w-full @md/field-group:flex-row @md/field-group:items-center @md/field-group:*:w-auto @md/field-group:has-[>[data-slot=field-content]]:items-start @md/field-group:*:data-[slot=field-label]:flex-auto [&>.sr-only]:w-auto @md/field-group:has-[>[data-slot=field-content]]:[&>[role=checkbox],[role=radio]]:mt-px", + }, + }, + defaultVariants: { + orientation: "vertical", + }, +}); + +const Field = React.forwardRef< + HTMLDivElement, + React.ComponentPropsWithoutRef<"div"> & VariantProps +>(({ className, orientation = "vertical", ...props }, ref) => ( +
+)); +Field.displayName = "Field"; + +const FieldContent = React.forwardRef>( + ({ className, ...props }, ref) => ( +
+ ), +); +FieldContent.displayName = "FieldContent"; + +const FieldLabel = React.forwardRef>( + ({ className, ...props }, ref) => ( +