scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 134

Merge remote-tracking branch 'origin/main' into litellm_fix_interrupted_anthropic_reasoning_usage

Browse source
This commit is contained in:
kerry 2026-09-16 23:13:06 +00:00
commit 43514a7ffe
66 changed files with 3908 additions and 1119 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
ARCHITECTURE.md
View file

@ -45,7 +45,7 @@ sequenceDiagram
ProxyServer->>Auth: user_api_key_auth()
Auth->>Redis: Check API key cache
Redis-->>Auth: Key info + spend limits
ProxyServer->>Hooks: max_budget_limiter, parallel_request_limiter
ProxyServer->>Hooks: parallel_request_limiter, cache_control_check
Hooks->>Redis: Check/increment rate limit counters
ProxyServer->>Router: route_request()
Router->>Main: litellm.acompletion()
@ -145,7 +145,6 @@ graph TD
| Hook | File | Purpose |
|------|------|---------|
| `max_budget_limiter` | `proxy/hooks/max_budget_limiter.py` | Enforce budget limits |
| `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user |
| `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation |
| `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation |

11
litellm/integrations/opentelemetry.py
View file

@ -2965,16 +2965,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
)
propagator: Final = TraceContextTextMapPropagator()
carrier: Final = {"traceparent": _traceparent}
carrier: Final = {key: headers[key] for key in ("traceparent", "tracestate") if headers.get(key) is not None}
_parent_context: Final = propagator.extract(carrier=carrier)
return _parent_context
def _get_span_context(self, kwargs, default_span: Span | None = None):
from opentelemetry import context, trace
from opentelemetry.trace.propagation.tracecontext import (
TraceContextTextMapPropagator,
)
litellm_params: Final = kwargs.get("litellm_params", {}) or {}
proxy_server_request: Final = litellm_params.get("proxy_server_request", {}) or {}
@ -2998,11 +2995,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
# Priority 2: HTTP traceparent header
if traceparent is not None:
verbose_logger.debug("OpenTelemetry: Using traceparent header for context propagation")
carrier: Final = {"traceparent": traceparent}
return (
TraceContextTextMapPropagator().extract(carrier=carrier),
None,
)
return self.get_traceparent_from_header(headers=headers), None
# Priority 3: Active span from global context (auto-detection)
try:

32
litellm/integrations/otel/plumbing/context.py
View file

@ -26,6 +26,7 @@ if TYPE_CHECKING:
from litellm.integrations.otel.model.destination import OtelDestination
_PROPAGATOR: Final = TraceContextTextMapPropagator()
_W3C_TRACE_HEADERS: Final = frozenset(("traceparent", "tracestate"))
# The request's root span — the FastAPI-owned SERVER span — captured ONCE when the
# proxy first resolves it, so request-level spans (the LLM call, guardrails) can
@ -310,6 +311,37 @@ def extract_traceparent(headers: Mapping[str, str]) -> Context | None:
return _PROPAGATOR.extract(carrier)
def _outgoing_trace_context(parent_span: object) -> Context | None:
if isinstance(parent_span, Span) and is_recordable_span(parent_span):
return context_from_span(parent_span)
root: Final = request_root_span()
if root is not None:
return context_from_span(root)
current: Final = get_current()
if is_recordable_span(get_current_span(current)):
return current
return None
def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]:
"""``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span.
Parent preference: ``parent_span`` (the request span auth stashed on the key), then
the anchored request root span, then the ambient active span. Only trace context is
injected, never Baggage. Unchanged when no valid span exists anywhere.
"""
context: Final = _outgoing_trace_context(parent_span)
if context is None:
return dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier
carrier: Final = { # mutable-ok: OpenTelemetry propagator requires a mutable carrier
key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS
}
_PROPAGATOR.inject(carrier, context=context)
return carrier
# The OTLP destinations this request's key or team pointed its traces at, resolved
# once during auth. A ``ContextVar`` for the same reason the root span above is one:
# it rides the request task's context into the ``asyncio.create_task`` children that

18
litellm/model_prices_and_context_window_backup.json
View file

@ -377,8 +377,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 1.5e-08
"supports_tool_choice": true
},
"amazon.nova-2-lite-v1:0": {
"cache_read_input_token_cost": 7.5e-08,
@ -561,8 +560,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 8.75e-09
"supports_tool_choice": true
},
"amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2e-07,
@ -578,8 +576,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 2e-07
"supports_tool_choice": true
},
"amazon.nova-sonic-v1:0": {
"deprecation_date": "2026-09-14",
@ -45794,8 +45791,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 1.5e-08
"supports_tool_choice": true
},
"us.amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 8.75e-09,
@ -45809,8 +45805,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 8.75e-09
"supports_tool_choice": true
},
"us.amazon.nova-premier-v1:0": {
"deprecation_date": "2026-09-14",
@ -45842,8 +45837,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 2e-07
"supports_tool_choice": true
},
"us.anthropic.claude-3-5-haiku-20241022-v1:0": {
"cache_creation_input_token_cost": 1e-06,

2
litellm/proxy/common_utils/proxy_rate_limit_error.py
View file

@ -11,7 +11,7 @@ exception types:
an upstream LLM provider returns 429.
* :class:`fastapi.HTTPException` (status 429) — raised directly by proxy hooks
such as ``parallel_request_limiter``, ``dynamic_rate_limiter``,
``batch_rate_limiter``, ``max_budget_limiter``, ``max_iterations_limiter``,
``batch_rate_limiter``, ``max_iterations_limiter``,
etc.
* :class:`litellm.llms.base_llm.chat.transformation.BaseLLMException` (status
429) — raised by some provider transports.

4
litellm/proxy/hooks/__init__.py
View file

@ -4,7 +4,6 @@ from typing import Final, Literal
from . import *
from .cache_control_check import _PROXY_CacheControlCheck
from .litellm_skills import SkillsInjectionHook
from .max_budget_limiter import _PROXY_MaxBudgetLimiter
from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler
from .max_iterations_limiter import _PROXY_MaxIterationsHandler
from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler
@ -18,7 +17,6 @@ from .sensitive_data_routing import _PROXY_SensitiveDataRoutingHandler
# transitively through `enterprise.enterprise_hooks` can resolve `PROXY_HOOKS`
# and `get_proxy_hook` from this partially-initialized module without circling.
PROXY_HOOKS: Final = {
"max_budget_limiter": _PROXY_MaxBudgetLimiter,
"parallel_request_limiter": _PROXY_MaxParallelRequestsHandler_v3,
"cache_control_check": _PROXY_CacheControlCheck,
"responses_id_security": ResponsesIDSecurity,
@ -35,7 +33,7 @@ if os.getenv("LEGACY_MULTI_INSTANCE_RATE_LIMITING", "false").lower() == "true":
def get_proxy_hook(
hook_name: Literal["max_budget_limiter", "managed_files", "parallel_request_limiter", "cache_control_check"] | str,
hook_name: Literal["managed_files", "parallel_request_limiter", "cache_control_check"] | str,
):
"""
Factory method to get a proxy hook instance by name

84
litellm/proxy/hooks/max_budget_limiter.py
View file

@ -1,84 +0,0 @@
from typing import Final
from fastapi import HTTPException
from litellm import verbose_logger
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.exceptions import RateLimitType
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
class _PROXY_MaxBudgetLimiter(CustomLogger):
# Class variables or attributes
def __init__(self):
pass
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
try:
verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook")
max_budget: Final = user_api_key_dict.user_max_budget
user_id: Final = user_api_key_dict.user_id
if max_budget is None or user_id is None:
return
from litellm.proxy.proxy_server import general_settings
if (
user_api_key_dict.team_id is not None
and general_settings.get("apply_user_budget_to_team_keys") is not True
):
return
# The reservation path admits at the strict-`<` boundary and
# atomically pre-fills the same counter we'd read here. Re-checking
# with `>=` would reject a request the reservation already admitted
# when the reservation fills the counter to exactly max_budget.
# Imported lazily to avoid a circular import via proxy.utils.
from litellm.proxy.spend_tracking.budget_reservation import (
get_reserved_counter_keys,
)
user_counter_key: Final = f"spend:user:{user_id}"
if user_counter_key in get_reserved_counter_keys(user_api_key_dict.budget_reservation):
return
from litellm.proxy.proxy_server import get_current_spend
curr_spend: Final = await get_current_spend(
counter_key=user_counter_key,
fallback_spend=user_api_key_dict.user_spend or 0.0,
)
verbose_proxy_logger.debug(
"MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f",
user_id,
curr_spend,
max_budget,
)
# CHECK IF REQUEST ALLOWED
if curr_spend >= max_budget:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None)
raise ProxyRateLimitError(
detail="Max budget limit reached.",
rate_limit_type=RateLimitType.BUDGET,
model=resolved_model,
llm_provider=llm_provider,
)
except HTTPException as e:
raise e
except Exception as e:
verbose_logger.exception(
"litellm.proxy.hooks.max_budget_limiter.py::async_pre_call_hook(): Exception occured - %s", e
)

49
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -986,6 +986,7 @@ async def pass_through_request(
headers=headers,
forward_headers=forward_headers,
)
upstream_headers: Final = _with_trace_context(headers, parent_span=user_api_key_dict.parent_otel_span)
requested_query_params: dict | None = query_params or dict(request.query_params)
@ -1019,7 +1020,7 @@ async def pass_through_request(
verbose_proxy_logger.debug(
"Pass through endpoint sending request to \nURL %s\nheaders: %s\nbody: %s\n",
url,
headers,
upstream_headers,
_parsed_body,
)
@ -1257,7 +1258,7 @@ async def pass_through_request(
additional_args={
"complete_input_dict": _parsed_body,
"api_base": str(logging_url),
"headers": headers,
"headers": upstream_headers,
},
)
stream = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
@ -1274,7 +1275,7 @@ async def pass_through_request(
request=request,
async_client=async_client,
url=url,
headers=headers,
headers=upstream_headers,
requested_query_params=requested_query_params,
stream=True,
)
@ -1286,7 +1287,7 @@ async def pass_through_request(
request.method,
url,
params=requested_query_params,
headers=headers,
headers=upstream_headers,
content=state_raw_body,
)
if state_raw_body is not None
@ -1294,7 +1295,7 @@ async def pass_through_request(
request.method,
url,
params=requested_query_params,
headers=headers,
headers=upstream_headers,
json=_parsed_body,
)
)
@ -1371,7 +1372,7 @@ async def pass_through_request(
raw_body_request: Final = async_client.build_request(
request.method,
url,
headers=headers,
headers=upstream_headers,
params=requested_query_params,
content=state_raw_body,
)
@ -1381,7 +1382,7 @@ async def pass_through_request(
request=request,
async_client=async_client,
url=url,
headers=headers,
headers=upstream_headers,
requested_query_params=requested_query_params,
_parsed_body=_parsed_body,
forward_multipart=is_multipart,
@ -2158,6 +2159,17 @@ def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None:
return upstream_close
_WEBSOCKET_FORWARDED_HEADERS: Final = frozenset(("authorization", "x-api-key", "x-goog-user-project"))
def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict[str, str]:
try:
from litellm.integrations.otel.plumbing.context import inject_trace_context
except ImportError:
return dict(headers) # mutable-ok: matches inject_trace_context's carrier return type
return inject_trace_context(headers, parent_span=parent_span)
async def websocket_passthrough_request(
websocket: WebSocket,
target: str,
@ -2200,20 +2212,15 @@ async def websocket_passthrough_request(
await websocket.accept()
verbose_proxy_logger.debug("WebSocket passthrough (%s): WebSocket connection accepted", endpoint)
# Prepare headers for the upstream connection
upstream_headers: Final = custom_headers.copy()
if forward_headers:
# Forward relevant headers from the incoming request
incoming_headers: Final = dict(websocket.headers)
for header_name, header_value in incoming_headers.items():
# Only forward certain headers to avoid conflicts
if header_name.lower() in [
"authorization",
"x-api-key",
"x-goog-user-project",
]:
upstream_headers[header_name] = header_value
forwarded_headers: Final = { # mutable-ok: one-shot upstream header dict, read as a Mapping
**custom_headers,
**{
header_name: header_value
for header_name, header_value in websocket.headers.items()
if forward_headers and header_name.lower() in _WEBSOCKET_FORWARDED_HEADERS
},
}
upstream_headers: Final = _with_trace_context(forwarded_headers, parent_span=user_api_key_dict.parent_otel_span)
# Initialize logging object similar to HTTP passthrough
team_callbacks: Final = _resolve_team_callback_wiring(

4
litellm/proxy/utils.py
View file

@ -164,7 +164,6 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrai
)
from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.parallel_request_limiter import (
_PROXY_MaxParallelRequestsHandler,
)
@ -982,7 +981,6 @@ class ProxyLogging:
dual_cache=DualCache(default_in_memory_ttl=1) # ping redis cache every 1s
)
self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler(self.internal_usage_cache)
self.max_budget_limiter = _PROXY_MaxBudgetLimiter()
self.cache_control_check = _PROXY_CacheControlCheck()
self.alerting: list[str] | None = None
self.alerting_threshold: float = 300 # default to 5 min. threshold
@ -3580,7 +3578,7 @@ class ProxyLogging:
caps: Final = ProxyLogging._callback_capabilities()
post_call_pipelines: Final = _streamable_post_call_pipelines(request_data, user_api_key_dict)
# Fast path: no real overrides. Internal proxy CustomLogger callbacks
# (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default
# (e.g. _PROXY_CacheControlCheck, ManagedFiles) inherit the default
# ``async for chunk: yield chunk`` body, so wrapping the iterator
# through each of them adds N pass-through trampolines per chunk for
# zero behavior change. Skip the chain entirely and stream through.

18
model_prices_and_context_window.json
View file

@ -377,8 +377,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 1.5e-08
"supports_tool_choice": true
},
"amazon.nova-2-lite-v1:0": {
"cache_read_input_token_cost": 7.5e-08,
@ -561,8 +560,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 8.75e-09
"supports_tool_choice": true
},
"amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2e-07,
@ -578,8 +576,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 2e-07
"supports_tool_choice": true
},
"amazon.nova-sonic-v1:0": {
"deprecation_date": "2026-09-14",
@ -45794,8 +45791,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 1.5e-08
"supports_tool_choice": true
},
"us.amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 8.75e-09,
@ -45809,8 +45805,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 8.75e-09
"supports_tool_choice": true
},
"us.amazon.nova-premier-v1:0": {
"deprecation_date": "2026-09-14",
@ -45842,8 +45837,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 2e-07
"supports_tool_choice": true
},
"us.anthropic.claude-3-5-haiku-20241022-v1:0": {
"cache_creation_input_token_cost": 1e-06,

54
tests/_live_test_helpers.py
View file

@ -1,6 +1,8 @@
import os
from datetime import date
import pytest
from pydantic import BaseModel, ConfigDict
def _skip_live_prompt_caching_test():
@ -8,3 +10,55 @@ def _skip_live_prompt_caching_test():
pytest.skip("Live prompt-caching E2E tests are opt-in")
if os.environ.get("CASSETTE_REDIS_URL"):
pytest.skip("Live prompt-caching E2E tests cannot run under VCR replay")
class TogetherCostEntry(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
litellm_provider: str | None = None
mode: str | None = None
deprecation_date: str | None = None
input_cost_per_token: float | None = None
output_cost_per_token: float | None = None
supports_function_calling: bool | None = None
supports_response_schema: bool | None = None
def cheapest_together_chat_model(
*, function_calling: bool = False, response_schema: bool = False
) -> str:
import litellm
today = date.today().isoformat()
def qualifies(name: str, entry: TogetherCostEntry) -> bool:
return (
name.startswith("together_ai/")
and entry.litellm_provider == "together_ai"
and entry.mode == "chat"
and (entry.deprecation_date is None or entry.deprecation_date > today)
and (entry.input_cost_per_token or 0.0) > 0
and (entry.output_cost_per_token or 0.0) > 0
and (not function_calling or bool(entry.supports_function_calling))
and (not response_schema or bool(entry.supports_response_schema))
)
registry: dict[str, TogetherCostEntry] = {
name: TogetherCostEntry.model_validate(raw)
for name, raw in litellm.model_cost.items()
if isinstance(raw, dict) and name.startswith("together_ai/")
}
candidates = sorted(
(name for name, entry in registry.items() if qualifies(name, entry)),
key=lambda name: (
registry[name].input_cost_per_token or 0.0,
registry[name].output_cost_per_token or 0.0,
name,
),
)
assert candidates, (
"no live together_ai chat model in the cost map satisfies "
f"function_calling={function_calling} response_schema={response_schema}"
)
return candidates[0]

962
tests/code_coverage_tests/test_provider_cache.py
View file

File diff suppressed because it is too large Load diff

44
tests/e2e/PROVIDER_CACHE.md
View file

@ -1,11 +1,41 @@
# Shared provider-response cache
`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live
`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live
The edge caches complete successful POST responses for `/v1/chat/completions` and `/v1/messages`, including streams. Unsupported endpoints pass through. It matches the method, original URL, effective outbound headers (including authentication and HTTP-library defaults), body presence and exact body bytes using a full keyed digest. It sends the same prepared request used for matching. No prompts, random markers, JSON values or credentials are normalized away. Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies
The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored
Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, cache too. AWS frames those as binary `vnd.amazon.eventstream` rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so streaming is most of the suite's Bedrock traffic
Two details of that rule are worth knowing before changing it. A ConverseStream ends with `metadata`, not with `messageStop`, and the `metadata` frame is what carries the token usage litellm prices the call from, so the rule requires it: a stream cut between the two still names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser, which yields the frames it did receive and silently discards a trailing partial one, so the body is also checked against the frame lengths it declares. A stream cut one byte short parses clean and has to be caught that way
## Request identity
A recording belongs to one test. The key is a keyed digest over the test's node id, the method, the URL, the effective outbound headers (including authentication and HTTP-library defaults), body presence and the body bytes, with one normalization: a 12-hex-digit run, the shape `unique_marker()` mints, is replaced by a placeholder in both the URL and a UTF-8 body. Nothing else is normalized away. No prompts, JSON values or credentials are rewritten, and the rule is the one `fixture_canonical.py` already applies for record/replay, so there is a single definition of what a marker is
Requests that differ only by their markers therefore share a canonical identity, which is what makes the cache reusable across builds: every e2e test salts its prompt afresh, so an exact-byte key would miss on every call. Within one test, calls that share a canonical identity are still recorded and replayed separately, by a FIFO slot index appended to the key. That matters because a replayed response carries the recorded provider response id, `LiteLLM_SpendLogs.request_id` is that id, and one shared recording answering two calls would collapse two spend rows into one
Two different tests never share a recording, and a provider call made outside any test (fixtures, session setup) is never cached, because the identity has no test node id to bind to
A client that varies its own request between runs defeats that identity without breaking any rule, and the Claude Code compat cells did. The CLI sends a device id and a session id in `metadata.user_id`, and its system prompt names both its memory directory and its working directory, adding the branch and recent commits when that directory is a git repository. Driven with a fresh HOME and the checkout as its working directory, every cell sent different bytes every build. The fix belongs in the driver rather than here: `claude_code/cli_driver.py` pins the config directory, the working directory and both identifiers, which is why the cache needs no rule for any of it. Normalizing them instead would have hidden a real defect class, since a rule cannot tell a client's own churn from a value a test means to assert on
Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies
An eligible miss calls the provider. A complete successful response is stored immediately even if a later test assertion fails. Provider errors, malformed responses, truncated streams and cancelled captures are not stored. Cache reads, writes and lease failures fall through to normal provider behavior; they introduce no provider retry. An already-started response cannot be restarted after a delivery failure
## Bedrock
Bedrock could not be mounted before because SigV4 signs the `Host` header, so a rewritten `api_base` failed signature verification at the provider. The edge now re-signs: it drops the proxy's signature headers, signs the upstream request with the run pod's own AWS identity from its EKS Pod Identity association, and forwards that. The signature headers are excluded from the key, since `x-amz-date` is a timestamp and keying on it would make every Bedrock call a permanent miss
Almost every Bedrock deployment in the suite declares its region as `os.environ/AWS_REGION`, which only the proxy can resolve, and the run pod does not share that environment. A `us.` inference profile fans out across the US regions and is reachable from any of them, so those route to the default mount whatever the proxy resolved. A model that is not cross-region and declares its region that way keeps its direct path rather than being sent to a region it may not exist in.
Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove
Which models route is an explicit allowlist in `provider_cache_routing.py`, mirroring the runner role's IAM policy, which names its models one by one. That coupling is deliberate: the edge re-signs with the run pod's identity, so a model the role cannot invoke comes back 403 from Bedrock rather than falling back. An unlisted model keeps its direct path and loses only caching, so adding a Bedrock model to the suite can never turn it red. Adding one to the edge is a policy edit in litellm-ops plus a line here
Vertex and Gemini are not mounted, for different reasons. litellm grafts the default Vertex path onto an `api_base` only when that `api_base` has no path of its own, so a path-prefixed Vertex mount instead becomes `{api_base}:{endpoint}`, dropping project, location and model. Vertex needs a root-mounted edge on its own port, or a change in litellm
Gemini reaches a path-prefixed mount perfectly well and was mounted for one build, then backed out, because litellm's two Gemini endpoints disagree about what `api_base` means. Chat composes `{api_base}/models/{model}:{endpoint}` and defaults `api_base` to `https://generativelanguage.googleapis.com/v1beta`, so the version has to be inside it. File upload composes `{api_base}/upload/v1beta/files` and defaults to the host root, so the version has to be outside it. One `api_base` cannot satisfy both, and a deployment gives no signal at registration time about which it will be used for, so mounting Gemini turned `TestGeminiFiles::test_gemini_file_upload` red in build 227. Anyone pointing litellm's Gemini provider at an AI gateway or a corporate proxy hits the same thing; it is a litellm bug rather than a cache limitation, and mounting Gemini is one line once it is fixed
Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires
## Configuration
@ -18,16 +48,18 @@ The trusted runner receives:
- `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision
- `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory
Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits
Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_error_status` (the provider answered, with an error), `rejected_incomplete` (the body arrived whole with a success status and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached at all). A mount whose rejections are nearly all of one kind is a different problem from one whose rejections are nearly all of another, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits
Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay
Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read.
One more class needs it, and it is the cost of normalizing the marker. A test that mints a fresh marker, sends it, and then asserts the provider's answer contains that exact value is asserting on the marker rather than using it as a salt. The key treats two such requests as the same identity, so a stale recording matches and answers with the marker from the run that recorded it. `TestOpenAIMessagesToolContinuation` is the one in the suite today: it sends a freshly minted receipt through a tool result and asserts the model echoes it back verbatim. If you add a test that asserts a provider echoed your own unique value, it belongs on the live path. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay
## Recorded response semantics
Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Existing spend reconciliation requests use distinct prompt markers and retain their distinct-ID and row-count assertions; accounting tests are not automatically excluded from caching
Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Spend reconciliation keeps its distinct-ID and row-count assertions: its prompts differ by an index as well as a marker, so they stay distinct once markers are normalized, and calls that are canonically equal within one test take separate FIFO slots and separate recordings anyway. Accounting tests are not automatically excluded from caching
Provider remaining-quota headers describe the captured response. Metrics derived from them are historical on a cache hit, not a measurement of current provider capacity. Gateway-generated API-key quota headers are a separate contract. A test of fresh provider quota or timing must use the live-provider policy; replay can still exercise how the proxy processes the recorded headers
## Qualification
`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence
`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis, including the marker-canonical key, the FIFO slot index, per-test isolation, SigV4 re-signing against a local upstream, and each endpoint's completeness rule. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence

162
tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py Normal file
View file

@ -0,0 +1,162 @@
"""The CLI must send the same request bytes from one build to the next.
Markerless harness test: it drives the real `claude` binary against a local
stub instead of a proxy, so it carries no `e2e` marker. The binary is a
prerequisite of this whole suite, so a missing one is a failure rather than a
skip.
Two builds differ in ways the driver does not control: a fresh pod, so no CLI
state survives, and a different candidate checked out at a different commit.
Both used to reach the request body, through the memory path the system prompt
names and through the git block the CLI adds for its working directory, so the
shared provider cache missed on every Claude Code cell. This replays those two
differences across a pair of invocations and holds the bytes equal.
A pinned session id is what makes the second test necessary. The matrix runs
its cells across xdist workers, and the CLI refuses to start a session id that
another live process already holds, so pinning one without also opting out of
session persistence turns most of a parallel run red.
"""
from __future__ import annotations
import json
import os
import shutil
import subprocess
import threading
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import List, Tuple
import pytest
from claude_code.cli_driver import _FIXED_CLI_USER_ID, _seed_cli_identity, _stable_cli_state, run_claude
from claude_code.rate_limiter import RateLimiter
pytestmark = pytest.mark.cli_determinism
_STUB_REPLY = {
"id": "msg_stub",
"type": "message",
"role": "assistant",
"model": "claude-haiku-4-5",
"content": [{"type": "text", "text": "ok"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 2},
}
def _make_repo(root: Path, subject: str) -> Path:
root.mkdir(parents=True, exist_ok=True)
identity = {"NAME": "t", "EMAIL": "t@e2e"}
env = dict(
os.environ,
**{f"GIT_{role}_{key}": value for role in ("AUTHOR", "COMMITTER") for key, value in identity.items()},
)
(root / "file.txt").write_text(subject, encoding="utf-8")
for args in (["init", "-q"], ["add", "."], ["commit", "-q", "-m", subject]):
subprocess.run(["git", *args], cwd=root, env=env, check=True, capture_output=True)
return root
@pytest.fixture(name="captured")
def _captured() -> Tuple[str, List[bytes]]:
bodies: List[bytes] = []
lock = threading.Lock()
class Handler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self) -> None:
raw = self.rfile.read(int(self.headers.get("content-length") or 0))
if "count_tokens" not in self.path:
with lock:
bodies.append(raw)
payload = json.dumps({"input_tokens": 10} if "count_tokens" in self.path else _STUB_REPLY).encode()
self.send_response(200)
self.send_header("content-type", "application/json")
self.send_header("content-length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
def log_message(self, *_args: object) -> None:
return
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
threading.Thread(target=server.serve_forever, daemon=True).start()
try:
yield f"http://127.0.0.1:{server.server_address[1]}", bodies
finally:
server.shutdown()
def test_two_builds_send_the_same_request_bytes(captured: Tuple[str, List[bytes]], tmp_path: Path) -> None:
base_url, bodies = captured
limiter = RateLimiter(state_dir=tmp_path / "limiter")
checkouts = (_make_repo(tmp_path / "build-1", "first"), _make_repo(tmp_path / "build-2", "second"))
origin = Path.cwd()
sent = []
for checkout in checkouts:
shutil.rmtree(Path(_stable_cli_state()[0]).parent, ignore_errors=True)
os.chdir(checkout)
try:
before = len(bodies)
run_claude(
prompt="say ok",
model="claude-haiku-4-5",
base_url=base_url,
api_key="stub",
extra_env={"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"},
rate_limiter=limiter,
)
sent.append(bodies[before:])
finally:
os.chdir(origin)
assert sent[0], "the CLI sent no request to the stub, so there is nothing to compare"
assert sent[0] == sent[1]
def test_concurrent_cells_do_not_collide_on_the_pinned_session(
captured: Tuple[str, List[bytes]], tmp_path: Path
) -> None:
base_url, bodies = captured
limiter = RateLimiter(state_dir=tmp_path / "limiter")
def one(_index: int) -> int:
return run_claude(
prompt="say ok",
model="claude-haiku-4-5",
base_url=base_url,
api_key="stub",
extra_env={"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"},
rate_limiter=limiter,
).exit_code
with ThreadPoolExecutor(max_workers=4) as pool:
codes = list(pool.map(one, range(4)))
assert codes == [0, 0, 0, 0]
assert bodies, "the CLI sent no request to the stub, so there is nothing to compare"
assert set(Counter(bodies).values()) == {4}
def test_seeding_the_device_id_survives_threads_racing_on_the_same_directory(tmp_path: Path) -> None:
"""`run_claude_models_parallel` drives several models from one process, so the
seed's staged file has to be unique per thread and not merely per process."""
config_dir = tmp_path / "config"
config_dir.mkdir()
seeded = config_dir / ".claude.json"
for _round in range(20):
seeded.unlink(missing_ok=True)
with ThreadPoolExecutor(max_workers=16) as pool:
for outcome in [pool.submit(_seed_cli_identity, str(config_dir)) for _ in range(16)]:
outcome.result()
assert json.loads(seeded.read_text(encoding="utf-8"))["userID"] == _FIXED_CLI_USER_ID
assert sorted(entry.name for entry in config_dir.iterdir()) == [".claude.json"]

62
tests/e2e/claude_code/cli_driver.py
View file

@ -132,6 +132,62 @@ def _make_isolated_home() -> str:
return tempfile.mkdtemp(prefix="claude-cli-home-")
_FIXED_CLI_USER_ID = "0" * 64
_FIXED_CLI_SESSION_ID = "00000000-0000-4000-8000-000000000000"
def _seed_cli_identity(config_dir: str) -> None:
"""Pin the device id the CLI would otherwise mint per config directory.
It mints 32 random bytes on first run, writes them to `.claude.json` as
`userID`, and sends them in `metadata.user_id` forever after, so the value
is stable for exactly as long as that file lives. Pinning it, and the
session id passed beside it, costs nothing: both feed abuse detection
rather than quota, caching or continuity.
The staged name has to be unique per *thread*, not per process:
`run_claude_models_parallel` drives several models from one process, so a
pid-suffixed name lets one thread rename the file another is still
writing, and the loser dies on a missing path."""
path = os.path.join(config_dir, ".claude.json")
try:
with open(path, encoding="utf-8") as handle:
if json.load(handle).get("userID") == _FIXED_CLI_USER_ID:
return
except (OSError, ValueError):
pass
handle_fd, staged = tempfile.mkstemp(dir=config_dir, prefix=".claude.json.")
with os.fdopen(handle_fd, "w", encoding="utf-8") as handle:
json.dump({"userID": _FIXED_CLI_USER_ID}, handle)
os.replace(staged, path)
def _stable_cli_state() -> Tuple[str, str]:
"""Config directory and working directory for the CLI, at fixed paths.
Both reach the request body. The memory directory the system prompt
names is `$CLAUDE_CONFIG_DIR/projects/<cwd slug>/memory`, and a working
directory inside a git repository also contributes its branch and recent
commits. So a per-invocation config directory rewrites every body, and
inheriting the checkout rewrites every body once per candidate, which is
why the shared provider cache could never serve a Claude Code cell.
Pinning both makes the bodies repeatable across builds.
This narrows what survives rather than widening it: HOME stays fresh and
empty per invocation, so the isolation `_make_isolated_home` describes is
unchanged, and the CLI's own state no longer outlives the pod either. The
working directory is deliberately not the checkout, so a model-directed
`Read` sees an empty directory instead of the repository.
"""
root = os.path.join(tempfile.gettempdir(), f"litellm-e2e-claude-{os.getuid()}")
config_dir = os.path.join(root, "config")
workspace = os.path.join(root, "workspace")
for path in (root, config_dir, workspace):
os.makedirs(path, mode=0o700, exist_ok=True)
_seed_cli_identity(config_dir)
return config_dir, workspace
class ClaudeCLIError(RuntimeError):
"""Raised when the `claude` CLI cannot be invoked or returns a fatal error."""
@ -222,6 +278,9 @@ def run_claude(
"--verbose",
"--model",
model,
"--session-id",
_FIXED_CLI_SESSION_ID,
"--no-session-persistence",
]
if extra_args:
cmd.extend(extra_args)
@ -244,6 +303,8 @@ def run_claude(
# regardless of how the subprocess exits.
isolated_home = _make_isolated_home()
env["HOME"] = isolated_home
config_dir, workspace = _stable_cli_state()
env["CLAUDE_CONFIG_DIR"] = config_dir
if extra_env:
env.update(extra_env)
@ -262,6 +323,7 @@ def run_claude(
completed = run_fn(
cmd,
env=env,
cwd=workspace,
input=stdin_input,
capture_output=True,
text=True,

13
tests/e2e/conftest.py
View file

@ -23,6 +23,7 @@ from typing import Final
import pytest
import requests
from e2e_config import (
CLI_DETERMINISM_OPT_IN_ENV,
CONTROL_PLANE_BASE_URL,
FIXTURE_DIR,
FIXTURE_MODE_RAW,
@ -53,6 +54,7 @@ OPT_IN_MARKERS: Final = MappingProxyType(
"managed_files": MANAGED_FILES_OPT_IN_ENV,
"prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV,
"redis_chaos": REDIS_CHAOS_OPT_IN_ENV,
"cli_determinism": CLI_DETERMINISM_OPT_IN_ENV,
}
)
@ -85,7 +87,11 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient)
def pytest_configure(config: pytest.Config) -> None:
config.addinivalue_line("markers", "provider_live: requires actual provider timing, limits or state; bypass shared cache")
config.addinivalue_line(
"markers",
"provider_live: requires actual provider timing, limits, state, or a response that echoes this"
" run's own unique value; bypass shared cache",
)
config.addinivalue_line(
"markers",
"e2e: live test that requires a running proxy and real provider keys",
@ -116,6 +122,11 @@ def pytest_configure(config: pytest.Config) -> None:
"prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including "
"prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set",
)
config.addinivalue_line(
"markers",
"cli_determinism: drives the real claude CLI for several seconds, which widens the window in which "
"another test's in-flight upstream call is attributed to it; deselected unless E2E_CLI_DETERMINISM is set",
)
config.addinivalue_line(
"markers",
"redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from "

1
tests/e2e/e2e_config.py
View file

@ -145,6 +145,7 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY"
MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK"
PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK"
REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS"
CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM"
ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))
ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3"))

5
tests/e2e/fixture_canonical.py
View file

@ -51,6 +51,9 @@ SECRET_FIELD_SUFFIXES: Final[tuple[str, ...]] = (
)
SECRET_PLACEHOLDER: Final = "<secret>"
MARKER_PATTERN: Final = re.compile(r"(?<![0-9a-fA-F])[0-9a-f]{12}(?![0-9a-fA-F])")
MARKER_PLACEHOLDER: Final = "<marker>"
PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = (
(re.compile(r"(?<![0-9a-fA-F])[0-9a-f]{64}(?![0-9a-fA-F])"), "<sha256>"),
(
@ -67,7 +70,7 @@ PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = (
re.compile(r"\b(?:chatcmpl|msgbatch|msg|resp|batch|call|req|ftjob|gen|file)[-_][A-Za-z0-9]{8,}\b"),
"<id>",
),
(re.compile(r"(?<![0-9a-fA-F])[0-9a-f]{12}(?![0-9a-fA-F])"), "<marker>"),
(MARKER_PATTERN, MARKER_PLACEHOLDER),
)

1
tests/e2e/llm_translation/test_messages_e2e.py
View file

@ -372,6 +372,7 @@ def _request_tool(
class TestOpenAIMessagesToolContinuation:
@pytest.mark.provider_live
@pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"])
def test_required_tool_arguments_and_correlated_result(
self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool

1
tests/e2e/models.py
View file

@ -951,6 +951,7 @@ class LiteLLMParamsBody(BaseModel):
aws_access_key_id: str | None = None
aws_secret_access_key: str | None = None
aws_region_name: str | None = None
aws_bedrock_runtime_endpoint: str | None = None
vertex_project: str | None = None
vertex_location: str | None = None
vertex_credentials: str | None = None

459
tests/e2e/provider_cache.py
View file

@ -4,14 +4,18 @@ import base64
import hashlib
import hmac
import io
import json
import os
import threading
import time
from collections.abc import Callable, Generator, Mapping
from contextlib import closing
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import Final, Literal, Protocol
from urllib.parse import urlsplit
from botocore.eventstream import EventStreamBuffer, ParserError
from e2e_http import (
NetworkError,
StreamChunk,
@ -23,12 +27,36 @@ from e2e_http import (
prepare_forward,
primed_steps,
)
from fixture_canonical import MARKER_PATTERN, MARKER_PLACEHOLDER
from fixture_mode import SESSION_TEST_KEY, current_test_key
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
LIFETIME_SECONDS: Final = 86_400
MAX_REQUEST_BYTES: Final = 256 * 1024
MAX_RESPONSE_BYTES: Final = 8 * 1024 * 1024
UNRECORDED_RESPONSE_HEADERS: Final = frozenset({"set-cookie"})
SIGNATURE_HEADERS: Final = frozenset(
{"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"}
)
BEDROCK_MOUNT_PREFIX: Final = "bedrock"
BEDROCK_CONVERSE_SUFFIX: Final = "/converse"
BEDROCK_INVOKE_SUFFIX: Final = "/invoke"
BEDROCK_CONVERSE_STREAM_SUFFIX: Final = "/converse-stream"
BEDROCK_INVOKE_STREAM_SUFFIX: Final = "/invoke-with-response-stream"
BEDROCK_SUFFIXES: Final = (
BEDROCK_CONVERSE_SUFFIX,
BEDROCK_INVOKE_SUFFIX,
BEDROCK_CONVERSE_STREAM_SUFFIX,
BEDROCK_INVOKE_STREAM_SUFFIX,
)
EVENTSTREAM_PRELUDE_BYTES: Final = 4
CUT_SHORT: Final = "cut_short"
INCOMPLETE: Final = "incomplete"
UNREACHABLE: Final = "unreachable"
ERROR_STATUS: Final = "error_status"
EVENT_TYPE_HEADER: Final = ":event-type"
EVENTSTREAM_HEADERS: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str])
OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"})
JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
@ -56,6 +84,24 @@ class CacheUnavailable:
type CacheLookup = CacheHit | CaptureLease | CacheBusy | CacheUnavailable
type RequestSigner = Callable[[str, str, Mapping[str, str], bytes | None], dict[str, str]]
@dataclass(frozen=True, slots=True)
class MountPolicy:
"""What a mount needs beyond plain forwarding.
``sign`` mints a fresh credential over the upstream URL, for providers whose
auth covers the Host the edge rewrote. ``unkeyed_headers`` names headers that
must stay out of the cache key because they change on every call and would
otherwise make the mount a permanent miss: a minted signature, or an OAuth
token the provider rotates. Naming one costs the guarantee that a recording
can never cross credentials, so a mount with a rotating token relies on the
environment holding one identity for that provider. Mounts with a static API
key name nothing here and keep the guarantee whole."""
sign: RequestSigner | None = None
unkeyed_headers: frozenset[str] = frozenset()
class ResponseStore(Protocol):
@ -83,28 +129,51 @@ class SignedResponse(BaseModel):
signature: str
def exact_key(secret: bytes, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> str:
def canonical_text(value: str) -> str:
return MARKER_PATTERN.sub(MARKER_PLACEHOLDER, value)
def canonical_body(body: bytes) -> bytes:
try:
return canonical_text(body.decode("utf-8")).encode("utf-8")
except UnicodeDecodeError:
return body
def request_identity(
secret: bytes, test_key: str, method: str, url: str, headers: Mapping[str, str], body: bytes | None,
) -> str:
fields: Final = (
b"provider-cache-exact-v1", method.encode(), url.encode(),
b"provider-cache-canonical-v2", test_key.encode(), method.encode(), canonical_text(url).encode(),
*(part.encode() for pair in sorted(headers.items()) for part in pair),
b"no-body" if body is None else b"body", b"" if body is None else body,
b"no-body" if body is None else b"body", b"" if body is None else canonical_body(body),
)
encoded: Final = b"".join(len(part).to_bytes(8, "big") + part for part in fields)
return hmac.new(secret, encoded, hashlib.sha256).hexdigest()
def cacheable_endpoint(method: str, url: str, body: bytes | None) -> bool:
return (
method == "POST"
and urlsplit(url).path in {"/v1/chat/completions", "/v1/messages"}
and body is not None
and len(body) <= MAX_REQUEST_BYTES
)
def slotted_key(secret: bytes, identity: str, slot: int) -> str:
return hmac.new(secret, f"{identity}:{slot}".encode(), hashlib.sha256).hexdigest()
def successful_response(url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool:
def is_bedrock(mount: str) -> bool:
return mount.partition("/")[0] == BEDROCK_MOUNT_PREFIX
def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> bool:
if method != "POST" or body is None or len(body) > MAX_REQUEST_BYTES:
return False
path: Final = urlsplit(url).path
if is_bedrock(mount):
return path.startswith("/model/") and path.endswith(BEDROCK_SUFFIXES)
return path in OPENAI_JSON_PATHS
def successful_response(mount: str, url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool:
if not 200 <= status < 300 or len(body) > MAX_RESPONSE_BYTES:
return False
if is_bedrock(mount):
return complete_bedrock_response(url, body)
streaming: Final = "text/event-stream" in headers.get("content-type", "").lower()
if streaming:
try:
@ -118,28 +187,33 @@ def successful_response(url: str, status: int, headers: Mapping[str, str], body:
values: Final = tuple(JSON_VALUE.validate_json(event) for event in events if event != "[DONE]")
except (UnicodeDecodeError, ValidationError):
return False
if not values or any(not isinstance(value, dict) or "error" in value or value.get("type") == "error" for value in values):
if not values or any(
not isinstance(value, dict) or value.get("error") is not None or value.get("type") == "error"
for value in values
):
return False
if urlsplit(url).path == "/v1/responses":
return complete_responses_stream(values)
if urlsplit(url).path == "/v1/chat/completions":
return events[-1] == "[DONE]" and "[DONE]" not in events[:-1] and complete_chat_stream(values)
return (
"[DONE]" not in events
and isinstance(values[0], dict) and values[0].get("type") == "message_start"
and isinstance(values[-1], dict) and values[-1].get("type") == "message_stop"
and any(
isinstance(value, dict) and value.get("type") == "message_delta"
and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str)
for value in values
)
)
return "[DONE]" not in events and complete_anthropic_stream(values)
try:
value: Final = JSON_VALUE.validate_json(body)
except ValidationError:
return False
if not isinstance(value, dict) or "error" in value:
if not isinstance(value, dict) or value.get("error") is not None:
return False
if urlsplit(url).path == "/v1/messages":
path: Final = urlsplit(url).path
if path == "/v1/messages":
return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str)
if path == "/v1/embeddings":
data: Final = value.get("data")
return isinstance(data, list) and bool(data) and isinstance(value.get("usage"), dict) and all(
isinstance(item, dict) and isinstance(item.get("embedding"), list) and bool(item["embedding"])
for item in data
)
if path == "/v1/responses":
return value.get("object") == "response" and value.get("status") == "completed"
choices: Final = value.get("choices")
return isinstance(choices, list) and bool(choices) and all(
isinstance(choice, dict) and isinstance(choice.get("message"), dict) and isinstance(choice.get("finish_reason"), str)
@ -147,6 +221,144 @@ def successful_response(url: str, status: int, headers: Mapping[str, str], body:
)
def complete_bedrock_response(url: str, body: bytes) -> bool:
"""Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an
Anthropic model answers the Anthropic message shape. Either way a truncated
or error body is missing the terminator field, which is what makes it safe to
record."""
path: Final = urlsplit(url).path
if path.endswith(BEDROCK_CONVERSE_STREAM_SUFFIX):
return complete_converse_stream(body)
if path.endswith(BEDROCK_INVOKE_STREAM_SUFFIX):
return complete_invoke_stream(body)
try:
value: Final = JSON_VALUE.validate_json(body)
except ValidationError:
return False
if not isinstance(value, dict) or "message" in value:
return False
if path.endswith(BEDROCK_CONVERSE_SUFFIX):
return isinstance(value.get("output"), dict) and isinstance(value.get("stopReason"), str)
return (
value.get("type") == "message"
and isinstance(value.get("content"), list)
and isinstance(value.get("stop_reason"), str)
)
def whole_eventstream_messages(body: bytes) -> bool:
"""Whether the body is exactly a whole number of eventstream messages.
A dropped connection is the failure this catches, and it has to be caught
here: botocore yields the messages it did receive and silently discards a
trailing partial one, so a stream cut a single byte short parses clean. Each
message declares its own total length in its first four bytes, so walking
those is enough to tell a complete body from a cut one."""
offset = 0 # rebind-ok: a cursor walking the declared frame lengths
while offset + EVENTSTREAM_PRELUDE_BYTES <= len(body):
total: int = int.from_bytes(body[offset : offset + EVENTSTREAM_PRELUDE_BYTES], "big")
if total <= 0 or offset + total > len(body):
return False
offset += total
return offset == len(body)
def eventstream_events(body: bytes) -> tuple[tuple[str, JsonValue], ...] | None:
"""The stream's (event type, decoded payload) pairs, or None if it is not a
complete, uncorrupted stream.
botocore validates both CRCs and raises ``ParserError`` rather than decoding
corruption into something plausible. A failure that began after Bedrock had
already answered 200 arrives as an ``exception`` frame in place of the
terminator, so it is the terminator rules below that reject it and this does
not need to inspect ``:message-type`` as well."""
if not body or not whole_eventstream_messages(body):
return None
buffer: Final = EventStreamBuffer()
buffer.add_data(body)
try:
return tuple(
(event_type(event.headers), JSON_VALUE.validate_json(event.payload))
for event in buffer
)
except (ParserError, ValidationError, ValueError):
return None
def event_type(headers: object) -> str:
"""botocore's eventstream headers come back untyped, so the one header this
reads is validated into a string rather than trusted."""
parsed: Final = EVENTSTREAM_HEADERS.validate_python(headers)
return parsed.get(EVENT_TYPE_HEADER, "")
def complete_converse_stream(body: bytes) -> bool:
"""ConverseStream ends with ``metadata``, not with ``messageStop``.
Requiring the metadata frame rather than the stop frame is deliberate: it
carries the token usage litellm prices the call from, so a stream cut between
the two still names a stop reason but would replay as a free call."""
events: Final = eventstream_events(body)
if not events or events[-1][0] != "metadata":
return False
return any(
event_type == "messageStop" and isinstance(payload, dict) and isinstance(payload.get("stopReason"), str)
for event_type, payload in events
)
def complete_invoke_stream(body: bytes) -> bool:
"""InvokeModelWithResponseStream wraps the ordinary Anthropic event grammar
in ``chunk`` frames, one base64 payload each, so it is held to the same
terminator rule as the Anthropic SSE path. A frame Bedrock sends instead of a
chunk, an exception among them, carries no such payload and fails the rule
without the frame type needing to be read."""
events: Final = eventstream_events(body)
if not events:
return False
values: Final = tuple(invoke_chunk_value(payload) for _, payload in events)
return all(value is not None for value in values) and complete_anthropic_stream(values)
def invoke_chunk_value(payload: JsonValue) -> JsonValue | None:
"""The Anthropic event inside one ``chunk`` frame, or None for a frame that
carries no readable one."""
if not isinstance(payload, dict) or not isinstance(encoded := payload.get("bytes"), str):
return None
try:
return JSON_VALUE.validate_json(base64.b64decode(encoded, validate=True))
except (ValidationError, ValueError):
return None
def complete_anthropic_stream(values: tuple[JsonValue, ...]) -> bool:
"""The Anthropic event grammar, shared by the SSE mounts and by Bedrock's
invoke stream, which carries the same events inside eventstream frames. A
``message_delta`` naming a stop reason is what separates a finished turn from
one the connection cut short."""
if not values:
return False
first: Final = values[0]
last: Final = values[-1]
return (
isinstance(first, dict) and first.get("type") == "message_start"
and isinstance(last, dict) and last.get("type") == "message_stop"
and any(
isinstance(value, dict) and value.get("type") == "message_delta"
and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str)
for value in values
)
)
def complete_responses_stream(values: tuple[JsonValue, ...]) -> bool:
"""The Responses API streams typed events and ends with ``response.completed``.
A run that failed, was cancelled, or ran out of tokens ends with a different
terminal event, so requiring that one keeps a half-finished response out."""
last: Final = values[-1]
return isinstance(last, dict) and last.get("type") == "response.completed"
def complete_chat_stream(values: tuple[JsonValue, ...]) -> bool:
if any(not isinstance(value, dict) or not isinstance(value.get("choices"), list) for value in values):
return False
@ -172,7 +384,7 @@ def encode_response(secret: bytes, response: CachedResponse) -> bytes:
return SignedResponse(response=raw, signature=hmac.new(secret, raw.encode(), hashlib.sha256).hexdigest()).model_dump_json().encode()
def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> CachedResponse | None:
def decode_response(secret: bytes, key: str, payload: bytes, mount: str, url: str) -> CachedResponse | None:
if len(payload) > 2 * MAX_RESPONSE_BYTES:
return None
try:
@ -183,11 +395,58 @@ def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> Cached
chunks: Final = tuple(base64.b64decode(chunk, validate=True) for chunk in response.chunks)
except (ValidationError, ValueError):
return None
if response.request_key != key or not successful_response(url, response.status_code, response.headers, b"".join(chunks)):
if response.request_key != key or not successful_response(
mount, url, response.status_code, response.headers, b"".join(chunks)
):
return None
return response
def component_digests(
test_key: str, method: str, url: str, headers: Mapping[str, str], body: bytes | None,
) -> dict[str, str]:
"""Per-component digests of everything the key covers.
A mount whose corpus never converges is a mount where one of these moves
between builds, and the flat key cannot say which. Values are digested, so
no payload or credential is written, and a JSON body contributes one digest
per top-level field so the field that moved can be named."""
parts: dict[str, str] = { # rebind-ok: a report assembled from three differently shaped sources
"test_key": test_key,
"method": method,
"url": short_digest(canonical_text(url).encode()),
}
for name, value in sorted(headers.items()):
parts[f"header:{name.lower()}"] = short_digest(value.encode())
canonical: Final = b"" if body is None else canonical_body(body)
parts["body"] = short_digest(canonical)
try:
parsed: Final = JSON_VALUE.validate_json(canonical)
except ValidationError:
return parts
if isinstance(parsed, dict):
for name, value in sorted(parsed.items()):
parts[f"body:{name}"] = short_digest(json.dumps(value, sort_keys=True).encode())
return parts
def short_digest(value: bytes) -> str:
return hashlib.sha256(value).hexdigest()[:16]
@dataclass(slots=True)
class KeyProbe:
"""Every keyed request's components, when a metrics directory is configured."""
rows: tuple[tuple[tuple[str, str], ...], ...] = ()
lock: threading.Lock = field(default_factory=threading.Lock)
def observe(self, mount: str, outcome: str, parts: Mapping[str, str]) -> None:
row: Final = tuple({"mount": mount, "outcome": outcome, **parts}.items())
with self.lock:
self.rows = (*self.rows, row)
@dataclass(slots=True)
class CacheCounters:
counts: tuple[tuple[str, int], ...] = ()
@ -199,6 +458,24 @@ class CacheCounters:
self.counts = tuple((current | {name: current.get(name, 0) + 1}).items())
@dataclass(slots=True)
class SlotCounter:
"""FIFO position of a request among the canonically identical ones its test
has already sent. Two calls in one test that differ only by ``unique_marker``
canonicalize the same, so without this they would share one recording and the
second would replay the first's provider response id."""
counts: tuple[tuple[str, int], ...] = ()
lock: threading.Lock = field(default_factory=threading.Lock)
def take(self, identity: str) -> int:
with self.lock:
current: Final = dict(self.counts)
taken: Final = current.get(identity, 0)
self.counts = tuple((current | {identity: taken + 1}).items())
return taken
@dataclass(slots=True)
class ResponseCapture:
buffer: io.BytesIO = field(default_factory=io.BytesIO)
@ -226,14 +503,21 @@ def response_steps(response: CachedResponse) -> Generator[StreamStep, None, None
yield StreamChunk(base64.b64decode(chunk, validate=True))
NO_POLICIES: Final[Mapping[str, MountPolicy]] = MappingProxyType({})
@dataclass(frozen=True, slots=True)
class CacheEdge:
store: ResponseStore
secret: bytes = field(repr=False)
counters: CacheCounters = field(default_factory=CacheCounters)
probe: KeyProbe = field(default_factory=KeyProbe)
slots: SlotCounter = field(default_factory=SlotCounter)
policies: Mapping[str, MountPolicy] = NO_POLICIES
wait_seconds: float = 2.0
clock: Callable[[], float] = time.monotonic
sleep: Callable[[float], None] = time.sleep
test_key: Callable[[], str] = current_test_key
def lookup(self, key: str) -> CacheLookup:
deadline: Final = self.clock() + self.wait_seconds
@ -241,59 +525,122 @@ class CacheEdge:
self.sleep(min(0.05, max(0, deadline - self.clock())))
return result
def forward(self, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float) -> StreamHead | NetworkError:
if not cacheable_endpoint(method, url, body):
self.counters.increment("bypass")
self.counters.increment("upstream_attempts")
return forward_stream(method, url, headers=headers, body=body, timeout=timeout)
prepared: Final = prepare_forward(method, url, headers, body)
def count(self, mount: str, name: str) -> None:
self.counters.increment(name)
self.counters.increment(f"mount:{mount}:{name}")
def record_key(
self, mount: str, outcome: str, test_key: str, method: str, url: str,
headers: Mapping[str, str], body: bytes | None,
) -> None:
if not os.environ.get("E2E_PROVIDER_CACHE_METRICS_DIR"):
return
self.probe.observe(mount, outcome, component_digests(test_key, method, url, headers, body))
def outbound(self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None) -> dict[str, str]:
"""The headers actually sent upstream. A signing mount gets a signature
minted over the upstream URL, because the edge rewrote the Host the proxy
signed and the provider verifies it."""
signer: Final = self.policies.get(mount, MountPolicy()).sign
return headers if signer is None else signer(method, url, headers, body)
def keyed(self, mount: str, headers: Mapping[str, str]) -> Mapping[str, str]:
"""Headers the cache key is built from. A mount keeps its credentials in
the key unless its policy names them unkeyed, so by default one account
can never read another's recording."""
unkeyed: Final = self.policies.get(mount, MountPolicy()).unkeyed_headers
if not unkeyed:
return headers
return {name: value for name, value in headers.items() if name.lower() not in unkeyed}
def forward(
self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float,
) -> StreamHead | NetworkError:
test_key: Final = self.test_key()
if test_key == SESSION_TEST_KEY or not cacheable_endpoint(mount, method, url, body):
self.count(mount, "bypass")
self.count(mount, "upstream_attempts")
return forward_stream(
method, url, headers=self.outbound(mount, method, url, headers, body), body=body, timeout=timeout,
)
prepared: Final = prepare_forward(method, url, self.outbound(mount, method, url, headers, body), body)
if isinstance(prepared, NetworkError):
self.counters.increment("rejected")
self.reject(mount, UNREACHABLE)
return prepared
key: Final = exact_key(self.secret, method, url, prepared.headers, body)
keyed_headers: Final = self.keyed(mount, prepared.headers)
identity: Final = request_identity(self.secret, test_key, method, url, keyed_headers, body)
key: Final = slotted_key(self.secret, identity, self.slots.take(identity))
found: Final = self.lookup(key)
if isinstance(found, CacheHit):
response: Final = decode_response(self.secret, key, found.payload, url)
response: Final = decode_response(self.secret, key, found.payload, mount, url)
if response is not None and self.clock() < found.valid_until:
self.counters.increment("hits")
self.count(mount, "hits")
self.record_key(mount, "hit", test_key, method, url, keyed_headers, body)
return StreamHead(response.status_code, response.headers, response_steps(response))
self.counters.increment("corrupt" if response is None else "expired")
self.count(mount, "corrupt" if response is None else "expired")
self.store.discard(key, found.payload)
capture_slot: Final = self.lookup(key) if isinstance(found, CacheHit) else found
self.counters.increment("misses")
self.count(mount, "misses")
self.record_key(mount, "miss", test_key, method, url, keyed_headers, body)
if isinstance(capture_slot, CacheUnavailable):
self.counters.increment("cache_errors")
self.counters.increment("upstream_attempts")
self.count(mount, "cache_errors")
self.count(mount, "upstream_attempts")
head: Final = forward_prepared_stream(prepared, timeout)
if not isinstance(capture_slot, CaptureLease):
return head
if isinstance(head, NetworkError):
self.store.release(key, capture_slot)
self.counters.increment("rejected")
self.reject(mount, UNREACHABLE)
return head
return StreamHead(head.status_code, head.headers, primed_steps(self.capture(key, capture_slot, url, head)))
return StreamHead(
head.status_code, head.headers, primed_steps(self.capture(mount, key, capture_slot, url, head)),
)
def capture(self, key: str, lease: CaptureLease, url: str, head: StreamHead) -> Generator[StreamStep, None, None]:
def capture(
self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead,
) -> Generator[StreamStep, None, None]:
capture: Final = ResponseCapture()
reason = CUT_SHORT # rebind-ok: a consumer that walks away never reaches the settle call below
try:
with closing(head.steps):
yield StreamChunk(b"")
for step in head.steps:
yield step
capture.observe(step)
chunks: Final = capture.chunks() if capture.eligible else ()
headers: Final = {
name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS
}
if not capture.eligible or not successful_response(url, head.status_code, headers, b"".join(chunks)):
self.counters.increment("rejected")
return
response: Final = CachedResponse(
request_key=key, status_code=head.status_code, headers=headers,
chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks),
)
published: Final = self.store.publish(key, lease, encode_response(self.secret, response))
self.counters.increment("writes" if published else "write_failures")
reason = self.settle(mount, key, lease, url, head, capture)
finally:
self.reject(mount, reason)
self.store.release(key, lease)
capture.buffer.close()
def settle(
self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead, capture: ResponseCapture,
) -> str | None:
"""None once the response is stored, otherwise the reason it was not."""
if not capture.eligible:
return CUT_SHORT
headers: Final = {
name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS
}
if not 200 <= head.status_code < 300:
return ERROR_STATUS
chunks: Final = capture.chunks()
if not successful_response(mount, url, head.status_code, headers, b"".join(chunks)):
return INCOMPLETE
response: Final = CachedResponse(
request_key=key, status_code=head.status_code, headers=headers,
chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks),
)
published: Final = self.store.publish(key, lease, encode_response(self.secret, response))
self.count(mount, "writes" if published else "write_failures")
return None
def reject(self, mount: str, reason: str | None) -> None:
"""A flat rejection count cannot separate a connection that went away from
a body the provider finished sending and the rules turned down, and the two
have opposite fixes. A mount whose rejections are nearly all one or the
other is a different problem, so the report has to be able to say which."""
if reason is None:
return
self.count(mount, "rejected")
self.count(mount, f"rejected_{reason}")

4
tests/e2e/provider_cache_redis.py
View file

@ -134,6 +134,10 @@ def write_metrics(cache: CacheEdge) -> None:
root: Final = Path(directory)
root.mkdir(parents=True, exist_ok=True)
(root / f"{os.getpid()}.json").write_text(report + "\n")
if cache.probe.rows:
(root / f"keys-{os.getpid()}.json").write_text(
json.dumps([dict(row) for row in cache.probe.rows]) + "\n"
)
except OSError:
logging.getLogger(__name__).warning("provider cache metrics artifact unavailable")
logging.getLogger(__name__).info("%s", report)

70
tests/e2e/provider_cache_routing.py
View file

@ -8,14 +8,80 @@ from models import LiteLLMParamsBody, ModelMode
LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False)
DEFAULT_BEDROCK_REGION: Final = "us-east-1"
BEDROCK_CROSS_REGION_PREFIX: Final = "us."
BEDROCK_EDGE_MODELS: Final = frozenset(
{
"us.anthropic.claude-haiku-4-5-20251001-v1:0",
"us.anthropic.claude-sonnet-4-5-20250929-v1:0",
"us.anthropic.claude-sonnet-5",
"us.anthropic.claude-opus-4-7",
}
)
ENV_REFERENCE_PREFIX: Final = "os.environ/"
def bedrock_region(declared: str | None) -> str:
"""The region whose edge mount a deployment belongs to.
Most Bedrock deployments declare `os.environ/AWS_REGION`, which only the
proxy can resolve from its own environment; the run pod does not share it.
Answering those with the default mount is correct because every model on the
edge allowlist is a `us.` inference profile, which fans out across the US
regions and is reachable from any of them. That invariant is enforced on the
allowlist itself rather than re-checked per call."""
if declared is None or declared.startswith(ENV_REFERENCE_PREFIX):
return DEFAULT_BEDROCK_REGION
return declared
def bedrock_mount(params: LiteLLMParamsBody) -> str | None:
"""The edge mount a Bedrock deployment belongs to, or None.
The allowlist mirrors the runner role's IAM policy, which names its models
one by one. A model outside it would be re-signed with an identity that
cannot invoke it and come back 403 from Bedrock, so an unlisted model keeps
its direct path and loses only caching. Adding a model is a policy edit in
litellm-ops and a line here."""
route: Final = params.model.partition("/")[2]
model: Final = route.partition("/")[2] or route
if model not in BEDROCK_EDGE_MODELS:
return None
return f"bedrock/{bedrock_region(params.aws_region_name)}"
def route_bedrock(
params: LiteLLMParamsBody, base_for: Callable[[str], str | None], mode: ModelMode | None,
) -> LiteLLMParamsBody:
"""Deployments that carry their own AWS identity stay off the edge. The edge
re-signs with the run pod's role, so routing an `aws_role_name` deployment
would quietly replace the very assume-role chain that test exists to prove."""
if mode is not None or params.aws_role_name is not None or params.aws_access_key_id is not None:
return params
if params.api_base is not None or params.aws_bedrock_runtime_endpoint is not None:
return params
mount: Final = bedrock_mount(params)
if mount is None:
return params
base: Final = base_for(mount)
if base is None:
return params
return params.model_copy(update={"aws_bedrock_runtime_endpoint": base})
def route_cache_model(
params: LiteLLMParamsBody, base_for: Callable[[str], str | None], *, enabled: bool, mode: ModelMode | None = None,
) -> LiteLLMParamsBody:
if not enabled or mode == "realtime" or LIVE_PROVIDER_REQUIRED.get() or params.api_base is not None or params.mock_response is not None:
if not enabled or LIVE_PROVIDER_REQUIRED.get() or params.mock_response is not None:
return params
if params.litellm_credential_name is not None:
return params
provider: Final = params.model.partition("/")[0]
if provider not in {"openai", "anthropic"} or params.litellm_credential_name is not None:
if provider == "bedrock":
return route_bedrock(params, base_for, mode)
if mode == "realtime" or params.api_base is not None:
return params
if provider not in {"openai", "anthropic"}:
return params
base: Final = base_for(provider)
if base is None:

67
tests/e2e/provider_edge.py
View file

@ -48,7 +48,7 @@ import threading
from collections import deque
from collections.abc import Generator, Mapping, Sequence
from contextlib import closing, contextmanager
from dataclasses import dataclass, field
from dataclasses import dataclass, field, replace
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from itertools import islice
from pathlib import Path
@ -94,17 +94,41 @@ from fixture_mode import (
parse_fixture_mode,
)
from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity
from provider_cache import CacheEdge
from provider_cache import SIGNATURE_HEADERS, CacheEdge, MountPolicy, is_bedrock
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
from pydantic import JsonValue, TypeAdapter
BEDROCK_REGIONS: Final[tuple[str, ...]] = ("us-east-1",)
EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType(
{
"openai": "https://api.openai.com",
"anthropic": "https://api.anthropic.com",
**{
f"bedrock/{region}": f"https://bedrock-runtime.{region}.amazonaws.com"
for region in BEDROCK_REGIONS
},
}
)
@dataclass(frozen=True, slots=True)
class ResolvedMount:
mount: str
upstream_base: str
upstream_path: str
def resolve_mount(path: str, mounts: Mapping[str, str]) -> ResolvedMount | None:
"""Longest mount prefix wins, so a region-qualified mount such as
``bedrock/us-east-1`` resolves whole instead of leaving the region as the
first segment of the upstream path."""
trimmed: Final = path.lstrip("/")
for mount in sorted(mounts, key=len, reverse=True):
if trimmed == mount or trimmed.startswith(f"{mount}/"):
return ResolvedMount(mount, mounts[mount], trimmed[len(mount):].lstrip("/"))
return None
REPLAY_MISS_STATUS: Final = 599
_HOP_BY_HOP_HEADERS: Final[frozenset[str]] = frozenset(
@ -754,14 +778,14 @@ def _handle_record(
def _handle_live(
method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float,
cache: CacheEdge | None = None,
cache: CacheEdge | None = None, mount: str = "",
) -> EdgeOutcome:
forwarded: Final = {
name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS
}
head: Final = (
forward_stream(method, url, headers=forwarded, body=body, timeout=timeout)
if cache is None else cache.forward(method, url, forwarded, body, timeout)
if cache is None else cache.forward(mount, method, url, forwarded, body, timeout)
)
match head:
case NetworkError(message=message):
@ -796,10 +820,13 @@ def handle_edge_request(
prefix, then record (forward + persist) or replay (serve from the bundle).
Socket-free so unit tests exercise every branch without a server."""
split: Final = urlsplit(raw_path)
mount, _, upstream_path = split.path.lstrip("/").partition("/")
upstream_base: Final = mounts.get(mount)
if upstream_base is None:
return _text_reply(404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}")
resolved: Final = resolve_mount(split.path, mounts)
if resolved is None:
unknown: Final = split.path.lstrip("/").partition("/")[0]
return _text_reply(404, f"unknown provider mount {unknown!r}; known mounts: {', '.join(sorted(mounts))}")
mount: Final = resolved.mount
upstream_base: Final = resolved.upstream_base
upstream_path: Final = resolved.upstream_path
profile: Final = (
backend.recorder.profile
if isinstance(backend, RecordEdge)
@ -830,7 +857,8 @@ def handle_edge_request(
match backend:
case CacheEdge():
return _handle_live(
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend,
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout,
backend, mount,
)
case LiveEdge():
return _handle_live(
@ -891,7 +919,7 @@ class _EdgeHandler(BaseHTTPRequestHandler):
)
if isinstance(edge_server.backend, CacheEdge) and duplicate_headers:
edge_server.backend.counters.increment("duplicate_header_bypass")
if urlsplit(self.path).path.lstrip("/").partition("/")[0] in edge_server.mounts:
if resolve_mount(urlsplit(self.path).path, edge_server.mounts) is not None:
edge_server.backend.counters.increment("upstream_attempts")
outcome: Final = handle_edge_request(
selected_backend,
@ -1079,6 +1107,8 @@ def provider_edge_api_base(
return _shared_cache_edge(bind_host, advertise_host, forward_timeout).api_base(mount)
return None
case "record" | "replay":
if is_bedrock(mount):
return None
if mount not in EDGE_MOUNTS:
raise ValueError(f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(EDGE_MOUNTS))}")
return _shared_edge(mode, bundle_dir, bind_host, advertise_host, forward_timeout, match_profile()).api_base(
@ -1108,7 +1138,22 @@ def configured_cache_backend() -> CacheEdge | None:
return None
from provider_cache_redis import configured_cache
return configured_cache()
cache: Final = configured_cache()
return None if cache is None else replace(cache, policies=bedrock_policies())
@functools.lru_cache(maxsize=1)
def bedrock_policies() -> Mapping[str, MountPolicy]:
"""One policy per mounted Bedrock region, built lazily so a run that never
mounts Bedrock neither imports botocore nor resolves an AWS identity."""
from provider_edge_bedrock import bedrock_signer
return MappingProxyType(
{
f"bedrock/{region}": MountPolicy(sign=bedrock_signer(region), unkeyed_headers=SIGNATURE_HEADERS)
for region in BEDROCK_REGIONS
}
)
@functools.lru_cache(maxsize=8)

72
tests/e2e/provider_edge_bedrock.py Normal file
View file

@ -0,0 +1,72 @@
"""SigV4 re-signing for Bedrock traffic routed through the provider edge.
Bedrock is the one provider the edge could never mount. SigV4 signs the Host
header, so rewriting ``api_base`` to point at the edge invalidates the proxy's
signature and Bedrock rejects the call before it reaches a model. The edge
therefore has to drop the proxy's signature and mint its own over the upstream
URL it is actually about to call.
The identity it signs with is the run pod's own, from the EKS Pod Identity
association on ServiceAccount ``buildkite-e2e-run``. That role carries Bedrock
invoke and converse on an allowlist of the Anthropic models the suite registers
and nothing else, so a re-signed call can reach exactly the models the suite
already uses. The proxy's own Bedrock credentials are not involved in a routed
deployment, which is why ``aws_role_name`` deployments stay off the edge: their
whole point is to prove the product's assume-role chain.
Signature headers are excluded from the cache key by the caller, and they have
to be: ``x-amz-date`` is a timestamp, so keying on it would make every Bedrock
request a permanent miss.
"""
from __future__ import annotations
import functools
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Final
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import Credentials
from botocore.session import Session
from provider_cache import SIGNATURE_HEADERS
BEDROCK_SERVICE: Final = "bedrock"
class MissingAwsCredentials(RuntimeError):
"""No AWS identity is resolvable, so the edge cannot sign for Bedrock."""
@dataclass(frozen=True, slots=True)
class BedrockSigner:
region: str
credentials: Callable[[], Credentials]
def __call__(self, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> dict[str, str]:
unsigned: Final = {
name: value for name, value in headers.items() if name.lower() not in SIGNATURE_HEADERS
}
request: Final = AWSRequest(method=method, url=url, headers=unsigned, data=body or b"")
SigV4Auth(self.credentials(), BEDROCK_SERVICE, self.region).add_auth(request)
return dict(request.headers)
@functools.lru_cache(maxsize=1)
def pod_credentials() -> Credentials:
"""The run pod's own identity, resolved once per process through botocore's
ordinary chain, which reaches Pod Identity at the ``container-role`` link."""
resolved: Final = Session().get_credentials()
if resolved is None: # pyright: ignore[reportUnnecessaryComparison] # stubs miss the empty-chain None
raise MissingAwsCredentials(
"the provider edge is mounted for Bedrock but no AWS credentials resolve; "
"the run pod gets them from the Pod Identity association on buildkite-e2e-run"
)
return resolved
def bedrock_signer(region: str, credentials: Callable[[], Credentials] = pod_credentials) -> BedrockSigner:
"""Credentials are resolved on the first signed request, not here, so a run
that mounts Bedrock but never calls it needs no AWS identity at all."""
return BedrockSigner(region, credentials)

1
tests/e2e/pytest.ini
View file

@ -10,4 +10,5 @@ markers =
weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set
managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set
prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set
cli_determinism: drives the real claude CLI for several seconds, which widens the window in which another test's in-flight upstream call is attributed to it; deselected unless E2E_CLI_DETERMINISM is set
redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set

2
tests/e2e/quota_management/budgets/BUDGET_TEST_COVERAGE_MATRIX.md
View file

@ -21,7 +21,7 @@ on the shared lifecycle (every entity it creates is deleted on teardown).
| Entity | Unit | Pre-existing live | This suite (live) | Status |
|--------|------|-------------------|-------------------|--------|
| API key | `test_budget_reservation.py`, `test_max_budget_limiter.py` | `otel_tests` | `test_budget_enforcement_e2e::test_key_budget_blocks` | **covered** |
| API key | `test_budget_reservation.py` | `otel_tests` | `test_budget_enforcement_e2e::test_key_budget_blocks` | **covered** |
| Team | `test_team_budget_limits.py` | `otel_tests` | (org test builds a team) | **covered** |
| Internal user | auth unit tests | - | `test_internal_user_budget_blocks` | **covered (new)** |
| Team member | `test_team_member_budget.py` | - | `test_team_member_budget_blocks` | **covered (new)** |

19
tests/e2e/test_provider_edge.py
View file

@ -1279,15 +1279,30 @@ class TestApiBaseSeam:
)
def test_unknown_mount_raises_naming_the_known_mounts(self, tmp_path: Path) -> None:
with pytest.raises(ValueError, match="unknown provider mount 'bedrock'"):
with pytest.raises(ValueError, match="unknown provider mount 'cohere'"):
provider_edge_api_base(
"bedrock",
"cohere",
mode_raw="record",
bundle_dir=tmp_path / "bundle",
bind_host="127.0.0.1",
advertise_host="127.0.0.1",
)
@pytest.mark.parametrize("mode_raw", ["record", "replay"])
def test_bedrock_never_wires_a_bundle_because_the_edge_cannot_sign_into_one(
self, tmp_path: Path, mode_raw: str,
) -> None:
"""Record and replay serve from a bundle without re-signing, so a Bedrock
deployment pointed at that edge would send the proxy's signature over a
rewritten Host. It keeps its direct route in both modes."""
assert provider_edge_api_base(
"bedrock/us-east-1",
mode_raw=mode_raw,
bundle_dir=tmp_path / "bundle",
bind_host="127.0.0.1",
advertise_host="127.0.0.1",
) is None
def test_record_mode_boots_one_shared_edge_and_prepares_the_bundle(self, tmp_path: Path) -> None:
root = tmp_path / "bundle"
first = provider_edge_api_base(

7
tests/llm_translation/test_together_ai.py
View file

@ -3,6 +3,7 @@ Test TogetherAI LLM
"""
from base_llm_unit_tests import BaseLLMChatTest
from tests._live_test_helpers import cheapest_together_chat_model
import json
import os
from datetime import datetime
@ -16,7 +17,11 @@ import pytest
class TestTogetherAI(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
litellm.set_verbose = True
return {"model": "together_ai/openai/gpt-oss-20b"}
return {
"model": cheapest_together_chat_model(
function_calling=True, response_schema=True
)
}
def test_tool_call_no_arguments(self, tool_call_no_arguments):
"""Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833"""

74
tests/local_testing/test_completion.py
View file

@ -57,23 +57,6 @@ def test_response_model_none():
assert isinstance(x, litellm.ModelResponse)
def test_completion_custom_provider_model_name():
try:
litellm.cache = None
response = completion(
model="together_ai/openai/gpt-oss-20b",
messages=messages,
logger_fn=logger_fn,
)
# Add assertions here to check the-response
print(response)
print(response["choices"][0]["finish_reason"])
except litellm.Timeout as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def _openai_mock_response(*args, **kwargs) -> litellm.ModelResponse:
new_response = MagicMock()
new_response.headers = {"hello": "world"}
@ -2803,41 +2786,6 @@ def test_completion_together_ai_llama():
# test_completion_together_ai()
def test_customprompt_together_ai():
try:
litellm.set_verbose = False
litellm.num_retries = 0
print("in test_customprompt_together_ai")
print(litellm.success_callback)
print(litellm._async_success_callback)
response = completion(
model="together_ai/openai/gpt-oss-20b",
messages=messages,
roles={
"system": {
"pre_message": "<|im_start|>system\n",
"post_message": "<|im_end|>",
},
"assistant": {
"pre_message": "<|im_start|>assistant\n",
"post_message": "<|im_end|>",
},
"user": {
"pre_message": "<|im_start|>user\n",
"post_message": "<|im_end|>",
},
},
)
print(response)
except litellm.exceptions.Timeout as e:
print(f"Timeout Error")
pass
except Exception as e:
print(f"ERROR TYPE {type(e)}")
pytest.fail(f"Error occurred: {e}")
# test_customprompt_together_ai()
def response_format_tests(response: litellm.ModelResponse):
@ -3644,28 +3592,6 @@ async def test_acompletion_stream_watsonx():
# test_maritalk()
def test_completion_together_ai_stream():
litellm.set_verbose = True
user_message = "Write 1pg about YC & litellm"
messages = [{"content": user_message, "role": "user"}]
try:
response = completion(
model="together_ai/openai/gpt-oss-20b",
messages=messages,
stream=True,
max_tokens=5,
)
print(response)
for chunk in response:
print(chunk)
# print(string_response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_completion_together_ai_stream()
def test_moderation():
response = litellm.moderation(input="i'm ishaan cto of litellm")
print(response)

3
tests/local_testing/test_text_completion.py
View file

@ -12,6 +12,7 @@ from unittest.mock import MagicMock, patch
import pytest
import litellm
from tests._live_test_helpers import cheapest_together_chat_model
from litellm import (
RateLimitError,
TextCompletionResponse,
@ -4030,7 +4031,7 @@ def test_async_text_completion_together_ai():
async def test_get_response():
try:
response = await litellm.atext_completion(
model="together_ai/openai/gpt-oss-20b",
model=cheapest_together_chat_model(),
prompt="good morning",
max_tokens=10,
)

26
tests/logging_callback_tests/test_generic_api_callback.py
View file

@ -10,6 +10,7 @@ import httpx
import json
import logging
import time
from typing import Final
from unittest.mock import AsyncMock, patch
import pytest
@ -98,8 +99,15 @@ async def test_generic_api_callback():
assert isinstance(actual_request, list), "Request body should be a list"
assert len(actual_request) > 0, "Request body list should not be empty"
# Validate the first payload item
payload_item: StandardLoggingPayload = StandardLoggingPayload(**actual_request[0])
this_test_messages: Final = [{"role": "user", "content": "Hello, world!"}]
mine: Final = [
item for item in actual_request if item.get("messages") == this_test_messages
]
assert (
len(mine) == 1
), f"Expected this test's single call in the batch, got {len(mine)} of {len(actual_request)}"
payload_item: StandardLoggingPayload = StandardLoggingPayload(**mine[0])
print("##########\n")
print(json.dumps(payload_item, indent=4))
print("##########\n")
@ -448,11 +456,17 @@ async def test_generic_api_callback_sumologic_uses_ndjson():
assert isinstance(ndjson_data, str), "Data should be a string for NDJSON"
lines = ndjson_data.strip().split("\n")
assert len(lines) == 2, f"Expected 2 lines of NDJSON, got {len(lines)}"
records: Final = [json.loads(line) for line in lines]
# Each line should be valid JSON
for line in lines:
json.loads(line) # Will raise if invalid JSON
this_test_messages: Final = [
[{"role": "user", "content": f"Test {i}"}] for i in range(2)
]
mine: Final = [
record for record in records if record.get("messages") in this_test_messages
]
assert (
len(mine) == 2
), f"Expected this test's 2 calls as NDJSON lines, got {len(mine)} of {len(records)}"
@pytest.mark.asyncio

12
tests/logging_callback_tests/test_unit_test_litellm_logging.py
View file

@ -8,8 +8,8 @@ from typing import Literal
import pytest
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler
from litellm._service_logger import ServiceLogging
import asyncio
@ -58,11 +58,11 @@ def test_is_internal_litellm_proxy_callback():
"""
Ensure we can determine if a callback is an internal litellm proxy callback
eg. `_PROXY_MaxBudgetLimiter`, `_PROXY_CacheControlCheck`
eg. `_PROXY_MaxIterationsHandler`, `_PROXY_CacheControlCheck`
"""
logging = setup_logging()
assert logging._is_internal_litellm_proxy_callback(_PROXY_MaxBudgetLimiter) == True
assert logging._is_internal_litellm_proxy_callback(_PROXY_MaxIterationsHandler) == True
# Test non-internal callbacks
def regular_callback():
@ -95,7 +95,7 @@ def test_should_run_sync_callbacks_for_async_calls():
assert logging._should_run_sync_callbacks_for_async_calls() == True
# Test with internal callback only
litellm.success_callback = [_PROXY_MaxBudgetLimiter]
litellm.success_callback = [_PROXY_MaxIterationsHandler]
assert logging._should_run_sync_callbacks_for_async_calls() == False
@ -107,7 +107,7 @@ def test_remove_internal_litellm_callbacks():
callbacks = [
regular_callback,
_PROXY_MaxBudgetLimiter,
_PROXY_MaxIterationsHandler,
_PROXY_CacheControlCheck,
"string_callback",
]
@ -116,5 +116,5 @@ def test_remove_internal_litellm_callbacks():
assert len(filtered) == 2 # Should only keep regular_callback and string_callback
assert regular_callback in filtered
assert "string_callback" in filtered
assert _PROXY_MaxBudgetLimiter not in filtered
assert _PROXY_MaxIterationsHandler not in filtered
assert _PROXY_CacheControlCheck not in filtered

117
tests/test_litellm/integrations/otel/test_otel_v2_components.py
View file

@ -5,6 +5,7 @@ builders, and the registry validator's failure paths. Needs the OTel SDK."""
import json
import threading
from collections.abc import Iterator
from contextvars import Context as ContextVarContext
from dataclasses import replace
from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer
@ -15,6 +16,8 @@ pytest.importorskip("opentelemetry")
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( # noqa: E402
ExportTraceServiceRequest,
)
from opentelemetry import baggage # noqa: E402
from opentelemetry.context import attach, detach # noqa: E402
from opentelemetry.sdk.metrics import MeterProvider # noqa: E402
from opentelemetry.sdk.metrics.export import InMemoryMetricReader # noqa: E402
from opentelemetry.sdk.trace import TracerProvider # noqa: E402
@ -26,7 +29,10 @@ from opentelemetry.sdk.trace.export import ( # noqa: E402
from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E402
InMemorySpanExporter,
)
from opentelemetry.trace import SpanKind # noqa: E402
from opentelemetry.trace import SpanKind, get_current_span # noqa: E402
from opentelemetry.trace.propagation.tracecontext import ( # noqa: E402
TraceContextTextMapPropagator,
)
from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402
from litellm.integrations.otel.plumbing import providers # noqa: E402
@ -464,6 +470,115 @@ def test_extract_traceparent():
assert ctx_mod.extract_traceparent({"x": "y"}) is None
def _test_tracer():
exporter = InMemorySpanExporter()
provider = TracerProvider()
provider.add_span_processor(SimpleSpanProcessor(exporter))
return provider.get_tracer("test")
def test_inject_trace_context_prefers_request_root_span():
def run():
tracer = _test_tracer()
with tracer.start_as_current_span("root") as root:
ctx_mod.set_request_root_span(root)
result = ctx_mod.inject_trace_context(
{"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"}
)
propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
return result, root, propagated
result, root, propagated = ContextVarContext().run(run)
assert result["traceparent"] != "00-11111111111111111111111111111111-2222222222222222-01"
assert propagated.get_span_context().trace_id == root.get_span_context().trace_id
assert propagated.get_span_context().span_id == root.get_span_context().span_id
def test_inject_trace_context_uses_ambient_span_without_request_root():
def run():
tracer = _test_tracer()
with tracer.start_as_current_span("ambient") as ambient:
result = ctx_mod.inject_trace_context({})
propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
return ambient, propagated
ambient, propagated = ContextVarContext().run(run)
assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id
assert propagated.get_span_context().span_id == ambient.get_span_context().span_id
def test_inject_trace_context_replaces_stale_trace_headers():
def run():
tracer = _test_tracer()
with tracer.start_as_current_span("ambient") as ambient:
headers = {
"Traceparent": "00-" + "a" * 32 + "-" + "b" * 16 + "-01",
"Tracestate": "vendor=old",
"x-keep": "1",
}
result = ctx_mod.inject_trace_context(headers)
propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
return result, ambient, propagated
result, ambient, propagated = ContextVarContext().run(run)
assert sum(key.lower() == "traceparent" for key in result) == 1
assert not any(key.lower() == "tracestate" for key in result)
assert result["x-keep"] == "1"
assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id
def test_inject_trace_context_prefers_explicit_parent_span_over_root_and_ambient():
def run():
tracer = _test_tracer()
parent = tracer.start_span("litellm_request")
with tracer.start_as_current_span("ambient") as ambient:
ctx_mod.set_request_root_span(ambient)
result = ctx_mod.inject_trace_context({}, parent_span=parent)
propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
return parent, ambient, propagated
parent, ambient, propagated = ContextVarContext().run(run)
assert propagated.get_span_context().trace_id == parent.get_span_context().trace_id
assert propagated.get_span_context().span_id == parent.get_span_context().span_id
assert propagated.get_span_context().span_id != ambient.get_span_context().span_id
def test_inject_trace_context_skips_unusable_parent_span():
def run():
tracer = _test_tracer()
with tracer.start_as_current_span("ambient") as ambient:
result = ctx_mod.inject_trace_context({}, parent_span=object())
propagated = get_current_span(TraceContextTextMapPropagator().extract(result))
return ambient, propagated
ambient, propagated = ContextVarContext().run(run)
assert propagated.get_span_context().span_id == ambient.get_span_context().span_id
def test_inject_trace_context_returns_headers_unchanged_without_context():
headers = {"x-custom": "value"}
result = ContextVarContext().run(lambda: ctx_mod.inject_trace_context(headers))
assert result == headers
assert "traceparent" not in result
assert result is not headers
def test_inject_trace_context_does_not_forward_baggage():
def run():
tracer = _test_tracer()
with tracer.start_as_current_span("ambient"):
token = attach(baggage.set_baggage("litellm.team.id", "team"))
try:
return ctx_mod.inject_trace_context({})
finally:
detach(token)
result = ContextVarContext().run(run)
assert "baggage" not in result
def test_set_request_baggage_empty_returns_context():
assert ctx_mod.set_request_baggage({}) is not None

59
tests/test_litellm/integrations/test_opentelemetry.py
View file

@ -5424,6 +5424,65 @@ class TestGetSpanContextLitellmMetadataFallback(unittest.TestCase):
self.assertIsNone(detected_span)
class TestInboundTraceContextKeepsCallerTracestate(unittest.TestCase):
"""The request span built from inbound W3C headers must carry the caller's
tracestate so outbound propagation (passthrough) re-emits it instead of
dropping it alongside the stripped stale header."""
CALLER_TRACEPARENT = "00-" + "a" * 32 + "-" + "b" * 16 + "-01"
CALLER_TRACESTATE = "vendor=abc,other=xyz"
def _otel(self):
provider = TracerProvider()
provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter()))
otel = OpenTelemetry()
otel.tracer = provider.get_tracer(__name__)
return otel
def test_request_span_propagates_caller_tracestate_downstream(self):
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from litellm.integrations.otel.plumbing.context import inject_trace_context
inbound = {"traceparent": self.CALLER_TRACEPARENT, "tracestate": self.CALLER_TRACESTATE}
span = self._otel().create_litellm_proxy_request_started_span(
start_time=datetime.now(timezone.utc), headers=inbound
)
outbound = inject_trace_context(inbound, parent_span=span)
span.end()
propagated = trace.get_current_span(TraceContextTextMapPropagator().extract(outbound)).get_span_context()
self.assertEqual(outbound["tracestate"], self.CALLER_TRACESTATE)
self.assertEqual(propagated.trace_id, span.get_span_context().trace_id)
self.assertEqual(propagated.span_id, span.get_span_context().span_id)
self.assertNotEqual(outbound["traceparent"], self.CALLER_TRACEPARENT)
def test_request_span_without_caller_tracestate_emits_none(self):
from litellm.integrations.otel.plumbing.context import inject_trace_context
inbound = {"traceparent": self.CALLER_TRACEPARENT}
span = self._otel().create_litellm_proxy_request_started_span(
start_time=datetime.now(timezone.utc), headers=inbound
)
outbound = inject_trace_context(inbound, parent_span=span)
span.end()
self.assertNotIn("tracestate", outbound)
self.assertNotEqual(outbound["traceparent"], self.CALLER_TRACEPARENT)
def test_span_context_from_header_keeps_caller_tracestate(self):
kwargs = {
"litellm_params": {
"proxy_server_request": {
"headers": {"traceparent": self.CALLER_TRACEPARENT, "tracestate": self.CALLER_TRACESTATE}
}
}
}
ctx, detected_span = self._otel()._get_span_context(kwargs)
self.assertIsNone(detected_span)
self.assertEqual(trace.get_current_span(ctx).get_span_context().trace_state.to_header(), self.CALLER_TRACESTATE)
class TestEndProxySpanLitellmMetadataFallback(unittest.TestCase):
"""
Tests for _end_proxy_span_from_kwargs() falling back to litellm_metadata.

49
tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py
View file

@ -1108,3 +1108,52 @@ def test_get_optional_params_preserves_max_for_declared_levels_model():
)
assert optional_params["reasoning_effort"] == "max"
def _together_chat_transport() -> tuple[HTTPHandler, list[httpx.Request]]:
captured_requests: list[httpx.Request] = []
def respond(request: httpx.Request) -> httpx.Response:
captured_requests.append(request)
return httpx.Response(
200,
json={
"id": "chatcmpl-together",
"object": "chat.completion",
"created": 1234567890,
"model": TOOL_CALLING_MODEL,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Hello!"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
},
)
client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond)))
return client, captured_requests
def test_custom_role_wrappers_never_reach_the_request():
client, captured_requests = _together_chat_transport()
messages = [{"role": "user", "content": "Hello!"}]
litellm.completion(
model=f"together_ai/{TOOL_CALLING_MODEL}",
messages=messages,
roles={
"system": {"pre_message": "<|im_start|>system\n", "post_message": "<|im_end|>"},
"assistant": {"pre_message": "<|im_start|>assistant\n", "post_message": "<|im_end|>"},
"user": {"pre_message": "<|im_start|>user\n", "post_message": "<|im_end|>"},
},
api_key="fake-key",
client=client,
)
request_body = json.loads(captured_requests[0].content)
assert request_body["messages"] == messages
assert "prompt" not in request_body
assert "roles" not in request_body

59
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -5605,6 +5605,65 @@ async def test_common_checks_personal_user_budget_blocks_in_gather():
assert "User=u1" in str(over.value)
async def _common_checks_for_over_budget_personal_key(*, model: str) -> bool:
from litellm import Router
from litellm.proxy.auth.auth_checks import _is_model_cost_zero, common_checks
llm_router: Final = Router(
model_list=[
{
"model_name": "free-model",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
"model_info": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0},
},
{
"model_name": "paid-model",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
},
]
)
user: Final = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=1.0)
token: Final = UserAPIKeyAuth(token="k1", user_id="u1")
async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs):
return 5.0 if counter_key == "spend:user:u1" else 0.0
proxy_logging_obj: Final = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
with (
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter),
):
result: Final = await common_checks(
request_body={"model": model, "messages": [{"role": "user", "content": "hi"}]},
team_object=None,
user_object=user,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route="/chat/completions",
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=token,
request=MagicMock(spec=Request),
skip_budget_checks=_is_model_cost_zero(model=model, llm_router=llm_router),
)
await asyncio.sleep(0)
return result
@pytest.mark.asyncio
async def test_common_checks_over_budget_user_can_still_call_zero_cost_model():
"""LIT-7464: an exhausted personal budget must not block a model priced at 0/0,
while the same user is still rejected on a priced model."""
assert await _common_checks_for_over_budget_personal_key(model="free-model") is True
with pytest.raises(litellm.BudgetExceededError) as over:
await _common_checks_for_over_budget_personal_key(model="paid-model")
assert "ExceededBudget: User=u1" in str(over.value)
async def _run_internal_user_budget_alert(
*,
spend: float,

237
tests/test_litellm/proxy/hooks/test_max_budget_limiter.py
View file

@ -1,237 +0,0 @@
"""
Unit tests for the personal-budget pre-call hook.
The reservation path (added in PR #26845) atomically pre-fills the same
`spend:user:{user_id}` counter this hook reads, admitting at a strict-`<`
boundary. Re-checking with `>=` after reservation would reject requests the
reservation already admitted when the reservation fills the counter to
exactly `max_budget` (e.g. requests with no `max_tokens` cap fall back to
reserving the smallest remaining headroom).
These tests pin the skip-when-reserved behavior and guard against drift.
"""
from unittest.mock import AsyncMock, patch
import pytest
from fastapi import HTTPException
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
def _make_user_api_key_auth(
user_id: str = "user-1",
user_max_budget: float = 10.0,
user_spend: float = 0.0,
team_id=None,
budget_reservation=None,
) -> UserAPIKeyAuth:
return UserAPIKeyAuth(
api_key="sk-test",
user_id=user_id,
user_max_budget=user_max_budget,
user_spend=user_spend,
team_id=team_id,
budget_reservation=budget_reservation,
)
@pytest.mark.asyncio
async def test_under_budget_passes():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=3.0),
):
result = await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert result is None
@pytest.mark.asyncio
async def test_over_budget_rejects_without_reservation():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.status_code == 429
assert "Max budget limit reached." in exc_info.value.detail
@pytest.mark.asyncio
async def test_skips_when_user_counter_is_reserved():
"""
Reservation atomically pre-fills `spend:user:{user_id}` and admits the
request. The legacy `>=` check must not double-enforce on the same
counter — that's what produced the boundary regression where a fresh
user with no `max_tokens` cap got 429'd on their first request.
"""
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(
user_id="user-1",
user_max_budget=10.0,
budget_reservation={
"reserved_cost": 10.0,
"entries": [
{
"counter_key": "spend:user:user-1",
"entity_type": "User",
"entity_id": "user-1",
"reserved_cost": 10.0,
"applied_adjustment": 0.0,
}
],
"finalized": False,
},
)
# `get_current_spend` would return 10.0 here (counter pre-filled by the
# reservation). The hook must skip without reading it.
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
) as mock_get_spend:
result = await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert result is None
mock_get_spend.assert_not_awaited()
@pytest.mark.asyncio
async def test_does_not_skip_when_reservation_covers_a_different_counter():
"""
A reservation that only covers e.g. `spend:team:{team_id}` (not the user
counter) must not exempt the user-budget check.
"""
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(
user_id="user-1",
user_max_budget=10.0,
budget_reservation={
"reserved_cost": 5.0,
"entries": [
{
"counter_key": "spend:team:team-x",
"entity_type": "Team",
"entity_id": "team-x",
"reserved_cost": 5.0,
"applied_adjustment": 0.0,
}
],
"finalized": False,
},
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.status_code == 429
@pytest.mark.asyncio
async def test_team_keys_skip_personal_budget():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(
user_max_budget=10.0,
team_id="team-1",
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=999.0),
) as mock_get_spend:
result = await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert result is None
mock_get_spend.assert_not_awaited()
@pytest.mark.asyncio
async def test_team_keys_enforce_personal_budget_when_flag_enabled():
"""This hook is the third personal-budget gate alongside common_checks and the
reservation path, so apply_user_budget_to_team_keys has to reach it too or an
opted-in deployment enforces in two places out of three."""
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = _make_user_api_key_auth(
user_max_budget=10.0,
team_id="team-1",
)
with patch.dict(
"litellm.proxy.proxy_server.general_settings",
{"apply_user_budget_to_team_keys": True},
), patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=999.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.status_code == 429
@pytest.mark.asyncio
async def test_no_max_budget_passes():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-test",
user_id="user-1",
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=999.0),
) as mock_get_spend:
result = await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert result is None
mock_get_spend.assert_not_awaited()

61
tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py
View file

@ -6,7 +6,7 @@ Background
----------
The proxy's internal rate-limit hooks (parallel_request_limiter,
parallel_request_limiter_v3, dynamic_rate_limiter, dynamic_rate_limiter_v3,
batch_rate_limiter, max_budget_limiter, max_iterations_limiter,
batch_rate_limiter, max_iterations_limiter,
max_budget_per_session_limiter) all fire from ``async_pre_call_hook`` —
*before* :func:`litellm.get_llm_provider` runs anywhere else in the request
lifecycle.
@ -50,7 +50,6 @@ from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHand
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
_PROXY_DynamicRateLimitHandlerV3,
)
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.max_budget_per_session_limiter import (
_PROXY_MaxBudgetPerSessionHandler,
)
@ -830,64 +829,6 @@ async def test_batch_rate_limiter_unknown_model_falls_back():
assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
# ---------------------------------------------------------------------------
# max_budget_limiter
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_max_budget_limiter_populates_provider():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-budget",
user_id="user-1",
user_max_budget=10.0,
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={"model": "gpt-4o-mini"},
call_type="completion",
)
exc = exc_info.value
assert exc.status_code == 429
assert isinstance(exc, RateLimitError)
assert exc.llm_provider == "openai"
assert exc.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_max_budget_limiter_no_model_falls_back():
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-budget",
user_id="user-1",
user_max_budget=10.0,
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=10.0),
):
with pytest.raises(HTTPException) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK
assert exc_info.value.model == ""
# ---------------------------------------------------------------------------
# max_iterations_limiter
# ---------------------------------------------------------------------------

122
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -2,6 +2,7 @@ import asyncio
import json
import logging
import os
import sys
from collections.abc import Callable
from contextlib import ExitStack, contextmanager
from io import BytesIO
@ -29,6 +30,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
resolve_pass_through_request_timeout,
resolve_llm_passthrough_timeout,
websocket_passthrough_request,
_with_trace_context,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -46,6 +48,15 @@ import litellm
MESSAGE_START_SSE_FRAME = b'event: message_start\ndata: {"type": "message_start"}\n\n'
def test_with_trace_context_without_opentelemetry(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setitem(sys.modules, "litellm.integrations.otel.plumbing.context", None)
headers = _with_trace_context({"authorization": "x"}, parent_span=None)
assert headers == {"authorization": "x"}
assert "traceparent" not in headers
# Test is_multipart
def test_is_multipart():
# Test with multipart content type
@ -4270,6 +4281,47 @@ def _relay_client_request(method="GET"):
return mock_request
@pytest.mark.asyncio
@pytest.mark.parametrize("span_source", ["auth_parent_span", "ambient_span"])
async def test_pass_through_request_propagates_active_trace_context(span_source: str):
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace import get_current_span
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
captured: dict[str, httpx.Headers] = {}
def transport_handler(upstream_request: httpx.Request) -> httpx.Response:
captured["headers"] = upstream_request.headers
return httpx.Response(200, json={"ok": True}, request=upstream_request)
fake_client, cleanup = _inject_fake_passthrough_client(httpx.MockTransport(transport_handler), timeout=None)
tracer = TracerProvider().get_tracer("test")
try:
with ExitStack() as stack:
_enter_relay_logging_mocks(stack, {})
if span_source == "auth_parent_span":
span = tracer.start_span("litellm_request")
stack.callback(span.end)
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", parent_otel_span=span)
else:
span = stack.enter_context(tracer.start_as_current_span("passthrough"))
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test")
response = await pass_through_request(
request=_relay_client_request(method="POST"),
target="http://internal-api.test/v1/generate",
custom_headers={},
user_api_key_dict=user_api_key_dict,
)
finally:
cleanup()
await fake_client.aclose()
assert response.status_code == 200
propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"]))
assert propagated.get_span_context().trace_id == span.get_span_context().trace_id
assert propagated.get_span_context().span_id == span.get_span_context().span_id
@pytest.mark.asyncio
async def test_pass_through_request_relays_non_json_body_without_buffering():
"""
@ -4866,6 +4918,76 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame():
assert all(call.kwargs.get("code") != 1011 for call in websocket.close.await_args_list)
@pytest.mark.asyncio
@pytest.mark.parametrize("forward_headers", [True, False])
@pytest.mark.parametrize("span_source", ["auth_parent_span", "ambient_span"])
async def test_websocket_passthrough_propagates_active_trace_context(
monkeypatch, forward_headers: bool, span_source: str
):
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace import get_current_span
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from starlette.websockets import WebSocketState
captured: dict[str, dict[str, str]] = {}
upstream_ws = FakeUpstreamWebSocket(b"{}")
def fake_connect(target, additional_headers):
captured["headers"] = additional_headers
return FakeUpstreamConnect(upstream_ws)
websocket = MagicMock()
websocket.accept = AsyncMock()
websocket.send_text = AsyncMock()
websocket.send_bytes = AsyncMock()
websocket.receive = AsyncMock(return_value={"type": "websocket.disconnect"})
websocket.close = AsyncMock()
websocket.headers = {"authorization": "Bearer client"}
websocket.client_state = WebSocketState.CONNECTED
websocket.application_state = WebSocketState.CONNECTED
tracer = TracerProvider().get_tracer("test")
mock_proxy_logging = MagicMock()
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
mock_proxy_logging.post_call_success_hook = AsyncMock()
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_worker = MagicMock()
mock_worker.ensure_initialized_and_enqueue = MagicMock(
side_effect=lambda async_coroutine: async_coroutine.close()
)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)
monkeypatch.setattr(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect",
fake_connect,
)
monkeypatch.setattr(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER",
mock_worker,
)
with ExitStack() as stack:
if span_source == "auth_parent_span":
span = tracer.start_span("litellm_request")
stack.callback(span.end)
user_api_key_dict = UserAPIKeyAuth(parent_otel_span=span)
else:
span = stack.enter_context(tracer.start_as_current_span("websocket_passthrough"))
user_api_key_dict = UserAPIKeyAuth()
await websocket_passthrough_request(
websocket=websocket,
target="wss://upstream.example.test/v1/realtime",
custom_headers={},
user_api_key_dict=user_api_key_dict,
forward_headers=forward_headers,
endpoint="/realtime",
accept_websocket=True,
)
propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"]))
assert propagated.get_span_context().trace_id == span.get_span_context().trace_id
assert propagated.get_span_context().span_id == span.get_span_context().span_id
assert captured["headers"].get("authorization") == ("Bearer client" if forward_headers else None)
class ClosingUpstreamWebSocket:
def __init__(self, close_exc: Exception):
self._close_exc = close_exc

4
tests/test_litellm/proxy/proxy_server/test_routes_config.py
View file

@ -1397,7 +1397,7 @@ def test_get_config_callbacks_excludes_internal_runtime_callbacks(client, auth_a
from litellm.integrations.s3_v2 import S3Logger
from litellm.integrations.sqs import SQSLogger
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import VectorStorePreCallHook
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
from litellm.router import Router
class _InventoryTestGuardrail(CustomGuardrail):
@ -1425,7 +1425,7 @@ def test_get_config_callbacks_excludes_internal_runtime_callbacks(client, auth_a
litellm,
"callbacks",
[
_PROXY_MaxBudgetLimiter(),
_PROXY_CacheControlCheck(),
_PROXY_LiteLLMManagedFiles(internal_usage_cache=MagicMock(), prisma_client=MagicMock()),
ServiceLogging(),
VectorStorePreCallHook(),

2
tests/test_litellm/proxy/test_proxy_logging_hook_detection.py
View file

@ -73,7 +73,7 @@ async def test_post_call_response_headers_hook_returns_early_without_callbacks(
def test_callback_capabilities_skips_default_custom_logger(monkeypatch):
"""
Internal proxy hooks (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit
Internal proxy hooks (e.g. _PROXY_CacheControlCheck, ManagedFiles) inherit
the default ``async_post_call_streaming_iterator_hook`` body. The
capability scanner must NOT report them as iterator overrides — wrapping
the chunk stream through every no-op layer was responsible for ~10x

10
tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py
View file

@ -220,7 +220,7 @@ def test_add_proxy_hooks_registers_callbacks(proxy_logging, monkeypatch):
what gets registered. Verifies that the resulting instances land in
``proxy_logging.proxy_hook_mapping`` keyed by hook name.
"""
hook_keys = ["cache_control_check", "max_budget_limiter"]
hook_keys = ["cache_control_check", "max_iterations_limiter"]
registered: List[Any] = []
from litellm.proxy import utils as utils_mod
@ -362,22 +362,22 @@ def test_add_proxy_hooks_unknown_hook_raises(proxy_logging, monkeypatch):
def test_get_proxy_hook_returns_registered_instance(proxy_logging):
s_cache = MagicMock()
s_budget = MagicMock()
s_iterations = MagicMock()
s_parallel = MagicMock()
proxy_logging.proxy_hook_mapping = {
"cache_control_check": s_cache,
"max_budget_limiter": s_budget,
"max_iterations_limiter": s_iterations,
"max_parallel_request_limiter": s_parallel,
}
snapshot = {
"cache_control_check": proxy_logging.get_proxy_hook("cache_control_check") is s_cache,
"max_budget_limiter": proxy_logging.get_proxy_hook("max_budget_limiter") is s_budget,
"max_iterations_limiter": proxy_logging.get_proxy_hook("max_iterations_limiter") is s_iterations,
"max_parallel_request_limiter": proxy_logging.get_proxy_hook("max_parallel_request_limiter") is s_parallel,
"unknown_returns_none": proxy_logging.get_proxy_hook("unknown") is None,
}
assert snapshot == {
"cache_control_check": True,
"max_budget_limiter": True,
"max_iterations_limiter": True,
"max_parallel_request_limiter": True,
"unknown_returns_none": True,
}

34
tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py
View file

@ -3,7 +3,7 @@
from __future__ import annotations
from typing import Any, Dict
from unittest.mock import AsyncMock, MagicMock
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
@ -400,6 +400,37 @@ def test_has_pre_call_guardrails_counts_a_content_enforcer(proxy_logging, monkey
assert proxy_logging.has_pre_call_guardrails({}) is True
@pytest.mark.asyncio
async def test_registered_hooks_do_not_enforce_user_budget(proxy_logging, monkeypatch):
"""
Personal budget is auth's job (`_user_max_budget_check`), which exempts
zero-cost models. A hook re-checking the same counter without that
exemption is what 429'd free models once a user was over budget.
"""
monkeypatch.setattr(litellm, "callbacks", [])
with patch("litellm.proxy.proxy_server.prisma_client", None):
proxy_logging._add_proxy_hooks(llm_router=None)
ProxyLogging._callback_capabilities_cache.clear()
over_budget_user = UserAPIKeyAuth(
api_key="sk-personal",
user_id="user-over-budget",
user_max_budget=1.0,
user_spend=5.0,
team_id=None,
)
data = {"model": "free-model", "messages": [{"role": "user", "content": "hi"}]}
with patch("litellm.proxy.proxy_server.get_current_spend", new=AsyncMock(return_value=5.0)):
out = await proxy_logging.pre_call_hook(
user_api_key_dict=over_budget_user,
data=data,
call_type="completion",
)
assert out == data
def test_every_pre_call_customlogger_is_deliberately_classified():
"""
A ledger, so a new hook cannot land unclassified.
@ -415,7 +446,6 @@ def test_every_pre_call_customlogger_is_deliberately_classified():
"_ENTERPRISE_BlockedUserList",
}
counts_or_shapes_the_request = {
"_PROXY_MaxBudgetLimiter",
"_PROXY_MaxParallelRequestsHandler_v3",
"_PROXY_MaxIterationsHandler",
"_PROXY_MaxBudgetPerSessionHandler",

69
tests/test_litellm/test_rate_limit_error_unification.py
View file

@ -221,28 +221,6 @@ class TestProxyHookCategoryWiring:
"""End-to-end check that every proxy-side rate limiter raises the unified
class with a sensible category, not a bare HTTPException."""
def test_max_budget_limiter_raises_proxy_rate_limit_error(self):
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
limiter = _PROXY_MaxBudgetLimiter()
# The simplest deterministic path: directly raise from the conditional
# branch by calling into the helper's exception construction. We
# round-trip through the public class to assert the shape.
with pytest.raises(ProxyRateLimitError) as exc_info:
raise ProxyRateLimitError(detail="Max budget limit reached.")
assert exc_info.value.status_code == 429
assert exc_info.value.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT
# And it's also a RateLimitError + HTTPException (the unification).
assert isinstance(exc_info.value, RateLimitError)
assert isinstance(exc_info.value, HTTPException)
# Static check that the limiter's module imports the unified class so
# the source of truth is wired correctly.
from litellm.proxy.hooks import max_budget_limiter
assert hasattr(max_budget_limiter, "ProxyRateLimitError")
assert max_budget_limiter.ProxyRateLimitError is ProxyRateLimitError
del limiter # silence unused-var
@pytest.mark.parametrize(
"module_path",
[
@ -251,7 +229,6 @@ class TestProxyHookCategoryWiring:
"litellm.proxy.hooks.dynamic_rate_limiter",
"litellm.proxy.hooks.dynamic_rate_limiter_v3",
"litellm.proxy.hooks.batch_rate_limiter",
"litellm.proxy.hooks.max_budget_limiter",
"litellm.proxy.hooks.max_budget_per_session_limiter",
"litellm.proxy.hooks.max_iterations_limiter",
],
@ -542,44 +519,6 @@ class TestProxyHooksActuallyRaiseProxyRateLimitError:
assert isinstance(e, RateLimitError)
assert isinstance(e, HTTPException)
@pytest.mark.asyncio
async def test_max_budget_limiter_raises_proxy_rate_limit_error(self):
"""
Drive `_PROXY_MaxBudgetLimiter` past the user budget and assert it
raises the unified class. Mocks `get_current_spend` so we don't need
the proxy DB.
"""
from unittest.mock import patch
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.max_budget_limiter import (
_PROXY_MaxBudgetLimiter,
)
handler = _PROXY_MaxBudgetLimiter()
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-test-budget",
user_id="user-budget-1",
user_max_budget=1.0,
user_spend=2.0,
)
with patch(
"litellm.proxy.proxy_server.get_current_spend",
return_value=5.0,
):
with pytest.raises(ProxyRateLimitError) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data={},
call_type="completion",
)
e = exc_info.value
assert e.status_code == 429
assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT
assert "max budget" in str(e.detail).lower()
@pytest.mark.asyncio
async def test_dynamic_rate_limiter_v1_raises_proxy_rate_limit_error(self):
"""
@ -1156,14 +1095,6 @@ class TestProxyHooksWireTypeCorrectly:
max-iterations) without grepping the error message.
"""
def test_max_budget_limiter_emits_budget_type(self):
e = ProxyRateLimitError(
detail="Max budget limit reached.",
rate_limit_type=RateLimitType.BUDGET,
)
assert e.category == "litellm_rate_limit"
assert e.rate_limit_type == "budget"
def test_max_iterations_limiter_emits_max_iterations_type(self):
e = ProxyRateLimitError(
detail="Max iterations exceeded for session abc.",

82
ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx
View file

@ -1,82 +0,0 @@
/* @vitest-environment jsdom */
import { renderHook } from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";
const { mockPush, navState } = vi.hoisted(() => ({
mockPush: vi.fn(),
navState: { pathname: "/logs" },
}));
vi.mock("next/navigation", () => ({
usePathname: () => navState.pathname,
useRouter: () => ({ push: mockPush }),
}));
vi.mock("@/components/networking", () => ({ serverRootPath: "" }));
import { createTabRoutes } from "@/utils/tabRoutes";
import { useTabRouting } from "./useTabRouting";
const routes = createTabRoutes("logs", ["audit", "deleted-keys", "deleted-teams"] as const);
const render = (ready = true) => {
const config = {
routes,
baseTabKey: "request-logs",
visibleKeys: ["audit", "deleted-keys", "deleted-teams"],
ready,
};
return renderHook(() => useTabRouting(config));
};
describe("useTabRouting", () => {
beforeEach(() => {
navState.pathname = "/logs";
mockPush.mockClear();
});
it("maps the base path to the base tab key", () => {
const { result } = render();
expect(result.current.activeSlug).toBe("");
expect(result.current.activeKey).toBe("request-logs");
});
it("uses the slug itself as the active key for a known nested tab", () => {
navState.pathname = "/ui/logs/audit";
const { result } = render();
expect(result.current.activeKey).toBe("audit");
});
it("falls back to the base tab key for an unknown slug", () => {
navState.pathname = "/ui/logs/bogus";
const { result } = render();
expect(result.current.activeKey).toBe("request-logs");
});
it("redirects an unknown slug to the base href once ready", () => {
const replaceMock = vi.fn();
const originalLocation = window.location;
Object.defineProperty(window, "location", { configurable: true, value: { replace: replaceMock } });
navState.pathname = "/ui/logs/bogus";
render(true);
expect(replaceMock).toHaveBeenCalledWith("/ui/logs/");
Object.defineProperty(window, "location", { configurable: true, value: originalLocation });
});
it("does not redirect while not ready (role/creds still loading)", () => {
const replaceMock = vi.fn();
const originalLocation = window.location;
Object.defineProperty(window, "location", { configurable: true, value: { replace: replaceMock } });
navState.pathname = "/ui/logs/bogus";
render(false);
expect(replaceMock).not.toHaveBeenCalled();
Object.defineProperty(window, "location", { configurable: true, value: originalLocation });
});
it("pushes the tab href on change, mapping the base key back to the empty slug", () => {
const { result } = render();
result.current.onTabChange("audit");
expect(mockPush).toHaveBeenCalledWith("/ui/logs/audit/");
result.current.onTabChange("request-logs");
expect(mockPush).toHaveBeenCalledWith("/ui/logs/");
});
});

38
ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts
View file

@ -1,38 +0,0 @@
import { useEffect } from "react";
import { usePathname, useRouter } from "next/navigation";
import type { TabRoutes } from "@/utils/tabRoutes";
interface UseTabRoutingArgs {
routes: Pick<TabRoutes<string>, "tabHref" | "slugFromPathname">;
baseTabKey: string;
visibleKeys: readonly string[];
ready?: boolean;
}
interface TabRoutingState {
activeSlug: string;
activeKey: string;
onTabChange: (key: string) => void;
}
export function useTabRouting({ routes, baseTabKey, visibleKeys, ready = true }: UseTabRoutingArgs): TabRoutingState {
const { tabHref, slugFromPathname } = routes;
const pathname = usePathname();
const router = useRouter();
const activeSlug = slugFromPathname(pathname);
const isKnownSlug = activeSlug === "" || visibleKeys.includes(activeSlug);
const activeKey = isKnownSlug ? activeSlug || baseTabKey : baseTabKey;
useEffect(() => {
if (ready && activeSlug !== "" && !isKnownSlug) {
window.location.replace(tabHref(""));
}
}, [ready, activeSlug, isKnownSlug, tabHref]);
const onTabChange = (key: string) => {
router.push(tabHref(key === baseTabKey ? "" : key));
};
return { activeSlug, activeKey, onTabChange };
}

53
ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx
View file

@ -1,5 +1,8 @@
import { render, screen } from "@testing-library/react";
import { screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import type { OnUrlUpdateFunction } from "nuqs/adapters/testing";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { renderWithProviders } from "../../../../tests/test-utils";
import PlaygroundPage from "./page";
const authState = { userRole: "Admin" };
@ -35,14 +38,17 @@ vi.mock("@/app/(dashboard)/playground/components/chat_ui/AgentBuilderView", () =
default: () => <div data-testid="agent-builder" />,
}));
describe("PlaygroundPage role guard", () => {
beforeEach(() => {
authState.userRole = "Admin";
});
const lastUrlUpdate = (onUrlUpdate: ReturnType<typeof vi.fn<OnUrlUpdateFunction>>) =>
onUrlUpdate.mock.calls.at(-1)?.[0];
beforeEach(() => {
authState.userRole = "Admin";
});
describe("PlaygroundPage role guard", () => {
it.each(["Internal Viewer", "Admin Viewer"])("blocks the entire playground for %s", (role) => {
authState.userRole = role;
render(<PlaygroundPage />);
renderWithProviders(<PlaygroundPage />);
expect(screen.getByText("Access Denied")).toBeInTheDocument();
expect(screen.queryByRole("tab")).not.toBeInTheDocument();
@ -54,10 +60,43 @@ describe("PlaygroundPage role guard", () => {
it.each(["Admin", "Internal User", "Org Admin"])("renders the playground for %s", (role) => {
authState.userRole = role;
render(<PlaygroundPage />);
renderWithProviders(<PlaygroundPage />);
expect(screen.queryByText("Access Denied")).not.toBeInTheDocument();
expect(screen.getByRole("tab", { name: "Chat" })).toBeInTheDocument();
expect(screen.getByTestId("chat-ui")).toBeInTheDocument();
});
});
describe("PlaygroundPage ?tab= deep link", () => {
it("opens on Chat when the URL has no tab", () => {
renderWithProviders(<PlaygroundPage />);
expect(screen.getByRole("tab", { name: "Chat" })).toHaveAttribute("aria-selected", "true");
});
it("activates the tab named in ?tab=", () => {
renderWithProviders(<PlaygroundPage />, { searchParams: { tab: "compare" } });
expect(screen.getByRole("tab", { name: "Compare" })).toHaveAttribute("aria-selected", "true");
expect(screen.getByRole("tab", { name: "Chat" })).toHaveAttribute("aria-selected", "false");
});
it("falls back to Chat when ?tab= is not a playground tab", () => {
renderWithProviders(<PlaygroundPage />, { searchParams: { tab: "settings" } });
expect(screen.getByRole("tab", { name: "Chat" })).toHaveAttribute("aria-selected", "true");
});
it("clicking a tab writes ?tab= with history replace", async () => {
const user = userEvent.setup();
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<PlaygroundPage />, { onUrlUpdate });
await user.click(screen.getByRole("tab", { name: "Compliance" }));
expect(await screen.findByRole("tab", { name: "Compliance", selected: true })).toBeInTheDocument();
await waitFor(() => expect(lastUrlUpdate(onUrlUpdate)?.searchParams.get("tab")).toBe("compliance"));
expect(lastUrlUpdate(onUrlUpdate)?.options.history).toBe("replace");
});
});

10
ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx
View file

@ -9,6 +9,9 @@ import { DeprecationBanner } from "@/components/DeprecationBanner";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import { fetchProxySettings } from "@/utils/proxyUtils";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { useUrlTab } from "@/hooks/useUrlTab";
const PLAYGROUND_TABS = ["chat", "compare", "compliance", "agent-builder"] as const;
interface ProxySettings {
PROXY_BASE_URL?: string;
@ -18,6 +21,7 @@ interface ProxySettings {
export default function PlaygroundPage() {
const { accessToken, userRole, userId, disabledPersonalKeyCreation, token, isViewOnly } = useAuthorized();
const [proxySettings, setProxySettings] = useState<ProxySettings | undefined>(undefined);
const [activeTab, setActiveTab] = useUrlTab(PLAYGROUND_TABS, "chat");
useEffect(() => {
const initializeProxySettings = async () => {
@ -48,7 +52,11 @@ export default function PlaygroundPage() {
return (
<div className="flex h-full min-h-0 w-full min-w-0 flex-col overflow-hidden">
<Tabs defaultValue="chat" className="flex min-h-0 min-w-0 flex-1 flex-col gap-0 overflow-hidden">
<Tabs
value={activeTab}
onValueChange={setActiveTab}
className="flex min-h-0 min-w-0 flex-1 flex-col gap-0 overflow-hidden"
>
<TabsList variant="line" className="w-full shrink-0 justify-start overflow-x-auto pb-1">
<TabsTrigger value="chat" className="flex-none">
Chat

152
ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx
View file

@ -4,7 +4,7 @@ import type { OnUrlUpdateFunction } from "nuqs/adapters/testing";
import { vi, it, expect, beforeEach, describe, Mock, MockedFunction } from "vitest";
import { chooseSelectOption, renderWithProviders } from "../../../tests/test-utils";
import { VirtualKeysTable } from "./VirtualKeysTable";
import { KEY_TABLE_SORT_FIELDS } from "./keyTableColumns";
import { KEY_TABLE_HIDDEN_COLUMNS, KEY_TABLE_SORT_FIELDS } from "./keyTableColumns";
import { KeyResponse, Team } from "../key_team_helpers/key_list";
import { useKeyInfo } from "@/app/(dashboard)/hooks/keys/useKeyInfo";
import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
@ -187,6 +187,7 @@ const lastHistoryMode = (onUrlUpdate: Mock<OnUrlUpdateFunction>) => onUrlUpdate.
beforeEach(() => {
vi.clearAllMocks();
localStorage.clear();
mockUseKeys.mockReturnValue(keysResult([mockKey]));
mockUseKeyInfo.mockReturnValue(keyInfoResult(undefined));
@ -823,16 +824,27 @@ describe("table state lives in the URL so it survives leaving and returning to t
});
it("restores the drawer filters from the URL on mount", async () => {
renderWithProviders(<VirtualKeysTable />, { searchParams: { filter_team: "team-1", filter_user: "user-42" } });
const searchParams = {
filter_team: "team-1",
filter_org: "org-1",
filter_user: "user-42",
filter_key_id: mockKey.token,
};
const expectedKeyListOptions = {
teamID: "team-1",
organizationID: "org-1",
userID: "user-42",
keyHash: mockKey.token,
};
renderWithProviders(<VirtualKeysTable />, { searchParams });
await waitFor(() => {
expect(mockUseKeys).toHaveBeenLastCalledWith(
1,
50,
expect.objectContaining({ teamID: "team-1", userID: "user-42" }),
);
expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining(expectedKeyListOptions));
});
expect(screen.getByTestId("filter-chip-team_id")).toHaveTextContent("Test Team");
expect(screen.getByTestId("filter-chip-org_id")).toHaveTextContent("Test Organization");
expect(screen.getByTestId("filter-chip-user_id")).toHaveTextContent("user-42");
expect(screen.getByTestId("filter-chip-key_hash")).toHaveTextContent(mockKey.token);
});
it("restores the status filter from the URL and sends it to /key/list", async () => {
@ -853,6 +865,21 @@ describe("table state lives in the URL so it survives leaving and returning to t
expect(screen.queryByTestId("filter-chip-status")).not.toBeInTheDocument();
});
it("drops a hand-edited status from the URL when another filter chip is removed", async () => {
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<VirtualKeysTable />, {
searchParams: { filter_status: "bogus", filter_user: "user-42" },
onUrlUpdate,
});
fireEvent.click(await screen.findByTestId("filter-chip-remove-user_id"));
await waitFor(() => {
expect(lastSearchParam(onUrlUpdate, "filter_user")).toBeNull();
});
expect(lastSearchParam(onUrlUpdate, "filter_status")).toBeNull();
});
it("writes the search term to the URL", async () => {
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<VirtualKeysTable />, { onUrlUpdate });
@ -896,6 +923,40 @@ describe("table state lives in the URL so it survives leaving and returning to t
expect(screen.queryByTestId("filter-chip-user_id")).not.toBeInTheDocument();
});
it("writes the Organization and Key ID drawer filters to the URL and clears them again", async () => {
const user = userEvent.setup();
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<VirtualKeysTable />, { onUrlUpdate });
openFilters();
await chooseSelectOption(user, await screen.findByPlaceholderText(/Select an organization/), /Test Organization/);
fireEvent.change(screen.getByPlaceholderText(/Enter Key ID/), { target: { value: mockKey.token } });
fireEvent.click(screen.getByTestId("filter-drawer-apply"));
await waitFor(() => {
expect(lastSearchParam(onUrlUpdate, "filter_org")).toBe("org-1");
});
expect(lastSearchParam(onUrlUpdate, "filter_key_id")).toBe(mockKey.token);
expect(lastSearchParam(onUrlUpdate, "filter_org_id")).toBeNull();
expect(lastSearchParam(onUrlUpdate, "filter_key_hash")).toBeNull();
await waitFor(() => {
expect(mockUseKeys).toHaveBeenLastCalledWith(
1,
50,
expect.objectContaining({ organizationID: "org-1", keyHash: mockKey.token }),
);
});
fireEvent.click(screen.getByTestId("datatable-clear-filters"));
await waitFor(() => {
expect(lastSearchParam(onUrlUpdate, "filter_org")).toBeNull();
});
expect(lastSearchParam(onUrlUpdate, "filter_key_id")).toBeNull();
expect(screen.queryByTestId("filter-chip-org_id")).not.toBeInTheDocument();
expect(screen.queryByTestId("filter-chip-key_hash")).not.toBeInTheDocument();
});
it("returns to page 1 when the search term changes", async () => {
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<VirtualKeysTable />, { searchParams: { page: "3" }, onUrlUpdate });
@ -953,14 +1014,16 @@ describe("table state lives in the URL so it survives leaving and returning to t
});
});
it("falls back to the default sort when the URL names a column the table cannot sort by", async () => {
renderWithProviders(<VirtualKeysTable />, { searchParams: { sort_by: "totally_unknown_field" } });
it("falls back to the default sort column, keeping the URL's direction, when the table cannot sort by sort_by", async () => {
renderWithProviders(<VirtualKeysTable />, {
searchParams: { sort_by: "totally_unknown_field", sort_order: "asc" },
});
await waitFor(() => {
expect(mockUseKeys).toHaveBeenLastCalledWith(
1,
50,
expect.objectContaining({ sortBy: "created_at", sortOrder: "desc" }),
expect.objectContaining({ sortBy: "created_at", sortOrder: "asc" }),
);
});
expect(screen.getByText("Test Key Alias")).toBeInTheDocument();
@ -1003,3 +1066,72 @@ describe("table state lives in the URL so it survives leaving and returning to t
});
});
});
describe("column choices survive a reload", () => {
const STORAGE_KEY = "litellm_table_columns_virtual-keys";
const storedColumns = () => JSON.parse(localStorage.getItem(STORAGE_KEY) ?? "null");
it("hides a column that was hidden on a previous visit while the default-hidden columns stay hidden", () => {
localStorage.setItem(STORAGE_KEY, JSON.stringify({ budget_reset_at: false }));
renderWithProviders(<VirtualKeysTable />);
expect(screen.getByText("Test Key Alias")).toBeInTheDocument();
expect(screen.queryByText("Budget Reset")).not.toBeInTheDocument();
expect(screen.queryByText("Created By")).not.toBeInTheDocument();
});
it("writes a column toggled on through the Columns menu to storage and shows it again on the next mount", async () => {
const user = userEvent.setup();
const { unmount } = renderWithProviders(<VirtualKeysTable />);
expect(screen.queryByText("Created By")).not.toBeInTheDocument();
await user.click(screen.getByRole("button", { name: "Columns" }));
await user.click(await screen.findByText("Created By"));
await user.keyboard("{Escape}");
expect(storedColumns()).toEqual({ ...KEY_TABLE_HIDDEN_COLUMNS, created_by: true });
unmount();
renderWithProviders(<VirtualKeysTable />);
expect(screen.getByText("Created By")).toBeInTheDocument();
});
});
describe("a failed keys fetch does not rewrite the URL", () => {
const renderOnPage3OfMany = async () => {
mockUseKeys.mockReturnValue(keysResult([mockKey], { total_count: 200, total_pages: 4 }));
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
const view = renderWithProviders(<VirtualKeysTable />, { searchParams: { page: "3" }, onUrlUpdate });
await waitFor(() => {
expect(mockUseKeys).toHaveBeenLastCalledWith(3, 50, expect.anything());
});
return { ...view, onUrlUpdate };
};
it("keeps ?page=3 when the keys query errors, instead of snapping to page 1 on the empty count", async () => {
const { rerender, onUrlUpdate } = await renderOnPage3OfMany();
mockUseKeys.mockReturnValue(keysResult([], {}, { data: undefined, isError: true }));
rerender(<VirtualKeysTable />);
await new Promise((resolve) => setTimeout(resolve, 50));
expect(mockUseKeys).toHaveBeenLastCalledWith(3, 50, expect.anything());
expect(onUrlUpdate).not.toHaveBeenCalled();
});
it("still snaps ?page=3 back to the first page when the keys query succeeds with no rows", async () => {
const { rerender, onUrlUpdate } = await renderOnPage3OfMany();
mockUseKeys.mockReturnValue(keysResult([]));
rerender(<VirtualKeysTable />);
await waitFor(() => {
expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.anything());
});
await waitFor(() => {
expect(lastSearchParam(onUrlUpdate, "page")).toBeNull();
});
});
});

185
ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx
View file

@ -10,15 +10,18 @@ import {
DataTableFilterDrawer,
DataTableFilterField,
DataTableToolbar,
usePersistedColumnVisibility,
useUrlTableState,
type UrlTableStateOptions,
} from "@/components/shared/DataTable";
import { SearchSelect } from "@/components/shared/SearchSelect";
import { PageHeader } from "@/components/shared/PageHeader";
import { Input } from "@/components/ui/input";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
import { ColumnFiltersState, functionalUpdate, OnChangeFn } from "@tanstack/react-table";
import { KeyRound } from "lucide-react";
import { createParser, parseAsInteger, parseAsString, parseAsStringLiteral, useQueryState, useQueryStates } from "nuqs";
import { parseAsString, useQueryState } from "nuqs";
import React, { useCallback, useMemo, useState } from "react";
import { KeyResponse, Team } from "../key_team_helpers/key_list";
@ -56,44 +59,30 @@ const STATUS_FILTER_ITEMS = [
...KEY_STATUS_VALUES.map((value) => ({ value, label: KEY_STATUS_LABELS[value] })),
];
const isKeyStatusFilter = (value: string): value is KeyStatusFilter =>
(KEY_STATUS_VALUES as readonly string[]).includes(value);
const isKeyStatusFilter = (value: unknown): value is KeyStatusFilter =>
(KEY_STATUS_VALUES as readonly unknown[]).includes(value);
const DEFAULT_SORT_BY = "created_at";
const DEFAULT_SORT_ORDER = "desc";
const DEFAULT_PAGE_SIZE = 50;
const MAX_PAGE_SIZE = 100;
const MAX_PAGE = 100_000;
const isUsableFilter = (filter: ColumnFiltersState[number]): boolean =>
filter.id !== "status" || isKeyStatusFilter(filter.value);
const boundedInteger = (min: number, max: number, fallback: number) =>
createParser({
parse: (value: string) => {
const parsed = parseAsInteger.parse(value);
return parsed === null ? null : Math.min(Math.max(parsed, min), max);
},
serialize: String,
}).withDefault(fallback);
// The filters carry a prefix because /api-keys also takes team_id, key_alias and key_type
// as create-key prefills; an unprefixed filter would hijack those deep links.
const TABLE_STATE = {
key_search: parseAsString.withDefault(""),
sort_by: parseAsString.withDefault(DEFAULT_SORT_BY),
sort_order: parseAsStringLiteral(["asc", "desc"] as const).withDefault(DEFAULT_SORT_ORDER),
page: boundedInteger(1, MAX_PAGE, 1),
page_size: boundedInteger(1, MAX_PAGE_SIZE, DEFAULT_PAGE_SIZE),
filter_team: parseAsString.withDefault(""),
filter_org: parseAsString.withDefault(""),
filter_user: parseAsString.withDefault(""),
filter_key_id: parseAsString.withDefault(""),
filter_status: parseAsString.withDefault(""),
const TABLE_STATE_OPTIONS: UrlTableStateOptions<FilterColumn> = {
sortFields: KEY_TABLE_SORT_FIELDS,
defaultSort: { id: "created_at", desc: true },
defaultPageSize: 50,
maxPageSize: 100,
filterColumns: FILTER_COLUMNS,
urlKeys: {
search: "key_search",
filter_team_id: "filter_team",
filter_org_id: "filter_org",
filter_user_id: "filter_user",
filter_key_hash: "filter_key_id",
},
};
const toSortOrder = (active: SortingState[number]): "asc" | "desc" => (active.desc ? "desc" : "asc");
const filterValue = (filters: ColumnFiltersState, column: FilterColumn): string | null => {
const appliedFilter = (filters: ColumnFiltersState, column: FilterColumn): string | undefined => {
const value = filters.find((filter) => filter.id === column)?.value;
return (typeof value === "string" ? value.trim() : "") || null;
return typeof value === "string" ? value : undefined;
};
export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
@ -103,50 +92,38 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
const allTeams = useMemo<Team[]>(() => fetchedTeams ?? [], [fetchedTeams]);
const [selectedKeyId, setSelectedKeyId] = useQueryState("key", parseAsString.withOptions({ history: "push" }));
const [tableState, setTableState] = useQueryStates(TABLE_STATE);
const {
search: searchInput,
setSearch,
sorting,
onSortingChange,
pagination,
onPaginationChange,
columnFilters: urlColumnFilters,
onColumnFiltersChange: setUrlColumnFilters,
} = useUrlTableState(TABLE_STATE_OPTIONS);
const columnFilters = useMemo(() => urlColumnFilters.filter(isUsableFilter), [urlColumnFilters]);
const onColumnFiltersChange = useCallback<OnChangeFn<ColumnFiltersState>>(
(updaterOrValue) => setUrlColumnFilters(functionalUpdate(updaterOrValue, columnFilters)),
[columnFilters, setUrlColumnFilters],
);
const { columnVisibility, onColumnVisibilityChange } = usePersistedColumnVisibility(
"virtual-keys",
KEY_TABLE_HIDDEN_COLUMNS,
);
const [filtersOpen, setFiltersOpen] = useState(false);
const searchInput = tableState.key_search;
const [searchQuery] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS });
// A hand-edited sort_by the table cannot sort by would 400 at /key/list and leave the page loading.
const sortBy = KEY_TABLE_SORT_FIELDS.includes(tableState.sort_by) ? tableState.sort_by : DEFAULT_SORT_BY;
const sorting = useMemo<SortingState>(
() => [{ id: sortBy, desc: tableState.sort_order === "desc" }],
[sortBy, tableState.sort_order],
);
const tablePagination = useMemo<PaginationState>(
() => ({ pageIndex: tableState.page - 1, pageSize: tableState.page_size }),
[tableState.page, tableState.page_size],
);
const { filter_team, filter_org, filter_user, filter_key_id, filter_status } = tableState;
const appliedFilters = useMemo(
() => ({
team_id: filter_team.trim(),
org_id: filter_org.trim(),
user_id: filter_user.trim(),
key_hash: filter_key_id.trim(),
status: isKeyStatusFilter(filter_status) ? filter_status : "",
}),
[filter_team, filter_org, filter_user, filter_key_id, filter_status],
);
const columnFilters = useMemo<ColumnFiltersState>(
() =>
FILTER_COLUMNS.filter((column) => appliedFilters[column]).map((column) => ({
id: column,
value: appliedFilters[column],
})),
[appliedFilters],
);
const [activeSort] = sorting;
const keyListOptions = {
teamID: appliedFilters.team_id || undefined,
organizationID: appliedFilters.org_id || undefined,
teamID: appliedFilter(columnFilters, "team_id"),
organizationID: appliedFilter(columnFilters, "org_id"),
search: searchQuery.trim() || undefined,
userID: appliedFilters.user_id || undefined,
keyHash: appliedFilters.key_hash || undefined,
status: appliedFilters.status || undefined,
sortBy,
sortOrder: tableState.sort_order,
userID: appliedFilter(columnFilters, "user_id"),
keyHash: appliedFilter(columnFilters, "key_hash"),
status: appliedFilter(columnFilters, "status"),
sortBy: activeSort.id,
sortOrder: activeSort.desc ? "desc" : "asc",
expand: "user",
};
@ -155,55 +132,13 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
isPending,
isPlaceholderData,
isFetching,
isError,
refetch,
} = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, keyListOptions);
} = useKeys(pagination.pageIndex + 1, pagination.pageSize, keyListOptions);
const keyList = useMemo(() => keys?.keys ?? [], [keys]);
const rowCount = keys?.total_count ?? 0;
const handleSearchChange = useCallback(
(value: string) => {
void setTableState({ key_search: value || null, page: null });
},
[setTableState],
);
const handleSortingChange = useCallback<OnChangeFn<SortingState>>(
(updaterOrValue) => {
const active = functionalUpdate(updaterOrValue, sorting)[0];
void setTableState({
sort_by: active?.id ?? null,
sort_order: active ? toSortOrder(active) : null,
page: null,
});
},
[sorting, setTableState],
);
const handleColumnFiltersChange = useCallback<OnChangeFn<ColumnFiltersState>>(
(updaterOrValue) => {
const next = functionalUpdate(updaterOrValue, columnFilters);
const nextFilters = {
filter_team: filterValue(next, "team_id"),
filter_org: filterValue(next, "org_id"),
filter_user: filterValue(next, "user_id"),
filter_key_id: filterValue(next, "key_hash"),
filter_status: filterValue(next, "status"),
page: null,
};
void setTableState(nextFilters);
},
[columnFilters, setTableState],
);
const handlePaginationChange = useCallback<OnChangeFn<PaginationState>>(
(updaterOrValue) => {
const next = functionalUpdate(updaterOrValue, tablePagination);
void setTableState({ page: next.pageIndex + 1, page_size: next.pageSize });
},
[tablePagination, setTableState],
);
const columns = useMemo(
() => getKeyTableColumns({ allTeams, organizations, onSelectKey: (key) => void setSelectedKeyId(key.token) }),
[allTeams, organizations, setSelectedKeyId],
@ -296,20 +231,22 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
data={keyList}
columns={columns}
getRowId={(row) => row.token}
defaultColumnVisibility={KEY_TABLE_HIDDEN_COLUMNS}
columnVisibility={columnVisibility}
onColumnVisibilityChange={onColumnVisibilityChange}
sortingMode="server"
sorting={sorting}
onSortingChange={handleSortingChange}
onSortingChange={onSortingChange}
paginationMode="server"
pagination={tablePagination}
onPaginationChange={handlePaginationChange}
pagination={pagination}
onPaginationChange={onPaginationChange}
rowCount={rowCount}
filterMode="server"
columnFilters={columnFilters}
onColumnFiltersChange={handleColumnFiltersChange}
onColumnFiltersChange={onColumnFiltersChange}
enableColumnResizing
columnResizeMode="onChange"
isLoading={isPending || isPlaceholderData}
isError={isError}
loadingMessage="Loading keys..."
noDataMessage="No keys found"
fillHeight
@ -319,7 +256,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
<DataTableToolbar
table={table}
searchValue={searchInput}
onSearchChange={handleSearchChange}
onSearchChange={setSearch}
searchPlaceholder="Search by key alias or ID…"
onRefresh={() => refetch?.()}
isRefreshing={isFetching}

38
ui/litellm-dashboard/src/components/networking.test.ts
View file

@ -20,25 +20,39 @@ describe("networking - expired session handling", () => {
global.fetch = originalFetch;
});
it("should call clearTokenCookies on expired session", async () => {
const errorData = "Authentication Error - Expired Key";
const { toast } = await import("@/lib/toast");
const loadFreshHandleError = async () => {
vi.resetModules();
const fresh = await import("./networking");
return fresh.handleError;
};
if (errorData.includes("Authentication Error - Expired Key")) {
toast.info("UI Session Expired. Logging out.");
clearTokenCookies();
}
const stubLocation = (pathname: string, search: string, hash: string) => {
const location = { pathname, search, hash, href: "" };
vi.stubGlobal("window", { location });
return location;
};
afterEach(() => {
vi.unstubAllGlobals();
});
it("keeps the query string and hash on the redirect after session expiry", async () => {
const handleError = await loadFreshHandleError();
const location = stubLocation("/ui/api-keys/", "?filter_team=t1&page=2", "#row-3");
await handleError("Authentication Error - Expired Key");
expect(location.href).toBe("/ui/api-keys/?filter_team=t1&page=2#row-3");
expect(clearTokenCookies).toHaveBeenCalledOnce();
});
it("should not clear cookies for non-authentication errors", () => {
const errorData = "Some other error";
it("does not navigate or clear cookies for other errors", async () => {
const handleError = await loadFreshHandleError();
const location = stubLocation("/ui/api-keys/", "?filter_team=t1&page=2", "");
if (errorData.includes("Authentication Error - Expired Key")) {
clearTokenCookies();
}
await handleError("Some other error");
expect(location.href).toBe("");
expect(clearTokenCookies).not.toHaveBeenCalled();
});

2
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -383,7 +383,7 @@ export const handleError = async (errorData: string | any) => {
clearTokenCookies();
const browserLocation = getWindowLocation();
if (browserLocation) {
window.location.href = browserLocation.pathname;
window.location.href = browserLocation.pathname + browserLocation.search + browserLocation.hash;
}
}
lastErrorTime = currentTime;

27
ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx
View file

@ -1,4 +1,10 @@
import type { ColumnDef, PaginationState, RowSelectionState, SortingState } from "@tanstack/react-table";
import type {
ColumnDef,
PaginationState,
RowSelectionState,
SortingState,
VisibilityState,
} from "@tanstack/react-table";
import { DataTable } from "./DataTable";
@ -12,6 +18,7 @@ const columns: ColumnDef<Row, unknown>[] = [];
const sorting: SortingState = [{ id: "name", desc: false }];
const pagination: PaginationState = { pageIndex: 0, pageSize: 10 };
const rowSelection: RowSelectionState = { r1: true };
const columnVisibility: VisibilityState = { name: false };
const noop = () => {};
export const uncontrolled = <DataTable data={data} columns={columns} defaultSorting={sorting} />;
@ -32,6 +39,8 @@ export const controlled = (
onColumnFiltersChange={noop}
rowSelection={rowSelection}
onRowSelectionChange={noop}
columnVisibility={columnVisibility}
onColumnVisibilityChange={noop}
/>
);
@ -65,3 +74,19 @@ export const selectionWithoutHandler = (
// @ts-expect-error a controlled `rowSelection` needs `onRowSelectionChange` or selection changes are dropped
<DataTable data={data} columns={columns} rowSelection={rowSelection} />
);
export const visibilityWithoutHandler = (
// @ts-expect-error a controlled `columnVisibility` needs `onColumnVisibilityChange` or Columns-menu toggles are dropped
<DataTable data={data} columns={columns} columnVisibility={columnVisibility} />
);
export const bothVisibilitySources = (
// @ts-expect-error `defaultColumnVisibility` seeds uncontrolled visibility, so it cannot pair with a controlled `columnVisibility`
<DataTable
data={data}
columns={columns}
defaultColumnVisibility={columnVisibility}
columnVisibility={columnVisibility}
onColumnVisibilityChange={noop}
/>
);

171
ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx
View file

@ -1,4 +1,4 @@
import type { ColumnDef, ExpandedState, OnChangeFn, PaginationState } from "@tanstack/react-table";
import type { ColumnDef, ExpandedState, OnChangeFn, PaginationState, VisibilityState } from "@tanstack/react-table";
import { render, screen, waitFor, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { useState } from "react";
@ -278,11 +278,18 @@ describe("DataTable pagination", () => {
type ServerPageHarnessProps = {
rowCount: number;
isLoading?: boolean;
isError?: boolean;
initialPageIndex: number;
onChange: (next: PaginationState) => void;
};
function ServerPageHarness({ rowCount, isLoading = false, initialPageIndex, onChange }: ServerPageHarnessProps) {
function ServerPageHarness({
rowCount,
isLoading = false,
isError = false,
initialPageIndex,
onChange,
}: ServerPageHarnessProps) {
const [pagination, setPagination] = useState<PaginationState>({ pageIndex: initialPageIndex, pageSize: 10 });
const handleChange: OnChangeFn<PaginationState> = (updater) => {
const next = typeof updater === "function" ? updater(pagination) : updater;
@ -298,6 +305,7 @@ describe("DataTable pagination", () => {
onPaginationChange={handleChange}
rowCount={rowCount}
isLoading={isLoading}
isError={isError}
/>
);
}
@ -339,6 +347,102 @@ describe("DataTable pagination", () => {
expect(onChange).toHaveBeenCalledTimes(1);
expect(screen.getByText("Page 2 of 2")).toBeInTheDocument();
});
it("server mode keeps a deep-linked page when the fetch failed, instead of snapping to page 1 on rowCount 0", async () => {
const onChange = vi.fn();
render(<ServerPageHarness rowCount={0} isError initialPageIndex={2} onChange={onChange} />);
expect(screen.getByText("Page 3 of 1")).toBeInTheDocument();
await new Promise((resolve) => setTimeout(resolve, 20));
expect(onChange).not.toHaveBeenCalled();
});
type ClientPageHarnessProps = {
data: Person[];
isLoading?: boolean;
initialPageIndex: number;
onChange: (next: PaginationState) => void;
};
function ClientPageHarness({ data, isLoading = false, initialPageIndex, onChange }: ClientPageHarnessProps) {
const [pagination, setPagination] = useState<PaginationState>({ pageIndex: initialPageIndex, pageSize: 2 });
const handleChange: OnChangeFn<PaginationState> = (updater) => {
const next = typeof updater === "function" ? updater(pagination) : updater;
onChange(next);
setPagination(next);
};
return (
<DataTable
data={data}
columns={nameCellColumns}
paginationMode="client"
pageSizeOptions={[2]}
pagination={pagination}
onPaginationChange={handleChange}
isLoading={isLoading}
/>
);
}
it("client mode keeps a controlled page when rows arrive after loading and when they are refetched", async () => {
const onChange = vi.fn();
const { rerender } = render(<ClientPageHarness data={[]} isLoading initialPageIndex={1} onChange={onChange} />);
rerender(<ClientPageHarness data={fivePeople} initialPageIndex={1} onChange={onChange} />);
await new Promise((resolve) => setTimeout(resolve, 20));
expect(names()).toEqual(["P2", "P3"]);
rerender(<ClientPageHarness data={[...fivePeople]} initialPageIndex={1} onChange={onChange} />);
await new Promise((resolve) => setTimeout(resolve, 20));
expect(names()).toEqual(["P2", "P3"]);
expect(onChange).not.toHaveBeenCalled();
});
it("client mode snaps a controlled page past the end back to the last page", async () => {
const onChange = vi.fn();
render(<ClientPageHarness data={fivePeople} initialPageIndex={5} onChange={onChange} />);
await waitFor(() => expect(onChange).toHaveBeenCalledWith({ pageIndex: 2, pageSize: 2 }));
expect(onChange).toHaveBeenCalledTimes(1);
expect(names()).toEqual(["P4"]);
});
it("client mode leaves a controlled page alone while there are no rows to page through", async () => {
const onChange = vi.fn();
render(<ClientPageHarness data={[]} initialPageIndex={3} onChange={onChange} />);
await new Promise((resolve) => setTimeout(resolve, 20));
expect(onChange).not.toHaveBeenCalled();
});
it("client mode without a controlled page still returns to the first page when the rows change", async () => {
const user = userEvent.setup();
const { rerender } = render(
<DataTable data={fivePeople} columns={nameCellColumns} paginationMode="client" pageSizeOptions={[2]} />,
);
await user.click(screen.getByTestId("pagination-next"));
expect(names()).toEqual(["P2", "P3"]);
rerender(
<DataTable data={[...fivePeople]} columns={nameCellColumns} paginationMode="client" pageSizeOptions={[2]} />,
);
await waitFor(() => expect(names()).toEqual(["P0", "P1"]));
});
it("server mode resumes clamping once the error clears and a real rowCount arrives", async () => {
const onChange = vi.fn();
const { rerender } = render(<ServerPageHarness rowCount={0} isError initialPageIndex={2} onChange={onChange} />);
await new Promise((resolve) => setTimeout(resolve, 20));
expect(onChange).not.toHaveBeenCalled();
rerender(<ServerPageHarness rowCount={15} initialPageIndex={2} onChange={onChange} />);
await waitFor(() => expect(onChange).toHaveBeenCalledWith({ pageIndex: 1, pageSize: 10 }));
expect(onChange).toHaveBeenCalledTimes(1);
expect(screen.getByText("Page 2 of 2")).toBeInTheDocument();
});
});
describe("DataTable filtering", () => {
@ -555,6 +659,69 @@ describe("DataTable column visibility", () => {
expect(await screen.findByTestId("view-option-email")).toBeInTheDocument();
expect(screen.queryByTestId("view-option-name")).not.toBeInTheDocument();
});
it("uncontrolled mode seeds hidden columns from defaultColumnVisibility and still toggles internally", async () => {
const user = userEvent.setup();
render(
<DataTable
data={CHARLIE_ALICE_BOB}
columns={nameEmailColumns}
defaultColumnVisibility={{ email: false }}
toolbar={(table) => <DataTableViewOptions table={table} />}
/>,
);
expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument();
await user.click(screen.getByTestId("view-options-trigger"));
await user.click(await screen.findByTestId("view-option-email"));
expect(await screen.findByRole("columnheader", { name: "Email" })).toBeInTheDocument();
});
it("controlled mode hides columns from the prop and reports toggles without changing them locally", async () => {
const user = userEvent.setup();
const onColumnVisibilityChange = vi.fn<OnChangeFn<VisibilityState>>();
render(
<DataTable
data={CHARLIE_ALICE_BOB}
columns={nameEmailColumns}
columnVisibility={{ email: false }}
onColumnVisibilityChange={onColumnVisibilityChange}
toolbar={(table) => <DataTableViewOptions table={table} />}
/>,
);
expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument();
await user.click(screen.getByTestId("view-options-trigger"));
await user.click(await screen.findByTestId("view-option-email"));
expect(onColumnVisibilityChange).toHaveBeenCalledTimes(1);
const updater = onColumnVisibilityChange.mock.calls[0]?.[0];
const next = typeof updater === "function" ? updater({ email: false }) : updater;
expect(next).toEqual({ email: true });
expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument();
});
it("controlled mode reveals the column once the parent applies the reported change", async () => {
const user = userEvent.setup();
const Harness = () => {
const [columnVisibility, setColumnVisibility] = useState<VisibilityState>({ email: false });
return (
<DataTable
data={CHARLIE_ALICE_BOB}
columns={nameEmailColumns}
columnVisibility={columnVisibility}
onColumnVisibilityChange={setColumnVisibility}
toolbar={(table) => <DataTableViewOptions table={table} />}
/>
);
};
render(<Harness />);
expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument();
await user.click(screen.getByTestId("view-options-trigger"));
await user.click(await screen.findByTestId("view-option-email"));
expect(await screen.findByRole("columnheader", { name: "Email" })).toBeInTheDocument();
});
});
describe("DataTable pinned columns", () => {

47
ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx
View file

@ -425,7 +425,7 @@ function useControllable<T>(
return { value: internal, onChange: setInternal };
}
function useServerPageClamp(
function usePageClamp(
active: boolean,
rowCount: number | undefined,
pagination: { value: PaginationState; onChange: OnChangeFn<PaginationState> },
@ -457,6 +457,7 @@ function useDataTableInstance<TData extends RowData, TValue>(
onPaginationChange,
rowCount,
isLoading = false,
isError,
pageSizeOptions = DEFAULT_PAGE_SIZE_OPTIONS,
filterMode = "none",
columnFilters,
@ -466,6 +467,8 @@ function useDataTableInstance<TData extends RowData, TValue>(
onGlobalFilterChange,
enableColumnResizing = false,
columnResizeMode = "onEnd",
columnVisibility,
onColumnVisibilityChange,
defaultColumnVisibility,
getRowCanExpand,
renderSubComponent,
@ -481,7 +484,6 @@ function useDataTableInstance<TData extends RowData, TValue>(
pageIndex: 0,
pageSize: pageSizeOptions[0] ?? 25,
});
useServerPageClamp(paginationMode === "server" && !isLoading, rowCount, paginationState);
const filterState = useControllable<ColumnFiltersState>(
columnFilters,
onColumnFiltersChange,
@ -490,7 +492,11 @@ function useDataTableInstance<TData extends RowData, TValue>(
const globalFilterState = useControllable<string>(globalFilter, onGlobalFilterChange, "");
const expandedState = useControllable<ExpandedState>(expanded, onExpandedChange, {});
const rowSelectionState = useControllable<RowSelectionState>(rowSelection, onRowSelectionChange, {});
const [columnVisibility, setColumnVisibility] = useState<VisibilityState>(defaultColumnVisibility ?? {});
const columnVisibilityState = useControllable<VisibilityState>(
columnVisibility,
onColumnVisibilityChange,
defaultColumnVisibility ?? {},
);
const [columnSizing, setColumnSizing] = useState<ColumnSizingState>({});
const columnPinning = React.useMemo(() => derivePinning(columns), [columns]);
const expansionGuard = renderSubComponent !== undefined ? getRowCanExpand : undefined;
@ -505,7 +511,7 @@ function useDataTableInstance<TData extends RowData, TValue>(
globalFilter: globalFilterState.value,
expanded: expandedState.value,
rowSelection: rowSelectionState.value,
columnVisibility,
columnVisibility: columnVisibilityState.value,
columnSizing,
},
initialState: { columnPinning },
@ -521,7 +527,7 @@ function useDataTableInstance<TData extends RowData, TValue>(
onGlobalFilterChange: globalFilterState.onChange,
onExpandedChange: expandedState.onChange,
onRowSelectionChange: rowSelectionState.onChange,
onColumnVisibilityChange: setColumnVisibility,
onColumnVisibilityChange: columnVisibilityState.onChange,
onColumnSizingChange: setColumnSizing,
getColumnCanGlobalFilter: (column) => columnCanGlobalFilter(data[0], column),
getCoreRowModel: getCoreRowModel(),
@ -529,9 +535,38 @@ function useDataTableInstance<TData extends RowData, TValue>(
...(getRowId !== undefined ? { getRowId } : {}),
...(enableRowSelection !== undefined ? { enableRowSelection } : {}),
...(paginationMode === "server" && rowCount !== undefined ? { rowCount } : {}),
autoResetPageIndex: pagination === undefined && paginationMode !== "server",
};
return useReactTable(tableOptions);
const table = useReactTable(tableOptions);
const clampOptions: SettledPageClampOptions = {
paginationMode,
controlled: pagination !== undefined,
settled: !isLoading && !isError,
rowCount,
pagination: paginationState,
};
useSettledPageClamp(table, clampOptions);
return table;
}
type SettledPageClampOptions = {
paginationMode: PaginationMode;
controlled: boolean;
settled: boolean;
rowCount: number | undefined;
pagination: { value: PaginationState; onChange: OnChangeFn<PaginationState> };
};
function useSettledPageClamp<TData extends RowData>(table: Table<TData>, options: SettledPageClampOptions): void {
const { paginationMode, controlled, settled, rowCount, pagination } = options;
const clientRowCount = paginationMode === "client" ? table.getPrePaginationRowModel().rows.length : 0;
const clientPageIsClampable = paginationMode === "client" && controlled && clientRowCount > 0;
usePageClamp(
settled && (paginationMode === "server" || clientPageIsClampable),
paginationMode === "server" ? rowCount : clientRowCount,
pagination,
);
}
export function DataTable<TData extends RowData, TValue>(props: DataTableProps<TData, TValue>) {

2
ui/litellm-dashboard/src/components/shared/DataTable/index.ts
View file

@ -12,6 +12,8 @@ export {
type DataTableSortVariant,
type DataTableSortField,
} from "./DataTableSortHeader";
export { usePersistedColumnVisibility } from "./usePersistedColumnVisibility";
export { useUrlTableState, type UrlTableState, type UrlTableStateOptions } from "./useUrlTableState";
export type { DataTablePaginationProps } from "./DataTablePagination";
export type {
ColumnPinnedSide,

19
ui/litellm-dashboard/src/components/shared/DataTable/types.ts
View file

@ -27,6 +27,7 @@ export interface DataTableResolvedProps<TData extends RowData, TValue> {
getRowId?: (row: TData, index: number, parent?: Row<TData>) => string;
isLoading?: boolean;
isError?: boolean;
loadingMessage?: string;
skeletonRowCount?: number;
noDataMessage?: React.ReactNode;
@ -53,6 +54,8 @@ export interface DataTableResolvedProps<TData extends RowData, TValue> {
enableColumnResizing?: boolean;
columnResizeMode?: ColumnResizeMode;
columnVisibility?: VisibilityState;
onColumnVisibilityChange?: OnChangeFn<VisibilityState>;
defaultColumnVisibility?: VisibilityState;
getRowCanExpand?: (row: Row<TData>) => boolean;
@ -96,6 +99,9 @@ type DataTableBaseProps<TData extends RowData, TValue> = Omit<
| "columnFilters"
| "onColumnFiltersChange"
| "defaultColumnFilters"
| "columnVisibility"
| "onColumnVisibilityChange"
| "defaultColumnVisibility"
| "rowSelection"
| "onRowSelectionChange"
>;
@ -142,6 +148,18 @@ type FilterProps =
defaultColumnFilters?: ColumnFiltersState;
};
type ColumnVisibilityProps =
| {
columnVisibility: VisibilityState;
onColumnVisibilityChange: OnChangeFn<VisibilityState>;
defaultColumnVisibility?: never;
}
| {
columnVisibility?: never;
onColumnVisibilityChange?: never;
defaultColumnVisibility?: VisibilityState;
};
type RowSelectionProps =
| { rowSelection: RowSelectionState; onRowSelectionChange: OnChangeFn<RowSelectionState> }
| { rowSelection?: never; onRowSelectionChange?: OnChangeFn<RowSelectionState> };
@ -150,4 +168,5 @@ export type DataTableProps<TData extends RowData, TValue> = DataTableBaseProps<T
SortingProps &
PaginationProps &
FilterProps &
ColumnVisibilityProps &
RowSelectionProps;

194
ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx Normal file
View file

@ -0,0 +1,194 @@
import type { VisibilityState } from "@tanstack/react-table";
import { act, renderHook } from "@testing-library/react";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { usePersistedColumnVisibility } from "./usePersistedColumnVisibility";
const keyFor = (tableId: string): string => `litellm_table_columns_${tableId}`;
const stored = (tableId: string): unknown => {
const raw = localStorage.getItem(keyFor(tableId));
return raw === null ? null : JSON.parse(raw);
};
const showEveryColumn = (previous: VisibilityState): VisibilityState =>
Object.fromEntries(Object.keys(previous).map((column) => [column, true]));
describe("usePersistedColumnVisibility", () => {
beforeEach(() => {
localStorage.clear();
});
afterEach(() => {
localStorage.clear();
vi.restoreAllMocks();
});
it("layers the stored choices over the defaults, so a default added after the snapshot still applies", () => {
localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false, spend: true }));
const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false, name: false }));
expect(result.current.columnVisibility).toEqual({ email: false, spend: true, name: false });
});
it("falls back to the defaults when nothing is stored, and to {} without defaults", () => {
const withDefaults = renderHook(() => usePersistedColumnVisibility("keys", { spend: false }));
expect(withDefaults.result.current.columnVisibility).toEqual({ spend: false });
const bare = renderHook(() => usePersistedColumnVisibility("keys"));
expect(bare.result.current.columnVisibility).toEqual({});
});
it("writes an object update to state and storage", () => {
const { result } = renderHook(() => usePersistedColumnVisibility("keys"));
act(() => result.current.onColumnVisibilityChange({ email: false }));
expect(result.current.columnVisibility).toEqual({ email: false });
expect(stored("keys")).toEqual({ email: false });
});
it("resolves a function updater against the current state before persisting", () => {
localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false }));
const { result } = renderHook(() => usePersistedColumnVisibility("keys"));
act(() => result.current.onColumnVisibilityChange((previous) => ({ ...previous, name: false })));
expect(result.current.columnVisibility).toEqual({ email: false, name: false });
expect(stored("keys")).toEqual({ email: false, name: false });
});
it("hands a function updater the default-hidden columns, so showing every column sticks", () => {
const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false }));
act(() => result.current.onColumnVisibilityChange(showEveryColumn));
expect(result.current.columnVisibility).toEqual({ spend: true });
expect(stored("keys")).toEqual({ spend: true });
});
it.each([
["truncated JSON", '{"email":fal'],
["a JSON scalar", "42"],
["a JSON array", "[true]"],
["non-boolean values", JSON.stringify({ email: "no" })],
])("falls back to the defaults when storage holds %s", (_label, raw) => {
localStorage.setItem(keyFor("keys"), raw);
const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false }));
expect(result.current.columnVisibility).toEqual({ spend: false });
});
it("keeps distinct tableIds isolated in state and storage", () => {
const keys = renderHook(() => usePersistedColumnVisibility("keys"));
const teams = renderHook(() => usePersistedColumnVisibility("teams"));
act(() => keys.result.current.onColumnVisibilityChange({ email: false }));
expect(keys.result.current.columnVisibility).toEqual({ email: false });
expect(teams.result.current.columnVisibility).toEqual({});
expect(stored("keys")).toEqual({ email: false });
expect(stored("teams")).toBeNull();
});
it("reads and writes the new table's columns after the tableId changes", () => {
localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false }));
localStorage.setItem(keyFor("teams"), JSON.stringify({ spend: false }));
const { result, rerender } = renderHook(({ tableId }) => usePersistedColumnVisibility(tableId), {
initialProps: { tableId: "keys" },
});
rerender({ tableId: "teams" });
expect(result.current.columnVisibility).toEqual({ spend: false });
act(() => result.current.onColumnVisibilityChange((previous) => ({ ...previous, name: false })));
expect(stored("teams")).toEqual({ spend: false, name: false });
expect(stored("keys")).toEqual({ email: false });
});
it("applies new defaults passed after mount", () => {
const initialProps: { defaults: VisibilityState } = { defaults: { spend: false } };
const { result, rerender } = renderHook(({ defaults }) => usePersistedColumnVisibility("keys", defaults), {
initialProps,
});
rerender({ defaults: { name: false } });
expect(result.current.columnVisibility).toEqual({ name: false });
});
it("shows a change another tab saved for the same table", () => {
const { result } = renderHook(() => usePersistedColumnVisibility("keys"));
act(() => {
localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false }));
window.dispatchEvent(new StorageEvent("storage", { key: keyFor("keys") }));
});
expect(result.current.columnVisibility).toEqual({ email: false });
});
it("keeps a toggle that storage refused, and saves the next one once storage accepts it", () => {
localStorage.setItem(keyFor("full"), JSON.stringify({ spend: false }));
vi.spyOn(console, "warn").mockImplementation(() => {});
vi.spyOn(Storage.prototype, "setItem").mockImplementationOnce(() => {
throw new Error("QuotaExceededError");
});
const { result } = renderHook(() => usePersistedColumnVisibility("full"));
act(() => result.current.onColumnVisibilityChange({ email: false }));
expect(result.current.columnVisibility).toEqual({ email: false });
expect(stored("full")).toEqual({ spend: false });
act(() => result.current.onColumnVisibilityChange({ name: false }));
expect(result.current.columnVisibility).toEqual({ name: false });
expect(stored("full")).toEqual({ name: false });
});
it("shows another tab's save over a toggle this tab could not save", () => {
vi.spyOn(console, "warn").mockImplementation(() => {});
vi.spyOn(Storage.prototype, "setItem").mockImplementationOnce(() => {
throw new Error("QuotaExceededError");
});
const { result } = renderHook(() => usePersistedColumnVisibility("shadowed"));
act(() => result.current.onColumnVisibilityChange({ email: false }));
act(() => {
localStorage.setItem(keyFor("shadowed"), JSON.stringify({ name: false }));
window.dispatchEvent(new StorageEvent("storage", { key: keyFor("shadowed") }));
});
expect(result.current.columnVisibility).toEqual({ name: false });
});
it("drops a toggle this tab could not save once another tab clears storage", () => {
vi.spyOn(console, "warn").mockImplementation(() => {});
vi.spyOn(Storage.prototype, "setItem").mockImplementationOnce(() => {
throw new Error("QuotaExceededError");
});
const { result } = renderHook(() => usePersistedColumnVisibility("cleared", { spend: false }));
act(() => result.current.onColumnVisibilityChange({ email: false }));
act(() => window.dispatchEvent(new StorageEvent("storage", { key: null })));
expect(result.current.columnVisibility).toEqual({ spend: false });
});
it("returns the defaults without throwing when storage is unavailable", () => {
vi.spyOn(console, "warn").mockImplementation(() => {});
vi.spyOn(Storage.prototype, "getItem").mockImplementation(() => {
throw new Error("SecurityError");
});
vi.spyOn(Storage.prototype, "setItem").mockImplementation(() => {
throw new Error("SecurityError");
});
const { result } = renderHook(() => usePersistedColumnVisibility("blocked", { spend: false }));
expect(result.current.columnVisibility).toEqual({ spend: false });
act(() => result.current.onColumnVisibilityChange((previous) => ({ ...previous, email: false })));
expect(result.current.columnVisibility).toEqual({ spend: false, email: false });
});
});

96
ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts Normal file
View file

@ -0,0 +1,96 @@
import type { OnChangeFn, VisibilityState } from "@tanstack/react-table";
import { useCallback, useMemo, useSyncExternalStore } from "react";
import {
LOCAL_STORAGE_EVENT,
emitLocalStorageChange,
getLocalStorageItem,
setLocalStorageItem,
} from "@/utils/localStorageUtils";
const STORAGE_KEY_PREFIX = "litellm_table_columns_";
const EMPTY_VISIBILITY: VisibilityState = {};
const unsavedWrites = new Map<string, string>();
function storageKey(tableId: string): string {
return `${STORAGE_KEY_PREFIX}${tableId}`;
}
function forgetUnsavedWrite(event: StorageEvent): void {
if (event.key === null) {
unsavedWrites.clear();
return;
}
unsavedWrites.delete(event.key);
}
function subscribe(onChange: () => void): () => void {
const onStorage = (event: StorageEvent): void => {
forgetUnsavedWrite(event);
onChange();
};
window.addEventListener("storage", onStorage);
window.addEventListener(LOCAL_STORAGE_EVENT, onChange);
return () => {
window.removeEventListener("storage", onStorage);
window.removeEventListener(LOCAL_STORAGE_EVENT, onChange);
};
}
function readRaw(key: string): string | null {
return unsavedWrites.get(key) ?? getLocalStorageItem(key);
}
function writeRaw(key: string, raw: string): void {
setLocalStorageItem(key, raw);
if (getLocalStorageItem(key) === raw) {
unsavedWrites.delete(key);
} else {
unsavedWrites.set(key, raw);
}
emitLocalStorageChange(key);
}
function isVisibilityState(value: unknown): value is VisibilityState {
if (typeof value !== "object" || value === null || Array.isArray(value)) {
return false;
}
return Object.values(value).every((visible) => typeof visible === "boolean");
}
function parseVisibility(raw: string | null, defaults: VisibilityState): VisibilityState {
if (raw === null) {
return defaults;
}
try {
const parsed: unknown = JSON.parse(raw);
return isVisibilityState(parsed) ? { ...defaults, ...parsed } : defaults;
} catch {
return defaults;
}
}
export function usePersistedColumnVisibility(
tableId: string,
defaults: VisibilityState = EMPTY_VISIBILITY,
): { columnVisibility: VisibilityState; onColumnVisibilityChange: OnChangeFn<VisibilityState> } {
const key = storageKey(tableId);
const raw = useSyncExternalStore(
subscribe,
() => readRaw(key),
() => null,
);
const columnVisibility = useMemo(() => parseVisibility(raw, defaults), [raw, defaults]);
const onColumnVisibilityChange = useCallback<OnChangeFn<VisibilityState>>(
(updater) => {
const next = typeof updater === "function" ? updater(parseVisibility(readRaw(key), defaults)) : updater;
writeRaw(key, JSON.stringify(next));
},
[key, defaults],
);
return { columnVisibility, onColumnVisibilityChange };
}

325
ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx Normal file
View file

@ -0,0 +1,325 @@
import { SortingState } from "@tanstack/react-table";
import { act, renderHook, waitFor } from "@testing-library/react";
import { withNuqsTestingAdapter, type OnUrlUpdateFunction } from "nuqs/adapters/testing";
import { describe, expect, it, Mock, vi } from "vitest";
import { useUrlTableState, type UrlTableStateOptions } from "./useUrlTableState";
const FILTER_COLUMNS = ["team_id", "user_id"] as const;
type FilterColumn = (typeof FILTER_COLUMNS)[number];
const BASE_OPTIONS: UrlTableStateOptions<FilterColumn> = {
sortFields: ["created_at", "spend", "key_alias"],
defaultSort: { id: "created_at", desc: true },
defaultPageSize: 50,
filterColumns: FILTER_COLUMNS,
};
const PREFIXED_AND_UNPREFIXED_PARAMS = {
audit_page: "2",
audit_page_size: "10",
audit_search: "prefixed",
audit_sort_by: "spend",
audit_sort_order: "asc",
audit_filter_team_id: "team-1",
page: "5",
search: "unprefixed",
filter_team_id: "other-team",
};
const RENAMED_AND_DEFAULT_PARAMS = {
key_search: "prod",
filter_team: "team-1",
search: "ignored",
filter_team_id: "ignored",
};
const flipDirection = (previous: SortingState): SortingState => previous.map((sort) => ({ ...sort, desc: !sort.desc }));
const renderTableState = (
searchParams: Record<string, string> = {},
overrides: Partial<UrlTableStateOptions<FilterColumn>> = {},
) => {
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
const options = { ...BASE_OPTIONS, ...overrides };
const hook = renderHook(() => useUrlTableState(options), {
wrapper: withNuqsTestingAdapter({ searchParams, onUrlUpdate, hasMemory: true }),
});
return { ...hook, onUrlUpdate };
};
const lastUrl = (onUrlUpdate: Mock<OnUrlUpdateFunction>) => {
const event = onUrlUpdate.mock.calls.at(-1)?.[0];
if (!event) throw new Error("no URL update was emitted");
return event;
};
const flushUrl = async (onUrlUpdate: Mock<OnUrlUpdateFunction>, write: () => void) => {
const callsBefore = onUrlUpdate.mock.calls.length;
await act(async () => {
write();
});
await waitFor(() => expect(onUrlUpdate.mock.calls.length).toBeGreaterThan(callsBefore));
return lastUrl(onUrlUpdate).searchParams;
};
describe("reading table state from the URL", () => {
it("falls back to the defaults when the URL carries no table state", () => {
const { result } = renderTableState();
expect(result.current.search).toBe("");
expect(result.current.sorting).toEqual([{ id: "created_at", desc: true }]);
expect(result.current.pagination).toEqual({ pageIndex: 0, pageSize: 50 });
expect(result.current.columnFilters).toEqual([]);
});
it("maps the 1-based page and page_size onto TanStack pagination", () => {
const { result } = renderTableState({ page: "3", page_size: "25" });
expect(result.current.pagination).toEqual({ pageIndex: 2, pageSize: 25 });
});
it.each(["0", "-3", "not-a-number"])("clamps a page of %s up to the first page", (page) => {
const { result } = renderTableState({ page });
expect(result.current.pagination.pageIndex).toBe(0);
});
it.each([
["1000", undefined, 100],
["1000", 20, 20],
["0", undefined, 1],
])("clamps a page_size of %s with maxPageSize %s to %s", (pageSize, maxPageSize, expected) => {
const { result } = renderTableState({ page_size: pageSize }, { maxPageSize });
expect(result.current.pagination.pageSize).toBe(expected);
});
it("reads a sortable sort_by and its sort_order", () => {
const { result } = renderTableState({ sort_by: "spend", sort_order: "asc" });
expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]);
});
it("resolves a sort_by outside the allow-list to the default column while keeping the URL's direction", () => {
const { result } = renderTableState({ sort_by: "totally_unknown", sort_order: "asc" });
expect(result.current.sorting).toEqual([{ id: "created_at", desc: false }]);
});
it("maps filter_<column> params onto columnFilters, trimming whitespace and dropping blanks", () => {
const { result } = renderTableState({ filter_team_id: "team-1", filter_user_id: " " });
expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]);
const trimmed = renderTableState({ filter_user_id: " user-42 " });
expect(trimmed.result.current.columnFilters).toEqual([{ id: "user_id", value: "user-42" }]);
});
it("reads the search term verbatim so the input can hold trailing spaces", () => {
const { result } = renderTableState({ search: "prod " });
expect(result.current.search).toBe("prod ");
});
it("reads every key under keyPrefix and ignores the unprefixed ones", () => {
const { result } = renderTableState(PREFIXED_AND_UNPREFIXED_PARAMS, { keyPrefix: "audit_" });
expect(result.current.pagination).toEqual({ pageIndex: 1, pageSize: 10 });
expect(result.current.search).toBe("prefixed");
expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]);
expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]);
});
it("reads renamed keys from urlKeys and ignores the default names", () => {
const { result } = renderTableState(RENAMED_AND_DEFAULT_PARAMS, {
urlKeys: { search: "key_search", filter_team_id: "filter_team" },
});
expect(result.current.search).toBe("prod");
expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]);
});
it("applies keyPrefix in front of a renamed key", () => {
const { result } = renderTableState(
{ audit_key_search: "prod", key_search: "ignored" },
{ keyPrefix: "audit_", urlKeys: { search: "key_search" } },
);
expect(result.current.search).toBe("prod");
});
});
describe("writing table state to the URL", () => {
it("resolves a function updater against the current pagination and replaces history", async () => {
const { result, onUrlUpdate } = renderTableState({ page: "2" });
const url = await flushUrl(onUrlUpdate, () =>
result.current.onPaginationChange((previous) => ({ ...previous, pageIndex: previous.pageIndex + 1 })),
);
expect(url.get("page")).toBe("3");
expect(url.has("page_size")).toBe(false);
expect(lastUrl(onUrlUpdate).options.history).toBe("replace");
expect(result.current.pagination).toEqual({ pageIndex: 2, pageSize: 50 });
});
it("writes page_size and drops it again once it returns to the default", async () => {
const { result, onUrlUpdate } = renderTableState();
const withSize = await flushUrl(onUrlUpdate, () =>
result.current.onPaginationChange({ pageIndex: 0, pageSize: 25 }),
);
expect(withSize.get("page_size")).toBe("25");
expect(withSize.has("page")).toBe(false);
const backToDefault = await flushUrl(onUrlUpdate, () =>
result.current.onPaginationChange({ pageIndex: 0, pageSize: 50 }),
);
expect(backToDefault.has("page_size")).toBe(false);
});
it("setSearch writes the term and returns to the first page", async () => {
const { result, onUrlUpdate } = renderTableState({ page: "3" });
const url = await flushUrl(onUrlUpdate, () => result.current.setSearch("prod"));
expect(url.get("search")).toBe("prod");
expect(url.has("page")).toBe(false);
expect(result.current.search).toBe("prod");
expect(result.current.pagination.pageIndex).toBe(0);
});
it("setSearch with an empty string removes the key", async () => {
const { result, onUrlUpdate } = renderTableState({ search: "prod" });
const url = await flushUrl(onUrlUpdate, () => result.current.setSearch(""));
expect(url.has("search")).toBe(false);
expect(result.current.search).toBe("");
});
it("onSortingChange writes sort_by and sort_order and returns to the first page", async () => {
const { result, onUrlUpdate } = renderTableState({ page: "3" });
const url = await flushUrl(onUrlUpdate, () => result.current.onSortingChange([{ id: "spend", desc: false }]));
expect(url.get("sort_by")).toBe("spend");
expect(url.get("sort_order")).toBe("asc");
expect(url.has("page")).toBe(false);
expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]);
});
it("onSortingChange drops the keys when the sort matches the default or is cleared", async () => {
const { result, onUrlUpdate } = renderTableState({ sort_by: "spend", sort_order: "asc" });
const explicitDefault = await flushUrl(onUrlUpdate, () =>
result.current.onSortingChange([{ id: "created_at", desc: true }]),
);
expect(explicitDefault.has("sort_by")).toBe(false);
expect(explicitDefault.has("sort_order")).toBe(false);
await flushUrl(onUrlUpdate, () => result.current.onSortingChange([{ id: "key_alias", desc: false }]));
const cleared = await flushUrl(onUrlUpdate, () => result.current.onSortingChange([]));
expect(cleared.has("sort_by")).toBe(false);
expect(cleared.has("sort_order")).toBe(false);
expect(result.current.sorting).toEqual([{ id: "created_at", desc: true }]);
});
it("onSortingChange resolves a function updater against the current sort", async () => {
const { result, onUrlUpdate } = renderTableState({ sort_by: "spend" });
const url = await flushUrl(onUrlUpdate, () => result.current.onSortingChange(flipDirection));
expect(url.get("sort_by")).toBe("spend");
expect(url.get("sort_order")).toBe("asc");
expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]);
});
it("onColumnFiltersChange writes trimmed filter_<column> keys and returns to the first page", async () => {
const { result, onUrlUpdate } = renderTableState({ page: "3" });
const url = await flushUrl(onUrlUpdate, () =>
result.current.onColumnFiltersChange([{ id: "team_id", value: " team-1 " }]),
);
expect(url.get("filter_team_id")).toBe("team-1");
expect(url.has("page")).toBe(false);
expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]);
});
it("onColumnFiltersChange removes the key for an empty value and for a filter no longer present", async () => {
const { result, onUrlUpdate } = renderTableState({ filter_team_id: "team-1", filter_user_id: "user-42" });
const url = await flushUrl(onUrlUpdate, () => result.current.onColumnFiltersChange([{ id: "team_id", value: "" }]));
expect(url.has("filter_team_id")).toBe(false);
expect(url.has("filter_user_id")).toBe(false);
expect(result.current.columnFilters).toEqual([]);
});
it("onColumnFiltersChange ignores a non-string filter value", async () => {
const { result, onUrlUpdate } = renderTableState({ filter_team_id: "team-1" });
const url = await flushUrl(onUrlUpdate, () =>
result.current.onColumnFiltersChange([{ id: "team_id", value: ["team-1", "team-2"] }]),
);
expect(url.has("filter_team_id")).toBe(false);
});
it("onColumnFiltersChange resolves a function updater against the current filters", async () => {
const { result, onUrlUpdate } = renderTableState({ filter_team_id: "team-1" });
const url = await flushUrl(onUrlUpdate, () =>
result.current.onColumnFiltersChange((previous) => [...previous, { id: "user_id", value: "user-42" }]),
);
expect(url.get("filter_team_id")).toBe("team-1");
expect(url.get("filter_user_id")).toBe("user-42");
});
it("writes prefixed and renamed keys only", async () => {
const { result, onUrlUpdate } = renderTableState(
{},
{ keyPrefix: "audit_", urlKeys: { search: "key_search", filter_team_id: "filter_team" } },
);
await flushUrl(onUrlUpdate, () => result.current.setSearch("prod"));
await flushUrl(onUrlUpdate, () => result.current.onSortingChange([{ id: "spend", desc: false }]));
const url = await flushUrl(onUrlUpdate, () =>
result.current.onColumnFiltersChange([{ id: "team_id", value: "team-1" }]),
);
expect(url.get("audit_key_search")).toBe("prod");
expect(url.get("audit_sort_by")).toBe("spend");
expect(url.get("audit_filter_team")).toBe("team-1");
expect([...url.keys()].filter((key) => !key.startsWith("audit_"))).toEqual([]);
expect(url.has("audit_search")).toBe(false);
expect(url.has("audit_filter_team_id")).toBe(false);
});
});
describe("referential stability", () => {
it("keeps the TanStack state and the page-clamp handler stable across rerenders while the URL is unchanged", () => {
const { result, rerender } = renderTableState({ page: "2", filter_team_id: "team-1", sort_by: "spend" });
const first = result.current;
rerender();
expect(result.current.sorting).toBe(first.sorting);
expect(result.current.pagination).toBe(first.pagination);
expect(result.current.columnFilters).toBe(first.columnFilters);
expect(result.current.onPaginationChange).toBe(first.onPaginationChange);
});
it("hands out new pagination and untouched sorting after a page change", async () => {
const { result, onUrlUpdate } = renderTableState({ sort_by: "spend" });
const first = result.current;
await flushUrl(onUrlUpdate, () => result.current.onPaginationChange({ pageIndex: 4, pageSize: 50 }));
expect(result.current.pagination).not.toBe(first.pagination);
expect(result.current.pagination.pageIndex).toBe(4);
expect(result.current.sorting).toBe(first.sorting);
});
});

232
ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts Normal file
View file

@ -0,0 +1,232 @@
import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
import { createParser, Nullable, parseAsInteger, parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs";
import { useCallback, useMemo } from "react";
const SORT_ORDERS = ["asc", "desc"] as const;
type SortOrder = (typeof SORT_ORDERS)[number];
const STANDARD_KEYS = ["search", "sort_by", "sort_order", "page", "page_size"] as const;
type StandardKey = (typeof STANDARD_KEYS)[number];
type FilterStateKey<F extends string> = `filter_${F}`;
type StateKey<F extends string> = StandardKey | FilterStateKey<F>;
const MAX_PAGE = 100_000;
const DEFAULT_MAX_PAGE_SIZE = 100;
export interface UrlTableStateOptions<F extends string> {
sortFields: readonly string[];
defaultSort: { id: string; desc: boolean };
defaultPageSize: number;
maxPageSize?: number;
filterColumns: readonly F[];
keyPrefix?: string;
urlKeys?: Partial<Record<StateKey<F>, string>>;
}
export interface UrlTableState {
search: string;
setSearch: (value: string) => void;
sorting: SortingState;
onSortingChange: OnChangeFn<SortingState>;
pagination: PaginationState;
onPaginationChange: OnChangeFn<PaginationState>;
columnFilters: ColumnFiltersState;
onColumnFiltersChange: OnChangeFn<ColumnFiltersState>;
}
const boundedInteger = (min: number, max: number, fallback: number) =>
createParser({
parse: (value: string) => {
const parsed = parseAsInteger.parse(value);
return parsed === null ? null : Math.min(Math.max(parsed, min), max);
},
serialize: String,
}).withDefault(fallback);
const optionalString = parseAsString.withDefault("");
type OptionalStringParser = typeof optionalString;
const sortOrderParser = (fallback: SortOrder) => parseAsStringLiteral(SORT_ORDERS).withDefault(fallback);
interface StandardValues {
search: string;
sort_by: string;
sort_order: SortOrder;
page: number;
page_size: number;
}
type FilterValues<F extends string> = Record<FilterStateKey<F>, string>;
type StandardUpdate = Partial<Nullable<StandardValues>>;
type FilterUpdate<F extends string> = Record<FilterStateKey<F>, string | null> & Pick<Nullable<StandardValues>, "page">;
type SetTableValues<F extends string> = (update: StandardUpdate | FilterUpdate<F> | null) => Promise<URLSearchParams>;
interface TableQueryState<F extends string> {
values: StandardValues;
filters: FilterValues<F>;
setValues: SetTableValues<F>;
}
type TableParsers<F extends string> = {
search: OptionalStringParser;
sort_by: OptionalStringParser;
sort_order: ReturnType<typeof sortOrderParser>;
page: ReturnType<typeof boundedInteger>;
page_size: ReturnType<typeof boundedInteger>;
} & Record<FilterStateKey<F>, OptionalStringParser>;
const useTableQueryStates = <F extends string>(
parsers: TableParsers<F>,
urlKeys: Record<StateKey<F>, string>,
): TableQueryState<F> => {
const [state, setState] = useQueryStates(parsers, { urlKeys });
return useMemo(
() => ({
values: state as StandardValues,
filters: state as FilterValues<F>,
setValues: setState as SetTableValues<F>,
}),
[state, setState],
);
};
const filterStateKey = <F extends string>(column: F): FilterStateKey<F> => `filter_${column}`;
const filterParsers = <F extends string>(filterColumns: readonly F[]) =>
Object.fromEntries(filterColumns.map((column) => [filterStateKey(column), optionalString])) as Record<
FilterStateKey<F>,
OptionalStringParser
>;
const resolveUrlKeys = <F extends string>(
filterColumns: readonly F[],
keyPrefix: string,
renamed: Partial<Record<StateKey<F>, string>>,
) => {
const stateKeys: readonly StateKey<F>[] = [
...STANDARD_KEYS,
...filterColumns.map((column) => filterStateKey(column)),
];
return Object.fromEntries(stateKeys.map((key) => [key, `${keyPrefix}${renamed[key] ?? key}`])) as Record<
StateKey<F>,
string
>;
};
const filterValue = (filters: ColumnFiltersState, column: string): string | null => {
const value = filters.find((filter) => filter.id === column)?.value;
return (typeof value === "string" ? value.trim() : "") || null;
};
const filterUpdates = <F extends string>(filterColumns: readonly F[], filters: ColumnFiltersState) =>
Object.fromEntries(filterColumns.map((column) => [filterStateKey(column), filterValue(filters, column)])) as Record<
FilterStateKey<F>,
string | null
>;
const toSortOrder = (active: SortingState[number]): SortOrder => (active.desc ? "desc" : "asc");
export function useUrlTableState<F extends string>(options: UrlTableStateOptions<F>): UrlTableState {
const {
sortFields,
defaultSort,
defaultPageSize,
maxPageSize = DEFAULT_MAX_PAGE_SIZE,
filterColumns,
keyPrefix = "",
urlKeys: renamedKeys,
} = options;
const defaultSortId = defaultSort.id;
const defaultSortOrder: SortOrder = defaultSort.desc ? "desc" : "asc";
const parsers = useMemo<TableParsers<F>>(
() => ({
search: optionalString,
sort_by: parseAsString.withDefault(defaultSortId),
sort_order: sortOrderParser(defaultSortOrder),
page: boundedInteger(1, MAX_PAGE, 1),
page_size: boundedInteger(1, maxPageSize, defaultPageSize),
...filterParsers(filterColumns),
}),
[defaultSortId, defaultSortOrder, defaultPageSize, maxPageSize, filterColumns],
);
const urlKeys = useMemo(
() => resolveUrlKeys(filterColumns, keyPrefix, renamedKeys ?? {}),
[filterColumns, keyPrefix, renamedKeys],
);
const { values, filters, setValues } = useTableQueryStates(parsers, urlKeys);
const sortBy = sortFields.includes(values.sort_by) ? values.sort_by : defaultSortId;
const sortDesc = values.sort_order === "desc";
const sorting = useMemo<SortingState>(() => [{ id: sortBy, desc: sortDesc }], [sortBy, sortDesc]);
const pagination = useMemo<PaginationState>(
() => ({ pageIndex: values.page - 1, pageSize: values.page_size }),
[values.page, values.page_size],
);
const columnFilters = useMemo<ColumnFiltersState>(
() =>
filterColumns.flatMap((column) => {
const value = filters[filterStateKey(column)].trim();
return value ? [{ id: column, value }] : [];
}),
[filterColumns, filters],
);
const setSearch = useCallback(
(value: string) => {
void setValues({ search: value || null, page: null });
},
[setValues],
);
const onSortingChange = useCallback<OnChangeFn<SortingState>>(
(updaterOrValue) => {
const active = functionalUpdate(updaterOrValue, sorting)[0];
void setValues({
sort_by: active?.id ?? null,
sort_order: active ? toSortOrder(active) : null,
page: null,
});
},
[setValues, sorting],
);
const onPaginationChange = useCallback<OnChangeFn<PaginationState>>(
(updaterOrValue) => {
const next = functionalUpdate(updaterOrValue, pagination);
void setValues({ page: next.pageIndex + 1, page_size: next.pageSize });
},
[pagination, setValues],
);
const onColumnFiltersChange = useCallback<OnChangeFn<ColumnFiltersState>>(
(updaterOrValue) => {
const next = functionalUpdate(updaterOrValue, columnFilters);
void setValues({ ...filterUpdates(filterColumns, next), page: null });
},
[columnFilters, filterColumns, setValues],
);
return useMemo<UrlTableState>(
() => ({
search: values.search,
setSearch,
sorting,
onSortingChange,
pagination,
onPaginationChange,
columnFilters,
onColumnFiltersChange,
}),
[
values.search,
setSearch,
sorting,
onSortingChange,
pagination,
onPaginationChange,
columnFilters,
onColumnFiltersChange,
],
);
}

100
ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx Normal file
View file

@ -0,0 +1,100 @@
import { act, renderHook, waitFor } from "@testing-library/react";
import { NuqsTestingAdapter, type OnUrlUpdateFunction } from "nuqs/adapters/testing";
import type { ReactNode } from "react";
import { describe, expect, it, vi } from "vitest";
import { useUrlTab } from "./useUrlTab";
const TABS = ["chat", "compare", "compliance"] as const;
type Tab = (typeof TABS)[number];
interface RenderArgs {
searchParams?: string;
onUrlUpdate?: OnUrlUpdateFunction;
key?: string;
}
const initialProps: { values: readonly Tab[] } = { values: TABS };
const renderUrlTab = ({ searchParams, onUrlUpdate, key }: RenderArgs = {}) =>
renderHook(({ values }: { values: readonly Tab[] }) => useUrlTab(values, "chat", key), {
initialProps,
wrapper: ({ children }: { children: ReactNode }) => (
<NuqsTestingAdapter
searchParams={searchParams}
onUrlUpdate={onUrlUpdate}
hasMemory
resetUrlUpdateQueueOnMount={false}
>
{children}
</NuqsTestingAdapter>
),
});
const lastUrlUpdate = (onUrlUpdate: ReturnType<typeof vi.fn<OnUrlUpdateFunction>>) =>
onUrlUpdate.mock.calls.at(-1)?.[0];
describe("useUrlTab", () => {
it("reads the active tab from the URL", () => {
const { result } = renderUrlTab({ searchParams: "?tab=compare" });
expect(result.current[0]).toBe("compare");
});
it("resolves a URL value outside the allowed tabs to the fallback and drops it from the URL", async () => {
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
const { result } = renderUrlTab({ searchParams: "?tab=settings&other=1", onUrlUpdate });
expect(result.current[0]).toBe("chat");
await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled());
expect(lastUrlUpdate(onUrlUpdate)?.searchParams.has("tab")).toBe(false);
expect(lastUrlUpdate(onUrlUpdate)?.searchParams.get("other")).toBe("1");
});
it("leaves a URL that names an allowed tab untouched", async () => {
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderUrlTab({ searchParams: "?tab=compare", onUrlUpdate });
await new Promise((resolve) => setTimeout(resolve, 100));
expect(onUrlUpdate).not.toHaveBeenCalled();
});
it("reads from the caller's key instead of the default one", () => {
const { result } = renderUrlTab({ searchParams: "?view=compliance&tab=compare", key: "view" });
expect(result.current[0]).toBe("compliance");
});
it("writes ?tab= with history replace when a tab is selected", async () => {
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
const { result } = renderUrlTab({ onUrlUpdate });
act(() => result.current[1]("compare"));
await waitFor(() => expect(lastUrlUpdate(onUrlUpdate)?.searchParams.get("tab")).toBe("compare"));
expect(lastUrlUpdate(onUrlUpdate)?.options.history).toBe("replace");
expect(result.current[0]).toBe("compare");
});
it("removes the param when the fallback tab is selected", async () => {
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
const { result } = renderUrlTab({ searchParams: "?tab=compare", onUrlUpdate });
act(() => result.current[1]("chat"));
await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled());
expect(lastUrlUpdate(onUrlUpdate)?.searchParams.has("tab")).toBe(false);
expect(result.current[0]).toBe("chat");
});
it("falls back and clears the param when the current tab is no longer among the allowed values", async () => {
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
const { result, rerender } = renderUrlTab({ searchParams: "?tab=compliance", onUrlUpdate });
expect(result.current[0]).toBe("compliance");
rerender({ values: ["chat", "compare"] });
expect(result.current[0]).toBe("chat");
await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled());
expect(lastUrlUpdate(onUrlUpdate)?.searchParams.has("tab")).toBe(false);
});
});

12
ui/litellm-dashboard/src/hooks/useUrlTab.ts Normal file
View file

@ -0,0 +1,12 @@
import { parseAsString, useQueryState } from "nuqs";
import { useCallback, useEffect } from "react";
export function useUrlTab<T extends string>(values: readonly T[], fallback: T, key = "tab"): [T, (tab: T) => void] {
const [urlTab, setUrlTab] = useQueryState(key, parseAsString.withDefault(fallback));
const tab = values.find((value) => value === urlTab) ?? fallback;
useEffect(() => {
if (urlTab !== tab) void setUrlTab(null);
}, [urlTab, tab, setUrlTab]);
const setTab = useCallback((next: T) => void setUrlTab(next), [setUrlTab]);
return [tab, setTab];
}

47
ui/litellm-dashboard/src/utils/tabRoutes.test.ts
View file

@ -1,47 +0,0 @@
/* @vitest-environment jsdom */
import { describe, expect, it, vi } from "vitest";
vi.mock("@/components/networking", () => ({ serverRootPath: "" }));
import { createTabRoutes } from "./tabRoutes";
const routes = createTabRoutes("logs", ["audit", "deleted-keys", "deleted-teams"] as const);
describe("createTabRoutes.slugFromPathname", () => {
it("returns empty string for the base path with or without a trailing slash", () => {
expect(routes.slugFromPathname("/logs")).toBe("");
expect(routes.slugFromPathname("/logs/")).toBe("");
});
it("extracts the tab slug from dev and proxy-mounted (/ui) paths", () => {
expect(routes.slugFromPathname("/logs/audit")).toBe("audit");
expect(routes.slugFromPathname("/ui/logs/deleted-teams/")).toBe("deleted-teams");
});
it("returns the raw segment for an unknown tab so the caller can redirect to base", () => {
expect(routes.slugFromPathname("/ui/logs/bogus")).toBe("bogus");
});
it("returns empty string when the base segment is not in the path", () => {
expect(routes.slugFromPathname("/teams")).toBe("");
});
});
describe("createTabRoutes.tabHref", () => {
it("builds the trailing-slash base href for the empty slug", () => {
expect(routes.tabHref("")).toBe("/ui/logs/");
});
it("builds a trailing-slash href for every tab slug (required by static export)", () => {
for (const slug of routes.slugs) {
expect(routes.tabHref(slug)).toBe(`/ui/logs/${slug}/`);
}
});
});
describe("createTabRoutes metadata", () => {
it("preserves the base segment and slug tuple", () => {
expect(routes.baseSegment).toBe("logs");
expect(routes.slugs).toEqual(["audit", "deleted-keys", "deleted-teams"]);
});
});

26
ui/litellm-dashboard/src/utils/tabRoutes.ts
View file

@ -1,26 +0,0 @@
import { uiHref } from "@/utils/uiHref";
export interface TabRoutes<Slug extends string> {
baseSegment: string;
slugs: readonly Slug[];
tabHref: (slug: string) => string;
slugFromPathname: (pathname: string) => string;
}
export function createTabRoutes<Slug extends string>(baseSegment: string, slugs: readonly Slug[]): TabRoutes<Slug> {
const tabHref = (slug: string): string => {
const base = uiHref(baseSegment);
return slug ? `${base}/${slug}/` : `${base}/`;
};
const slugFromPathname = (pathname: string): string => {
const parts = pathname.split("/").filter(Boolean);
const idx = parts.indexOf(baseSegment);
if (idx === -1) {
return "";
}
return parts[idx + 1] ?? "";
};
return { baseSegment, slugs, tabHref, slugFromPathname };
}