diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index b04e004aa1a..f418752d990 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -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 | diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 34b712d47e3..180929bcfd4 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -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: diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 21e61c71fb7..1282e654365 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -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 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b1c38d0350a..92e5b1c4ff7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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, diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index c109da6f571..888a6d077ad 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -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. diff --git a/litellm/proxy/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index f3542098f95..a504c2ba102 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -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 diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py deleted file mode 100644 index eaf37b0bcf1..00000000000 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ /dev/null @@ -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 - ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index de6f9cb7647..974312ab2a5 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 99cd78e0b58..18e58861266 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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. diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b1c38d0350a..92e5b1c4ff7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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, diff --git a/tests/_live_test_helpers.py b/tests/_live_test_helpers.py index a79b81e82c1..629f8ac9fdb 100644 --- a/tests/_live_test_helpers.py +++ b/tests/_live_test_helpers.py @@ -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] diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 828227ed239..e57d33a0406 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -1,13 +1,17 @@ from __future__ import annotations +import base64 +import binascii +import json import os import shutil import socket import subprocess +import struct import threading import time import uuid -from collections.abc import Generator +from collections.abc import Generator, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, replace @@ -17,22 +21,68 @@ from typing import Final from urllib.parse import urlsplit import pytest +from pydantic import JsonValue from e2e_http import NetworkError, PreparedForward, RawResponse, StreamChunk, StreamHead, forward, prepare_forward -from models import LiteLLMParamsBody -from provider_cache import CacheEdge, CacheHit, CaptureLease, exact_key, successful_response +from models import LiteLLMParamsBody, ModelMode +from botocore.credentials import Credentials +from botocore.eventstream import EventStreamBuffer +from provider_cache import ( + SIGNATURE_HEADERS, + CacheEdge, + CacheHit, + CaptureLease, + MountPolicy, + ResponseStore, + cacheable_endpoint, + request_identity, + slotted_key, + successful_response, +) from provider_cache_redis import PUBLISH, RedisCommands, RedisResponseStore, configured_cache, redis_store -from provider_cache_routing import LIVE_PROVIDER_REQUIRED, route_cache_model -from provider_edge import configured_cache_backend, start_provider_edge +from provider_cache_routing import ( + BEDROCK_CROSS_REGION_PREFIX, + BEDROCK_EDGE_MODELS, + LIVE_PROVIDER_REQUIRED, + bedrock_region, + route_cache_model, +) +from fixture_mode import SESSION_TEST_KEY +from provider_edge import EDGE_MOUNTS, configured_cache_backend, resolve_mount, start_provider_edge +from provider_edge_bedrock import bedrock_signer from redis.exceptions import ConnectionError as RedisConnectionError SECRET: Final = b"synthetic-cache-hmac-key-for-tests" BODY: Final = b'{"model":"test","messages":[{"role":"user","content":"hello"}]}' SUCCESS: Final = b'{"id":"provider-fixed-id","choices":[{"message":{"content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}' HEADERS: Final = {"content-type": "application/json", "authorization": "Bearer synthetic-account-one"} +TEST_KEY: Final = "tests/e2e/synthetic_suite.py::TestCase::test_case" +OTHER_TEST_KEY: Final = "tests/e2e/synthetic_suite.py::TestCase::test_other_case" + + +def marked(marker: str) -> bytes: + """One request body shaped like the suite's own: a fixed prompt salted with a + 12-lowercase-hex ``unique_marker()`` token, fresh on every run.""" + return b'{"model":"test","messages":[{"role":"user","content":"hello %s"}]}' % marker.encode() + + +MARKED: Final = marked("0a1b2c3d4e5f") +BEDROCK_MOUNT: Final = "bedrock/us-east-1" +BEDROCK_MODEL: Final = "us.anthropic.claude-haiku-4-5-20251001-v1%3A0" +BEDROCK_BODY: Final = b'{"messages":[{"role":"user","content":[{"text":"hello 0a1b2c3d4e5f"}]}]}' +CONVERSE_SUCCESS: Final = ( + b'{"output":{"message":{"role":"assistant","content":[{"text":"hi"}]}},' + b'"stopReason":"end_turn","usage":{"inputTokens":1,"outputTokens":1,"totalTokens":2}}' +) +INVOKE_SUCCESS: Final = ( + b'{"id":"msg_synthetic","type":"message","role":"assistant",' + b'"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}' +) +STATIC_CREDENTIALS: Final = Credentials("AKIAIOSFODNN7EXAMPLE", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY") class Provider(ThreadingHTTPServer): hits: tuple[tuple[str, bytes], ...] = () + authorizations: tuple[str, ...] = () response: bytes = SUCCESS status: int = 200 delay: float = 0 @@ -49,6 +99,7 @@ class Handler(BaseHTTPRequestHandler): assert isinstance(server, Provider) body: Final = self.rfile.read(int(self.headers.get("content-length", "0"))) server.hits += ((self.path, body),) + server.authorizations += (self.headers.get("authorization", ""),) time.sleep(server.delay) self.send_response(server.status) if server.stream: @@ -122,6 +173,31 @@ def store(redis_url: str) -> RedisResponseStore: return redis_store(redis_url, "test-" + uuid.uuid4().hex) +def cache_edge(store: ResponseStore, test_key: str = TEST_KEY) -> CacheEdge: + """A cache edge standing in for one pytest process. A fresh instance over the + same store is the next build running the same test: the recordings survive, + the per-test FIFO slot counters start over.""" + return CacheEdge(store, SECRET, test_key=lambda: test_key) + + +def slot_key( + url: str, slot: int = 0, body: bytes | None = BODY, + headers: dict[str, str] = HEADERS, test_key: str = TEST_KEY, +) -> str: + prepared: Final = prepare_forward("POST", url, headers, body) + assert isinstance(prepared, PreparedForward) + return slotted_key(SECRET, request_identity(SECRET, test_key, "POST", url, prepared.headers, body), slot) + + +def bedrock_cache_edge(store: ResponseStore, test_key: str = TEST_KEY) -> CacheEdge: + return CacheEdge( + store, SECRET, test_key=lambda: test_key, + policies={BEDROCK_MOUNT: MountPolicy( + sign=bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS), unkeyed_headers=SIGNATURE_HEADERS, + )}, + ) + + @contextmanager def edge(cache: CacheEdge, provider: Provider) -> Generator[str, None, None]: upstream: Final = f"http://127.0.0.1:{provider.server_port}" @@ -132,34 +208,54 @@ def edge(cache: CacheEdge, provider: Provider) -> Generator[str, None, None]: running.shutdown() +@contextmanager +def bedrock_edge(cache: CacheEdge, provider: Provider, action: str = "converse") -> Generator[str, None, None]: + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + running: Final = start_provider_edge(cache, mounts={BEDROCK_MOUNT: upstream}) + try: + yield f"{running.edge.api_base(BEDROCK_MOUNT)}/model/{BEDROCK_MODEL}/{action}" + finally: + running.shutdown() + + def call(url: str, body: bytes = BODY, headers: dict[str, str] = HEADERS) -> RawResponse: result: Final = forward("POST", url, headers=headers, body=body, timeout=5) assert isinstance(result, RawResponse), result return result -def test_success_is_reusable_across_fresh_edges(store: RedisResponseStore, provider: Provider) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: +def test_repeated_call_takes_its_own_slot_and_both_replay_next_run( + store: RedisResponseStore, provider: Provider, +) -> None: + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS assert call(url).body == SUCCESS - with edge(CacheEdge(store, SECRET), provider) as other: + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as other: assert call(other).body == SUCCESS - assert len(provider.hits) == 1 + assert call(other).body == SUCCESS + assert len(provider.hits) == 2 @pytest.mark.parametrize("body", [BODY + b" ", BODY.replace(b"hello", b"Hello"), BODY.replace(b"test", b"test2")]) def test_any_body_change_calls_live(store: RedisResponseStore, provider: Provider, body: bytes) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: call(url, body) + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as url: call(url, body) assert len(provider.hits) == 2 @pytest.mark.parametrize("name,value", [("authorization", "Bearer another-account"), ("x-request-id", "one"), ("anthropic-version", "new")]) def test_changed_header_cannot_reuse(store: RedisResponseStore, provider: Provider, name: str, value: str) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: call(url, headers=HEADERS | {name: value}) call(url + "?x=1") assert len(provider.hits) == 3 @@ -169,36 +265,59 @@ def test_changed_header_cannot_reuse(store: RedisResponseStore, provider: Provid def test_failed_provider_responses_never_enter_cache(store: RedisResponseStore, provider: Provider, status: int, response: bytes) -> None: provider.status = status provider.response = response - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: assert call(url).status_code == status + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: assert call(url).body == response assert len(provider.hits) == 2 def test_cookie_setting_success_is_reused_without_the_cookie(store: RedisResponseStore, provider: Provider) -> None: provider.cookie = "__cf_bm=synthetic-bot-management; Path=/; HttpOnly; Secure" - with edge(CacheEdge(store, SECRET), provider) as url: - replies: Final = tuple(call(url) for _ in range(2)) + with edge(cache_edge(store), provider) as url: + live: Final = call(url) + with edge(cache_edge(store), provider) as url: + replayed: Final = call(url) assert len(provider.hits) == 1 - assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in replies) + assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in (live, replayed)) def test_expiry_does_not_slide(store: RedisResponseStore, provider: Provider) -> None: short: Final = replace(store, lifetime_ms=250) - with edge(CacheEdge(short, SECRET), provider) as url: - call(url) - call(url) - time.sleep(0.3) - call(url) - call(url) + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + + def drain() -> None: + head = cache_edge(short).forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS + + drain() + assert len(provider.hits) == 1 + drain() + assert len(provider.hits) == 1 + time.sleep(0.3) + drain() assert len(provider.hits) == 2 -def test_concurrent_requests_publish_atomically(store: RedisResponseStore, provider: Provider) -> None: +def test_concurrent_builds_publish_one_recording_atomically( + store: RedisResponseStore, provider: Provider, +) -> None: + """Five processes running the same test at the same time all reach slot 0 of + one key, which is the only way the capture lease is contended now that a + repeat inside a single test takes its own slot.""" provider.delay = 0.15 - with edge(CacheEdge(store, SECRET), provider) as url: - with ThreadPoolExecutor(max_workers=5) as executor: - replies: Final = tuple(executor.map(lambda _: call(url).body, range(5))) + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + edges: Final = tuple(cache_edge(store) for _ in range(5)) + + def drain(cache: CacheEdge) -> bytes: + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + return b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) + + with ThreadPoolExecutor(max_workers=5) as executor: + replies: Final = tuple(executor.map(drain, edges)) assert replies == (SUCCESS,) * 5 assert len(provider.hits) == 1 @@ -231,9 +350,9 @@ def test_stream_completion_controls_publication(store: RedisResponseStore, provi provider.stream = True provider.truncated = truncated provider.response = b'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n' - with edge(CacheEdge(store, SECRET), provider) as url: - for _ in range(2): - result: Final = forward("POST", url, headers=HEADERS, body=BODY, timeout=5) + for _ in range(2): + with edge(cache_edge(store), provider) as url: + result = forward("POST", url, headers=HEADERS, body=BODY, timeout=5) if truncated: assert isinstance(result, NetworkError) else: @@ -246,9 +365,9 @@ def test_store_outage_preserves_provider_success(provider: Provider) -> None: probe.bind(("127.0.0.1", 0)) port: Final = probe.getsockname()[1] unavailable: Final = redis_store(f"redis://127.0.0.1:{port}/0", "unavailable") - with edge(CacheEdge(unavailable, SECRET), provider) as url: - assert call(url).body == SUCCESS - assert call(url).body == SUCCESS + for _ in range(2): + with edge(cache_edge(unavailable), provider) as url: + assert call(url).body == SUCCESS assert len(provider.hits) == 2 @@ -267,7 +386,8 @@ def test_old_lease_cannot_overwrite_new_owner(store: RedisResponseStore) -> None def test_identity_preserves_values_and_never_contains_credentials() -> None: variants: Final = (b'{}', b'{"a":null}', b'{"a":false}', b'{"a":0}', b'{"a":0.0}', b'{"a":"0"}', b' { }', None, b'') - keys: Final = tuple(exact_key(SECRET, "POST", "https://example.invalid/v1/chat/completions", HEADERS, body) for body in variants) + url: Final = "https://example.invalid/v1/chat/completions" + keys: Final = tuple(request_identity(SECRET, TEST_KEY, "POST", url, HEADERS, body) for body in variants) assert len(set(keys)) == len(variants) assert all(len(key) == 64 and "synthetic-account" not in key for key in keys) @@ -275,21 +395,22 @@ def test_identity_preserves_values_and_never_contains_credentials() -> None: @pytest.mark.parametrize("payload", [b"corrupt response", '{"response":"{}","signature":"é"}'.encode()]) def test_corrupt_entry_is_replaced_by_same_successful_request(store: RedisResponseStore, provider: Provider, payload: bytes) -> None: upstream: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - prepared: Final = prepare_forward("POST", upstream, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", upstream, prepared.headers, BODY) + key: Final = slot_key(upstream) lease: Final = store.lookup(key) assert isinstance(lease, CaptureLease) assert store.publish(key, lease, payload) - cache: Final = CacheEdge(store, SECRET) - for _ in range(2): - head = cache.forward("POST", upstream, HEADERS, BODY, 5) + caches: Final = tuple(cache_edge(store) for _ in range(2)) + for cache in caches: + head = cache.forward("openai", "POST", upstream, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS assert len(provider.hits) == 1 - assert dict(cache.counters.counts) == { - "corrupt": 1, "misses": 1, "upstream_attempts": 1, "writes": 1, "hits": 1, + assert dict(caches[0].counters.counts) == { + "corrupt": 1, "mount:openai:corrupt": 1, "misses": 1, "mount:openai:misses": 1, + "upstream_attempts": 1, "mount:openai:upstream_attempts": 1, + "writes": 1, "mount:openai:writes": 1, } + assert dict(caches[1].counters.counts) == {"hits": 1, "mount:openai:hits": 1} @pytest.mark.parametrize("payload", [ @@ -301,22 +422,494 @@ def test_corrupt_entry_is_replaced_by_same_successful_request(store: RedisRespon def test_malformed_success_stream_is_never_cached(store: RedisResponseStore, provider: Provider, payload: bytes) -> None: provider.stream = True provider.response = payload - with edge(CacheEdge(store, SECRET), provider) as url: - assert call(url).body == payload - assert call(url).body == payload + for _ in range(2): + with edge(cache_edge(store), provider) as url: + assert call(url).body == payload assert len(provider.hits) == 2 +def test_requests_differing_only_by_marker_share_one_recording_per_slot( + store: RedisResponseStore, provider: Provider, +) -> None: + """The whole point of the canonical key. Every e2e test salts its prompt with + a fresh ``unique_marker()``, so before this the same test could never reuse + anything across builds. The second run mints markers it has never sent, which + is what a later build actually does, and must still serve both from the two + slots the first run recorded.""" + with edge(cache_edge(store), provider) as url: + assert call(url, MARKED).body == SUCCESS + assert call(url, marked("f5e4d3c2b1a0")).body == SUCCESS + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as url: + assert call(url, marked("7c6b5a493827")).body == SUCCESS + assert call(url, marked("1122334455ff")).body == SUCCESS + assert len(provider.hits) == 2 + + +@pytest.mark.parametrize("body", [ + b'{"model":"test","messages":[{"role":"user","content":"hello 0a1b2c3d4e5"}]}', + b'{"model":"test","messages":[{"role":"user","content":"hello 0a1b2c3d4e5f0"}]}', + b'{"model":"test","messages":[{"role":"user","content":"hello 0A1B2C3D4E5F"}]}', + b'{"model":"0a1b2c3d4e5f","messages":[{"role":"user","content":"hello"}]}', +]) +def test_a_token_that_is_not_a_marker_keeps_its_own_key( + store: RedisResponseStore, provider: Provider, body: bytes, +) -> None: + """Too short, too long, upper case, or in another field: none of these is the + 12-lowercase-hex token ``unique_marker`` mints, so none may fold onto it.""" + with edge(cache_edge(store), provider) as url: + call(url, MARKED) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: + call(url, body) + assert len(provider.hits) == 2 + + +def test_another_test_never_reuses_this_tests_recording( + store: RedisResponseStore, provider: Provider, +) -> None: + with edge(cache_edge(store), provider) as url: + call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store, OTHER_TEST_KEY), provider) as url: + call(url) + assert len(provider.hits) == 2 + with edge(cache_edge(store, OTHER_TEST_KEY), provider) as url: + call(url) + assert len(provider.hits) == 2 + + +def test_calls_outside_any_test_are_never_cached( + store: RedisResponseStore, provider: Provider, +) -> None: + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + cache: Final = CacheEdge(store, SECRET, test_key=lambda: SESSION_TEST_KEY) + for _ in range(2): + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS + assert len(provider.hits) == 2 + assert dict(cache.counters.counts) == { + "bypass": 2, "mount:openai:bypass": 2, + "upstream_attempts": 2, "mount:openai:upstream_attempts": 2, + } + + +def test_counters_attribute_every_outcome_to_its_mount( + store: RedisResponseStore, provider: Provider, +) -> None: + """The build report needs per-provider hit counts, and the flat totals cannot + supply them. Anthropic is served a chat-shaped body here, which its validator + rejects, so one mount writes and the other does not.""" + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + cache: Final = cache_edge(store) + running: Final = start_provider_edge(cache, mounts={"openai": upstream, "anthropic": upstream}) + try: + call(running.edge.api_base("openai") + "/v1/chat/completions") + call(running.edge.api_base("anthropic") + "/v1/messages") + finally: + running.shutdown() + counts: Final = dict(cache.counters.counts) + assert counts["misses"] == 2 + assert counts["mount:openai:misses"] == 1 and counts["mount:anthropic:misses"] == 1 + assert counts["mount:openai:writes"] == 1 and "mount:anthropic:writes" not in counts + assert counts["mount:anthropic:rejected"] == 1 and "mount:openai:rejected" not in counts + + +def test_a_rejection_says_whether_the_body_was_cut_short_or_simply_unfinished( + store: RedisResponseStore, provider: Provider, +) -> None: + """One `rejected` count cannot tell a connection that dropped from a body the + provider finished sending and the rules turned down, and those have opposite + fixes: the first is the client going away mid-capture, the second is a grammar + the cache does not accept. A mount whose rejections are mostly one or the other + is a different problem, so the report has to be able to say which.""" + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + cut_short: Final = cache_edge(store) + provider.stream = True + provider.truncated = True + provider.response = b'data: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n' + running: Final = start_provider_edge(cut_short, mounts={"openai": upstream}) + try: + forward("POST", running.edge.api_base("openai") + "/v1/chat/completions", + headers=HEADERS, body=MARKED, timeout=5) + finally: + running.shutdown() + + unfinished: Final = cache_edge(store) + provider.stream = False + provider.truncated = False + provider.response = b'{"choices":[{"index":0,"message":{"content":"hi"}}]}' + second: Final = start_provider_edge(unfinished, mounts={"openai": upstream}) + try: + call(second.edge.api_base("openai") + "/v1/chat/completions", MARKED) + finally: + second.shutdown() + + refused: Final = cache_edge(store) + provider.status = 429 + provider.response = b'{"message":"Too many requests"}' + third: Final = start_provider_edge(refused, mounts={"openai": upstream}) + try: + call(third.edge.api_base("openai") + "/v1/chat/completions", MARKED) + finally: + third.shutdown() + + cut: Final = dict(cut_short.counters.counts) + turned_down: Final = dict(unfinished.counters.counts) + errored: Final = dict(refused.counters.counts) + assert cut["mount:openai:rejected"] == turned_down["mount:openai:rejected"] == errored["mount:openai:rejected"] == 1 + assert cut["mount:openai:rejected_cut_short"] == 1 + assert turned_down["mount:openai:rejected_incomplete"] == 1 + assert errored["mount:openai:rejected_error_status"] == 1 + assert not {"mount:openai:rejected_incomplete", "mount:openai:rejected_error_status"} & set(cut) + assert not {"mount:openai:rejected_cut_short", "mount:openai:rejected_error_status"} & set(turned_down) + assert not {"mount:openai:rejected_cut_short", "mount:openai:rejected_incomplete"} & set(errored) + + +EMBEDDING_SUCCESS: Final = ( + b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2]}],' + b'"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' +) +RESPONSE_SUCCESS: Final = ( + b'{"id":"resp_synthetic","object":"response","status":"completed","error":null,' + b'"incomplete_details":null,"output":[]}' +) +RESPONSE_STREAM_SUCCESS: Final = ( + b'data: {"type":"response.created","response":{"id":"resp_synthetic","error":null}}\n\n' + b'data: {"type":"response.completed","response":{"id":"resp_synthetic","status":"completed"},"error":null}\n\n' +) + + +@contextmanager +def openai_edge(cache: CacheEdge, provider: Provider, path: str) -> Generator[str, None, None]: + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + running: Final = start_provider_edge(cache, mounts={"openai": upstream}) + try: + yield running.edge.api_base("openai") + path + finally: + running.shutdown() + + +class TestNonChatOpenAiEndpoints: + """Chat and messages were the only cacheable paths. Embeddings and responses + are the other two JSON endpoints the suite drives through the same mount, and + each needs its own completeness rule: a chat response's ``choices`` check + would reject a perfectly good embedding.""" + + @pytest.mark.parametrize("path,response", [ + ("/v1/embeddings", EMBEDDING_SUCCESS), + ("/v1/responses", RESPONSE_SUCCESS), + ]) + def test_complete_responses_replay_on_the_next_run( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 1 + + def test_a_completed_response_stream_replays( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + provider.stream = True + provider.response = RESPONSE_STREAM_SUCCESS + for _ in range(2): + with openai_edge(cache_edge(store), provider, "/v1/responses") as url: + assert call(url, MARKED).body == RESPONSE_STREAM_SUCCESS + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("path,response", [ + ("/v1/embeddings", b'{"object":"list","data":[],"usage":{"prompt_tokens":0}}'), + ("/v1/embeddings", b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[]}],"usage":{}}'), + ("/v1/embeddings", b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1]}]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","status":"incomplete","output":[]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","status":"in_progress","output":[]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","output":[]}'), + ]) + def test_incomplete_bodies_never_enter_the_cache( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("payload", [ + b'data: {"type":"response.created","response":{"id":"resp_x"}}\n\n', + b'data: {"type":"response.created","response":{"id":"resp_x"}}\n\ndata: {"type":"response.failed"}\n\n', + b'data: {"type":"response.completed","response":{"id":"resp_x"}}\n\ndata: {"type":"response.created"}\n\n', + ]) + def test_a_response_stream_that_never_completed_is_never_cached( + self, store: RedisResponseStore, provider: Provider, payload: bytes, + ) -> None: + provider.stream = True + provider.response = payload + for _ in range(2): + with openai_edge(cache_edge(store), provider, "/v1/responses") as url: + assert call(url, MARKED).body == payload + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("path,response", [ + ("/v1/chat/completions", b'{"id":"x","error":null,"choices":[{"message":{"content":"hi"},' + b'"finish_reason":"stop"}]}'), + ("/v1/messages", b'{"id":"msg_x","type":"message","role":"assistant","error":null,' + b'"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}'), + ("/v1/responses", RESPONSE_SUCCESS), + ]) + def test_a_null_error_field_is_not_an_error( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + """Every OpenAI Responses body carries `error: null`, and testing the key's + presence rather than its value rejected all of them. The cost was silent: + nothing failed, the endpoint simply never cached.""" + assert b'"error":null' in response + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("path,response", [ + ("/v1/chat/completions", b'{"error":{"message":"rate limited","type":"rate_limit_error"}}'), + ("/v1/responses", b'{"object":"response","status":"completed","error":{"message":"bad"},"output":[]}'), + ]) + def test_a_populated_error_field_still_rejects( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("path,cacheable", [ + ("/v1/chat/completions", True), ("/v1/messages", True), + ("/v1/embeddings", True), ("/v1/responses", True), + ("/v1/audio/speech", False), ("/v1/images/generations", False), + ("/v1/files", False), ("/v1/batches", False), + ]) + def test_only_the_json_endpoints_are_cacheable(self, path: str, cacheable: bool) -> None: + assert cacheable_endpoint("openai", "POST", f"https://api.openai.com{path}", MARKED) is cacheable + + +BEDROCK_STREAM_MODEL: Final = "us.anthropic.claude-haiku-4-5-20251001-v1:0" +CONVERSE_STREAM_URL: Final = f"https://bedrock.invalid/model/{BEDROCK_STREAM_MODEL}/converse-stream" +INVOKE_STREAM_URL: Final = f"https://bedrock.invalid/model/{BEDROCK_STREAM_MODEL}/invoke-with-response-stream" + + +def eventstream_frame(headers: Mapping[str, str], payload: bytes) -> bytes: + """AWS eventstream wire framing, the shape `vnd.amazon.eventstream` bodies + arrive in. Built here rather than pasted from a capture so a test can express + the stream it means; `test_the_frames_these_tests_build_are_real_aws_framing` + holds it to botocore's own parser.""" + encoded: Final = b"".join( + bytes([len(name)]) + name.encode() + b"\x07" + struct.pack(">H", len(value)) + value.encode() + for name, value in headers.items() + ) + prelude: Final = struct.pack(">II", 16 + len(encoded) + len(payload), len(encoded)) + framed: Final = prelude + struct.pack(">I", binascii.crc32(prelude)) + encoded + payload + return framed + struct.pack(">I", binascii.crc32(framed)) + + +def eventstream_event(event_type: str, payload: JsonValue, message_type: str = "event") -> bytes: + return eventstream_frame( + {":event-type": event_type, ":message-type": message_type, ":content-type": "application/json"}, + json.dumps(payload).encode(), + ) + + +def invoke_chunk(inner: JsonValue) -> bytes: + return eventstream_event("chunk", {"bytes": base64.b64encode(json.dumps(inner).encode()).decode("ascii")}) + + +CONVERSE_STREAM_OK: Final = ( + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "hi"}}) + + eventstream_event("contentBlockStop", {"contentBlockIndex": 0}) + + eventstream_event("messageStop", {"stopReason": "end_turn"}) + + eventstream_event("metadata", {"usage": {"inputTokens": 12, "outputTokens": 6, "totalTokens": 18}}) +) +INVOKE_STREAM_OK: Final = ( + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x", "role": "assistant"}}) + + invoke_chunk({"type": "content_block_start", "index": 0}) + + invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"text": "hi"}}) + + invoke_chunk({"type": "content_block_stop", "index": 0}) + + invoke_chunk({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}) + + invoke_chunk({"type": "message_stop"}) +) + + +class TestBedrockSigning: + """Bedrock is the reason the edge could not mount it before: SigV4 covers the + Host header, so forwarding through a rewritten api_base invalidates the + proxy's signature. The edge mints its own over the upstream URL instead.""" + + def test_the_proxys_signature_is_replaced_not_forwarded(self) -> None: + signer: Final = bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS) + signed: Final = signer( + "POST", + f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/converse", + {"content-type": "application/json", "Authorization": "AWS4-HMAC-SHA256 Credential=PROXY/...", + "X-Amz-Date": "19700101T000000Z", "X-Amz-Security-Token": "proxy-session-token"}, + BEDROCK_BODY, + ) + assert "PROXY" not in str(signed) and "proxy-session-token" not in str(signed) + assert signed["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") + assert "/us-east-1/bedrock/aws4_request" in signed["Authorization"] + assert signed["X-Amz-Date"] != "19700101T000000Z" + assert signed["content-type"] == "application/json" + + def test_the_signed_url_reaches_the_wire_byte_for_byte(self) -> None: + """SigV4 hashes the canonical URI, so if the HTTP layer re-encoded the + colon in an inference-profile id after signing, every call would fail + with a signature mismatch rather than anything that names the cause.""" + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/converse" + signer: Final = bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS) + prepared: Final = prepare_forward("POST", url, signer("POST", url, dict(HEADERS), BEDROCK_BODY), BEDROCK_BODY) + assert isinstance(prepared, PreparedForward) + assert urlsplit(prepared.url).path == urlsplit(url).path + + def test_signature_headers_are_excluded_from_the_key( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + """A real signature is fresh on every call, so keying on it would make + every Bedrock request a permanent miss. The stub signer here varies its + stamp per call on purpose: the real one only varies once a second, which + would let this pass by luck when it should fail.""" + provider.response = CONVERSE_SUCCESS + stamps: Final = iter(("20260101T000000Z", "20260102T111111Z")) + + def varying(method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> dict[str, str]: + return dict(headers) | {"authorization": f"AWS4-HMAC-SHA256 {url}", "x-amz-date": next(stamps)} + + def signing_edge() -> CacheEdge: + return CacheEdge( + store, SECRET, test_key=lambda: TEST_KEY, + policies={BEDROCK_MOUNT: MountPolicy(sign=varying, unkeyed_headers=SIGNATURE_HEADERS)}, + ) + + for _ in range(2): + with bedrock_edge(signing_edge(), provider) as url: + assert call(url, BEDROCK_BODY).body == CONVERSE_SUCCESS + assert len(provider.hits) == 1 + assert provider.authorizations[0] == ( + f"AWS4-HMAC-SHA256 http://127.0.0.1:{provider.server_port}/model/{BEDROCK_MODEL}/converse" + ), "the signature must cover the upstream URL the edge calls, not the edge URL the proxy called" + + def test_a_mount_without_a_signer_still_keys_on_its_credentials( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + """The exclusion is per mount. Dropping authorization globally would let + one OpenAI account read another's recording.""" + cache: Final = bedrock_cache_edge(store) + assert "authorization" in SIGNATURE_HEADERS + assert "authorization" in cache.keyed("openai", HEADERS) + assert "authorization" not in cache.keyed(BEDROCK_MOUNT, HEADERS) + with edge(cache, provider) as url: + call(url) + with edge(bedrock_cache_edge(store), provider) as url: + call(url, headers=HEADERS | {"authorization": "Bearer synthetic-account-two"}) + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("action,response", [("converse", CONVERSE_SUCCESS), ("invoke", INVOKE_SUCCESS)]) + def test_complete_responses_replay_on_the_next_run( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("action,response", [ + ("converse", b'{"output":{"message":{}}}'), + ("converse", b'{"stopReason":"end_turn"}'), + ("converse", b'{"message":"The provided model identifier is invalid."}'), + ("converse", CONVERSE_SUCCESS[:-20]), + ("invoke", b'{"id":"msg_x","type":"message","content":[{"type":"text","text":"hi"}]}'), + ("invoke", b'{"id":"msg_x","type":"message","stop_reason":"end_turn"}'), + ("invoke", b'{"message":"Too many requests, please wait before trying again."}'), + ]) + def test_incomplete_or_error_bodies_never_enter_the_cache( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("action,response", [ + ("converse-stream", CONVERSE_STREAM_OK), + ("invoke-with-response-stream", INVOKE_STREAM_OK), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_finished_stream_is_served_from_the_cache_the_second_time( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + provider.response = response + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + replay: Final = bedrock_cache_edge(store) + with bedrock_edge(replay, provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + assert dict(replay.counters.counts)[f"mount:{BEDROCK_MOUNT}:hits"] == 1 + assert all( + sent.startswith("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") + for sent in provider.authorizations + ), provider.authorizations + + @pytest.mark.parametrize("action,response", [ + ("converse-stream", CONVERSE_STREAM_OK[:-1]), + ("invoke-with-response-stream", INVOKE_STREAM_OK[:-1]), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_stream_the_connection_cut_short_calls_the_provider_every_time( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + """The whole risk of caching an eventstream is recording a half-finished + one, so a truncated body has to be rejected rather than stored.""" + provider.response = response + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + replay: Final = bedrock_cache_edge(store) + with bedrock_edge(replay, provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 2 + assert dict(replay.counters.counts)[f"mount:{BEDROCK_MOUNT}:rejected"] == 1 + assert f"mount:{BEDROCK_MOUNT}:hits" not in dict(replay.counters.counts) + + @pytest.mark.parametrize("action", ["converse", "invoke", "converse-stream", "invoke-with-response-stream"]) + def test_every_anthropic_bedrock_action_is_cacheable(self, action: str) -> None: + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/{action}" + assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) + + @pytest.mark.parametrize("action", ["count-tokens", "invoke-async", "converse-stream-x"]) + def test_an_unknown_bedrock_action_is_not_cacheable(self, action: str) -> None: + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/{action}" + assert not cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) + + def test_a_region_mount_resolves_whole(self) -> None: + resolved: Final = resolve_mount(f"/{BEDROCK_MOUNT}/model/{BEDROCK_MODEL}/converse", EDGE_MOUNTS) + assert resolved is not None + assert resolved.mount == BEDROCK_MOUNT + assert resolved.upstream_base == "https://bedrock-runtime.us-east-1.amazonaws.com" + assert resolved.upstream_path == f"model/{BEDROCK_MODEL}/converse" + + def test_anthropic_stream_requires_start_finish_and_stop() -> None: start: Final = b'data: {"type":"message_start","message":{}}\n\n' finish: Final = b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"}}\n\n' stop: Final = b'data: {"type":"message_stop"}\n\n' url: Final = "https://example.invalid/v1/messages" headers: Final = {"content-type": "text/event-stream"} - assert successful_response(url, 200, headers, start + finish + stop) - assert not successful_response(url, 200, headers, start + stop) - assert not successful_response(url, 200, headers, finish + stop) - assert not successful_response(url, 200, headers, start + finish) + assert successful_response("anthropic", url, 200, headers, start + finish + stop) + assert not successful_response("anthropic", url, 200, headers, start + stop) + assert not successful_response("anthropic", url, 200, headers, finish + stop) + assert not successful_response("anthropic", url, 200, headers, start + finish) @pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")]) @@ -342,6 +935,93 @@ def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMPa assert route_cache_model(params, unexpected_edge, enabled=True) is params +@pytest.mark.parametrize("model,region", [ + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", None), + ("bedrock/converse/us.anthropic.claude-sonnet-5", None), + ("bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", None), + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "us-east-1"), + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "os.environ/AWS_REGION"), + ("bedrock/invoke/us.anthropic.claude-sonnet-5", "os.environ/AWS_REGION"), + ("bedrock/us.anthropic.claude-opus-4-7", "us-east-1"), + ("bedrock/converse/us.anthropic.claude-opus-4-7", "us-east-1"), +]) +def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str, region: str | None) -> None: + """Almost every Bedrock deployment in the suite declares its region as + `os.environ/AWS_REGION`, which only the proxy can resolve. Treating that + string as a region name would leave the whole Anthropic-on-Bedrock surface + off the edge, which is the point of mounting it at all.""" + params: Final = LiteLLMParamsBody(model=model, aws_region_name=region) + routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) + assert routed.aws_bedrock_runtime_endpoint == "http://edge.invalid/bedrock/us-east-1" + assert routed.api_base is None + assert routed.model_dump(exclude={"aws_bedrock_runtime_endpoint"}) == params.model_dump( + exclude={"aws_bedrock_runtime_endpoint"} + ) + + +@pytest.mark.parametrize("params", [ + LiteLLMParamsBody(model="bedrock/amazon.titan-embed-text-v2:0"), + LiteLLMParamsBody(model="bedrock/amazon.nova-canvas-v1:0"), + LiteLLMParamsBody(model="bedrock/amazon.nova-sonic-v1:0"), + LiteLLMParamsBody(model="bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_role_name="arn:aws:iam::1:role/x"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_access_key_id="AKIA"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", api_base="https://custom.invalid"), + LiteLLMParamsBody( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + aws_bedrock_runtime_endpoint="https://custom.invalid", + ), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_region_name="eu-west-1"), + LiteLLMParamsBody(model="bedrock/anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), + LiteLLMParamsBody(model="bedrock/invoke/eu.anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-opus-4-5", aws_region_name="us-east-1"), + LiteLLMParamsBody(model="bedrock/converse/us.anthropic.claude-haiku-9-9", aws_region_name="us-east-1"), +]) +def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(params: LiteLLMParamsBody) -> None: + """Non-Anthropic models the runner role cannot invoke, deployments carrying + their own AWS identity (routing those would replace the assume-role chain the + batch suite exists to prove), explicit endpoints, unmounted regions, and a + region only the proxy can resolve on a model that is not cross-region, whose + real region the harness cannot know.""" + routed: Final = route_cache_model( + params, lambda mount: None if mount not in EDGE_MOUNTS else f"http://edge.invalid/{mount}", enabled=True, + ) + assert routed is params + + +@pytest.mark.parametrize("declared,expected", [ + (None, "us-east-1"), + ("us-west-2", "us-west-2"), + ("eu-west-1", "eu-west-1"), + ("os.environ/AWS_REGION", "us-east-1"), + ("os.environ/ANY_OTHER_NAME", "us-east-1"), +]) +def test_a_region_only_the_proxy_can_resolve_falls_back_to_the_default_mount( + declared: str | None, expected: str, +) -> None: + """A declared literal region is the one the deployment meant. A region the + proxy resolves from its own environment is one the run pod cannot see, and + the default mount answers it.""" + assert bedrock_region(declared) == expected + + +def test_every_model_on_the_edge_allowlist_is_a_cross_region_profile() -> None: + """Answering an env-referenced region with the default mount is only correct + for a profile that fans out across the US regions and is reachable from any + of them. A single-region model on this list would be sent to a region it may + not exist in, so the list is where that is caught.""" + assert BEDROCK_EDGE_MODELS + assert all(model.startswith(BEDROCK_CROSS_REGION_PREFIX) for model in BEDROCK_EDGE_MODELS) + + +@pytest.mark.parametrize("mode", ["batch", "realtime", "image_generation"]) +def test_a_bedrock_deployment_with_a_mode_keeps_its_direct_route(mode: ModelMode) -> None: + params: Final = LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") + assert route_cache_model( + params, lambda mount: f"http://edge.invalid/{mount}", enabled=True, mode=mode, + ) is params + + def test_rollback_and_live_only_policy_keep_direct_provider_route() -> None: params: Final = LiteLLMParamsBody(model="openai/test") assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=False) is params @@ -366,14 +1046,16 @@ class PublishOutage: def test_write_outage_preserves_success_without_hidden_retry(store: RedisResponseStore, provider: Provider) -> None: unavailable: Final = replace(store, client=PublishOutage(store.client)) - cache: Final = CacheEdge(unavailable, SECRET) + cache: Final = cache_edge(unavailable) with edge(cache, provider) as url: assert call(url).body == SUCCESS assert call(url).body == SUCCESS assert len(provider.hits) == 2 assert dict(cache.counters.counts)["write_failures"] == 2 - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS + assert len(provider.hits) == 3 + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS assert len(provider.hits) == 3 @@ -382,45 +1064,41 @@ def test_connection_failure_releases_capture_lease(store: RedisResponseStore) -> with socket.socket() as unavailable: unavailable.bind(("127.0.0.1", 0)) url: Final = f"http://127.0.0.1:{unavailable.getsockname()[1]}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - assert isinstance(cache.forward("POST", url, HEADERS, BODY, 0.2), NetworkError) - prepared: Final = prepare_forward("POST", url, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY) - slot: Final = store.lookup(key) - assert isinstance(slot, CaptureLease) - assert store.release(key, slot) + cache: Final = cache_edge(store) + assert isinstance(cache.forward("openai", "POST", url, dict(HEADERS), BODY, 0.2), NetworkError) + key: Final = slot_key(url) + lease: Final = store.lookup(key) + assert isinstance(lease, CaptureLease) + assert store.release(key, lease) assert dict(cache.counters.counts)["rejected"] == 1 def test_close_before_first_chunk_releases_lease(store: RedisResponseStore, provider: Provider) -> None: url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - head: Final = cache.forward("POST", url, HEADERS, BODY, 5) + cache: Final = cache_edge(store) + head: Final = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) head.steps.close() - prepared: Final = prepare_forward("POST", url, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY) - slot: Final = store.lookup(key) - assert isinstance(slot, CaptureLease) - assert store.release(key, slot) + key: Final = slot_key(url) + lease: Final = store.lookup(key) + assert isinstance(lease, CaptureLease) + assert store.release(key, lease) def test_effective_account_change_cannot_reuse_cache( store: RedisResponseStore, provider: Provider, monkeypatch: pytest.MonkeyPatch, tmp_path, ) -> None: url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - for account in ("account-a", "account-b", "account-b"): + caches: Final = tuple(cache_edge(store) for _ in range(3)) + for account, cache in zip(("account-a", "account-b", "account-b"), caches, strict=True): netrc = tmp_path / account netrc.write_text(f"machine 127.0.0.1 login {account} password synthetic\n") monkeypatch.setenv("NETRC", str(netrc)) - head = cache.forward("POST", url, HEADERS, BODY, 5) + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS assert len(provider.hits) == 2 - assert dict(cache.counters.counts)["hits"] == 1 + assert dict(caches[2].counters.counts)["hits"] == 1 def test_enabled_environment_reuses_store_across_fresh_backends( @@ -449,7 +1127,7 @@ def test_enabled_environment_reuses_store_across_fresh_backends( def test_duplicate_headers_bypass_cache_and_count_live_calls( store: RedisResponseStore, provider: Provider, known_mount: bool, ) -> None: - cache: Final = CacheEdge(store, SECRET) + cache: Final = cache_edge(store) with edge(cache, provider) as url: parsed: Final = urlsplit(url) for _ in range(2): @@ -470,3 +1148,147 @@ def test_duplicate_headers_bypass_cache_and_count_live_calls( assert len(provider.hits) == (2 if known_mount else 0) assert dict(cache.counters.counts)["duplicate_header_bypass"] == 2 assert dict(cache.counters.counts).get("upstream_attempts", 0) == (2 if known_mount else 0) + + +class TestBedrockStreams: + def test_the_frames_these_tests_build_are_real_aws_framing(self) -> None: + buffer: Final = EventStreamBuffer() + buffer.add_data(CONVERSE_STREAM_OK) + assert [event.headers[":event-type"] for event in buffer] == [ + "messageStart", "contentBlockDelta", "contentBlockStop", "messageStop", "metadata", + ] + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + def test_a_finished_stream_is_recordable(self, url: str, body: bytes) -> None: + assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, b"{}") + assert successful_response(BEDROCK_MOUNT, url, 200, {}, body) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + @pytest.mark.parametrize("keep", [1, -1, -4]) + def test_a_stream_the_connection_cut_short_is_not_recordable( + self, url: str, body: bytes, keep: int, + ) -> None: + """botocore yields the frames it did receive and silently drops a trailing + partial one, so a stream cut a single byte short parses clean and only the + byte accounting and the terminator rule catch it.""" + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, body[:keep]) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + def test_a_corrupted_frame_is_not_recordable(self, url: str, body: bytes) -> None: + flipped: Final = bytearray(body) + flipped[len(body) // 2] ^= 0xFF + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, bytes(flipped)) + + def test_a_converse_stream_that_lost_its_usage_is_not_recordable(self) -> None: + """ConverseStream names its stop reason a frame before it reports usage, + and litellm prices the call from that usage, so a stream cut between the + two would replay as a free call.""" + without_metadata: Final = ( + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("messageStop", {"stopReason": "end_turn"}) + ) + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, without_metadata) + + def test_a_converse_stream_that_never_stopped_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}), + ) + + def test_a_stream_that_failed_after_answering_200_is_not_recordable(self) -> None: + """Bedrock reports a fault that began after the headers went out as an + exception frame in place of the terminator it never got to send.""" + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "hi"}}) + + eventstream_event("modelStreamErrorException", {"message": "boom"}, message_type="exception"), + ) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_stream_cut_after_its_terminator_is_not_recordable(self, url: str, body: bytes) -> None: + """The terminator rules cannot see this one. Every frame the stream owes + has arrived and the partial frame after them is the one botocore drops + without a word, so only counting the bytes against the frame lengths + tells this from a stream that ended where it meant to.""" + assert successful_response(BEDROCK_MOUNT, url, 200, {}, body) + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, body + b"\x00\x00\x02") + + def test_a_converse_stream_whose_stop_frame_names_no_reason_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("messageStop", {}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}), + ) + + def test_an_invoke_stream_carrying_a_frame_that_is_not_a_chunk_is_not_recordable(self) -> None: + """Every frame of an invoke stream is a `chunk` holding one base64 event. + A frame that is not one carries an event this rule cannot read, so the + stream can no longer be judged complete.""" + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}) + + invoke_chunk({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_a_frame_claiming_no_length_is_rejected_rather_than_walked_forever(self) -> None: + """A frame length of zero never advances the cursor. Rejecting it is what + keeps a corrupt body from spinning the edge instead of answering.""" + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, b"\x00\x00\x00\x00" * 4) + + @pytest.mark.parametrize("url,terminator", [ + (INVOKE_STREAM_URL, invoke_chunk({"type": "message_stop"})), + (CONVERSE_STREAM_URL, eventstream_event("metadata", {"usage": {"totalTokens": 18}})), + ], ids=["invoke-stream", "converse-stream"]) + def test_a_delta_that_names_no_stop_reason_does_not_finish_a_stream( + self, url: str, terminator: bytes, + ) -> None: + """A `message_delta` arriving without its stop reason is the shape of a + turn the connection cut short partway through the delta itself.""" + head: Final = ( + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + invoke_chunk({"type": "message_delta", "delta": {}}) + ) + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, head + terminator) + + def test_an_invoke_chunk_that_is_not_base64_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + eventstream_event("chunk", {"bytes": "not base64 at all !!"}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_an_invoke_stream_missing_its_stop_reason_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_an_empty_stream_is_not_recordable(self) -> None: + for url in (CONVERSE_STREAM_URL, INVOKE_STREAM_URL): + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, b"") + + def test_each_streaming_endpoint_is_held_to_its_own_grammar(self) -> None: + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, INVOKE_STREAM_OK) + assert not successful_response(BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, CONVERSE_STREAM_OK) + + def test_a_stream_that_errored_before_it_started_is_not_recordable(self) -> None: + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 503, {}, CONVERSE_STREAM_OK) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 8635c9ed9ae..d37b37eeba7 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -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 diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py new file mode 100644 index 00000000000..b7d330b7da6 --- /dev/null +++ b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py @@ -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"] diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index 447e8cc0bbb..a01d8ab3e7c 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -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//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, diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 829c84910a9..ac4cfb71407 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -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 " diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 896cb3e7efe..82ddb09f7f5 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -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")) diff --git a/tests/e2e/fixture_canonical.py b/tests/e2e/fixture_canonical.py index e76d63ca33b..019c011aa67 100644 --- a/tests/e2e/fixture_canonical.py +++ b/tests/e2e/fixture_canonical.py @@ -51,6 +51,9 @@ SECRET_FIELD_SUFFIXES: Final[tuple[str, ...]] = ( ) SECRET_PLACEHOLDER: Final = "" +MARKER_PATTERN: Final = re.compile(r"(?" + PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = ( (re.compile(r"(?"), ( @@ -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"), "", ), - (re.compile(r"(?"), + (MARKER_PATTERN, MARKER_PLACEHOLDER), ) diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index 44c416a3e78..09ec48daa2f 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -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 diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 7101438c5f8..7550bfdc150 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -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 diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 0c6eac75a43..22967869c78 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -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}") diff --git a/tests/e2e/provider_cache_redis.py b/tests/e2e/provider_cache_redis.py index be4e31b2c49..2c7419cfc0f 100644 --- a/tests/e2e/provider_cache_redis.py +++ b/tests/e2e/provider_cache_redis.py @@ -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) diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index 24599b5a313..f9775a2b152 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -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: diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index dda9e6f8e4f..2606b26fe99 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -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) diff --git a/tests/e2e/provider_edge_bedrock.py b/tests/e2e/provider_edge_bedrock.py new file mode 100644 index 00000000000..73e4a16d272 --- /dev/null +++ b/tests/e2e/provider_edge_bedrock.py @@ -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) diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index 1fdd3bd28ad..f6d23a3ec12 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -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 diff --git a/tests/e2e/quota_management/budgets/BUDGET_TEST_COVERAGE_MATRIX.md b/tests/e2e/quota_management/budgets/BUDGET_TEST_COVERAGE_MATRIX.md index 7ff920a8d6d..a07bdf3d4e9 100644 --- a/tests/e2e/quota_management/budgets/BUDGET_TEST_COVERAGE_MATRIX.md +++ b/tests/e2e/quota_management/budgets/BUDGET_TEST_COVERAGE_MATRIX.md @@ -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)** | diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py index 5d0c79f26f6..c8d70697182 100644 --- a/tests/e2e/test_provider_edge.py +++ b/tests/e2e/test_provider_edge.py @@ -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( diff --git a/tests/llm_translation/test_together_ai.py b/tests/llm_translation/test_together_ai.py index fd7ad40ed11..0b4e9d3952c 100644 --- a/tests/llm_translation/test_together_ai.py +++ b/tests/llm_translation/test_together_ai.py @@ -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""" diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index 43ed57f63af..25c6c50251d 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -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) diff --git a/tests/local_testing/test_text_completion.py b/tests/local_testing/test_text_completion.py index a814ce6d303..6808dfd768b 100644 --- a/tests/local_testing/test_text_completion.py +++ b/tests/local_testing/test_text_completion.py @@ -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, ) diff --git a/tests/logging_callback_tests/test_generic_api_callback.py b/tests/logging_callback_tests/test_generic_api_callback.py index 29d8f9e5694..d9853ebcb52 100644 --- a/tests/logging_callback_tests/test_generic_api_callback.py +++ b/tests/logging_callback_tests/test_generic_api_callback.py @@ -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 diff --git a/tests/logging_callback_tests/test_unit_test_litellm_logging.py b/tests/logging_callback_tests/test_unit_test_litellm_logging.py index 42ba4ff35f1..7709a823610 100644 --- a/tests/logging_callback_tests/test_unit_test_litellm_logging.py +++ b/tests/logging_callback_tests/test_unit_test_litellm_logging.py @@ -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 diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index ae41c74944d..72f6213e880 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -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 diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 7d1faf93116..bea9a38e9dd 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -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. diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py index 7eb7dc41d4f..1df8c96fb50 100644 --- a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py +++ b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py @@ -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 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 26ae28a57d2..6b8260e7f08 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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, diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py deleted file mode 100644 index 71671966d1a..00000000000 --- a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py +++ /dev/null @@ -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() diff --git a/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py index ec680317980..49bbd498cb9 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py @@ -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 # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 0fc961cf8c9..d6b9a782f99 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index 2d9c1bd8b46..dd3914e3ad5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -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(), diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 28ff4571b44..a3ff7f7447e 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -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 diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py index a97dcb41e44..40cb3f10d34 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py @@ -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, } diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index af89c424f8b..6e5cb7fcae3 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -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", diff --git a/tests/test_litellm/test_rate_limit_error_unification.py b/tests/test_litellm/test_rate_limit_error_unification.py index 99e9981857c..8241b29aff1 100644 --- a/tests/test_litellm/test_rate_limit_error_unification.py +++ b/tests/test_litellm/test_rate_limit_error_unification.py @@ -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.", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx deleted file mode 100644 index 24900bae798..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx +++ /dev/null @@ -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/"); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts deleted file mode 100644 index c17d71b4855..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts +++ /dev/null @@ -1,38 +0,0 @@ -import { useEffect } from "react"; -import { usePathname, useRouter } from "next/navigation"; -import type { TabRoutes } from "@/utils/tabRoutes"; - -interface UseTabRoutingArgs { - routes: Pick, "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 }; -} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx index 85e19d7d251..f12ebc0b831 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx @@ -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: () =>
, })); -describe("PlaygroundPage role guard", () => { - beforeEach(() => { - authState.userRole = "Admin"; - }); +const lastUrlUpdate = (onUrlUpdate: ReturnType>) => + 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(); + renderWithProviders(); 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(); + renderWithProviders(); 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(); + + expect(screen.getByRole("tab", { name: "Chat" })).toHaveAttribute("aria-selected", "true"); + }); + + it("activates the tab named in ?tab=", () => { + renderWithProviders(, { 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(, { 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(); + renderWithProviders(, { 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"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx index 78ca538d8b5..27a61415672 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx @@ -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(undefined); + const [activeTab, setActiveTab] = useUrlTab(PLAYGROUND_TABS, "chat"); useEffect(() => { const initializeProxySettings = async () => { @@ -48,7 +52,11 @@ export default function PlaygroundPage() { return (
- + Chat diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 6c742309eb7..4d963a2f603 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -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) => 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(, { 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(, { 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(); + renderWithProviders(, { + 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(); renderWithProviders(, { 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(); + renderWithProviders(, { 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(); renderWithProviders(, { 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(, { 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(, { + 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(); + + 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(); + 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(); + + 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(); + const view = renderWithProviders(, { 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(); + + 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(); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.anything()); + }); + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "page")).toBeNull(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 1f52bdd7335..39dd2cc5ab2 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -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 = { + 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(() => 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>( + (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( - () => [{ id: sortBy, desc: tableState.sort_order === "desc" }], - [sortBy, tableState.sort_order], - ); - const tablePagination = useMemo( - () => ({ 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( - () => - 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>( - (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>( - (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>( - (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) { refetch?.()} isRefreshing={isFetching} diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index 578e355b85d..3b2a17101ee 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -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(); }); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index cab073dc808..e77c8ba7e41 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -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; diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx index 7bbd4f918cd..c6c6a392aec 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx @@ -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[] = []; 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 = ; @@ -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 ); + +export const visibilityWithoutHandler = ( + // @ts-expect-error a controlled `columnVisibility` needs `onColumnVisibilityChange` or Columns-menu toggles are dropped + +); + +export const bothVisibilitySources = ( + // @ts-expect-error `defaultColumnVisibility` seeds uncontrolled visibility, so it cannot pair with a controlled `columnVisibility` + +); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 8ed8e392ae1..a7e4befa6ac 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -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({ pageIndex: initialPageIndex, pageSize: 10 }); const handleChange: OnChangeFn = (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(); + + 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({ pageIndex: initialPageIndex, pageSize: 2 }); + const handleChange: OnChangeFn = (updater) => { + const next = typeof updater === "function" ? updater(pagination) : updater; + onChange(next); + setPagination(next); + }; + return ( + + ); + } + + 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(); + + rerender(); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(names()).toEqual(["P2", "P3"]); + + rerender(); + 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(); + + 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(); + + 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( + , + ); + + await user.click(screen.getByTestId("pagination-next")); + expect(names()).toEqual(["P2", "P3"]); + + rerender( + , + ); + 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(); + + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(onChange).not.toHaveBeenCalled(); + + rerender(); + + 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( + } + />, + ); + + 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>(); + render( + } + />, + ); + + 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({ email: false }); + return ( + } + /> + ); + }; + render(); + + 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", () => { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index 26162a3f1f7..340f8d4f44f 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -425,7 +425,7 @@ function useControllable( return { value: internal, onChange: setInternal }; } -function useServerPageClamp( +function usePageClamp( active: boolean, rowCount: number | undefined, pagination: { value: PaginationState; onChange: OnChangeFn }, @@ -457,6 +457,7 @@ function useDataTableInstance( onPaginationChange, rowCount, isLoading = false, + isError, pageSizeOptions = DEFAULT_PAGE_SIZE_OPTIONS, filterMode = "none", columnFilters, @@ -466,6 +467,8 @@ function useDataTableInstance( onGlobalFilterChange, enableColumnResizing = false, columnResizeMode = "onEnd", + columnVisibility, + onColumnVisibilityChange, defaultColumnVisibility, getRowCanExpand, renderSubComponent, @@ -481,7 +484,6 @@ function useDataTableInstance( pageIndex: 0, pageSize: pageSizeOptions[0] ?? 25, }); - useServerPageClamp(paginationMode === "server" && !isLoading, rowCount, paginationState); const filterState = useControllable( columnFilters, onColumnFiltersChange, @@ -490,7 +492,11 @@ function useDataTableInstance( const globalFilterState = useControllable(globalFilter, onGlobalFilterChange, ""); const expandedState = useControllable(expanded, onExpandedChange, {}); const rowSelectionState = useControllable(rowSelection, onRowSelectionChange, {}); - const [columnVisibility, setColumnVisibility] = useState(defaultColumnVisibility ?? {}); + const columnVisibilityState = useControllable( + columnVisibility, + onColumnVisibilityChange, + defaultColumnVisibility ?? {}, + ); const [columnSizing, setColumnSizing] = useState({}); const columnPinning = React.useMemo(() => derivePinning(columns), [columns]); const expansionGuard = renderSubComponent !== undefined ? getRowCanExpand : undefined; @@ -505,7 +511,7 @@ function useDataTableInstance( globalFilter: globalFilterState.value, expanded: expandedState.value, rowSelection: rowSelectionState.value, - columnVisibility, + columnVisibility: columnVisibilityState.value, columnSizing, }, initialState: { columnPinning }, @@ -521,7 +527,7 @@ function useDataTableInstance( 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( ...(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 }; +}; + +function useSettledPageClamp(table: Table, 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(props: DataTableProps) { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/index.ts b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts index 39a887ba948..85cc5f287e0 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/index.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts @@ -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, diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts index c767a0a64c0..4529e3df164 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts @@ -27,6 +27,7 @@ export interface DataTableResolvedProps { getRowId?: (row: TData, index: number, parent?: Row) => string; isLoading?: boolean; + isError?: boolean; loadingMessage?: string; skeletonRowCount?: number; noDataMessage?: React.ReactNode; @@ -53,6 +54,8 @@ export interface DataTableResolvedProps { enableColumnResizing?: boolean; columnResizeMode?: ColumnResizeMode; + columnVisibility?: VisibilityState; + onColumnVisibilityChange?: OnChangeFn; defaultColumnVisibility?: VisibilityState; getRowCanExpand?: (row: Row) => boolean; @@ -96,6 +99,9 @@ type DataTableBaseProps = Omit< | "columnFilters" | "onColumnFiltersChange" | "defaultColumnFilters" + | "columnVisibility" + | "onColumnVisibilityChange" + | "defaultColumnVisibility" | "rowSelection" | "onRowSelectionChange" >; @@ -142,6 +148,18 @@ type FilterProps = defaultColumnFilters?: ColumnFiltersState; }; +type ColumnVisibilityProps = + | { + columnVisibility: VisibilityState; + onColumnVisibilityChange: OnChangeFn; + defaultColumnVisibility?: never; + } + | { + columnVisibility?: never; + onColumnVisibilityChange?: never; + defaultColumnVisibility?: VisibilityState; + }; + type RowSelectionProps = | { rowSelection: RowSelectionState; onRowSelectionChange: OnChangeFn } | { rowSelection?: never; onRowSelectionChange?: OnChangeFn }; @@ -150,4 +168,5 @@ export type DataTableProps = DataTableBaseProps `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 }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts new file mode 100644 index 00000000000..3cbf5c2a000 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts @@ -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(); + +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 } { + const key = storageKey(tableId); + const raw = useSyncExternalStore( + subscribe, + () => readRaw(key), + () => null, + ); + const columnVisibility = useMemo(() => parseVisibility(raw, defaults), [raw, defaults]); + + const onColumnVisibilityChange = useCallback>( + (updater) => { + const next = typeof updater === "function" ? updater(parseVisibility(readRaw(key), defaults)) : updater; + writeRaw(key, JSON.stringify(next)); + }, + [key, defaults], + ); + + return { columnVisibility, onColumnVisibilityChange }; +} diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx new file mode 100644 index 00000000000..ad46d18b5d8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx @@ -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 = { + 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 = {}, + overrides: Partial> = {}, +) => { + const onUrlUpdate = vi.fn(); + const options = { ...BASE_OPTIONS, ...overrides }; + const hook = renderHook(() => useUrlTableState(options), { + wrapper: withNuqsTestingAdapter({ searchParams, onUrlUpdate, hasMemory: true }), + }); + return { ...hook, onUrlUpdate }; +}; + +const lastUrl = (onUrlUpdate: Mock) => { + 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, 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_ 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_ 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); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts new file mode 100644 index 00000000000..1a423663936 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts @@ -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 = `filter_${F}`; +type StateKey = StandardKey | FilterStateKey; + +const MAX_PAGE = 100_000; +const DEFAULT_MAX_PAGE_SIZE = 100; + +export interface UrlTableStateOptions { + sortFields: readonly string[]; + defaultSort: { id: string; desc: boolean }; + defaultPageSize: number; + maxPageSize?: number; + filterColumns: readonly F[]; + keyPrefix?: string; + urlKeys?: Partial, string>>; +} + +export interface UrlTableState { + search: string; + setSearch: (value: string) => void; + sorting: SortingState; + onSortingChange: OnChangeFn; + pagination: PaginationState; + onPaginationChange: OnChangeFn; + columnFilters: ColumnFiltersState; + onColumnFiltersChange: OnChangeFn; +} + +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 = Record, string>; +type StandardUpdate = Partial>; +type FilterUpdate = Record, string | null> & Pick, "page">; +type SetTableValues = (update: StandardUpdate | FilterUpdate | null) => Promise; + +interface TableQueryState { + values: StandardValues; + filters: FilterValues; + setValues: SetTableValues; +} + +type TableParsers = { + search: OptionalStringParser; + sort_by: OptionalStringParser; + sort_order: ReturnType; + page: ReturnType; + page_size: ReturnType; +} & Record, OptionalStringParser>; + +const useTableQueryStates = ( + parsers: TableParsers, + urlKeys: Record, string>, +): TableQueryState => { + const [state, setState] = useQueryStates(parsers, { urlKeys }); + return useMemo( + () => ({ + values: state as StandardValues, + filters: state as FilterValues, + setValues: setState as SetTableValues, + }), + [state, setState], + ); +}; + +const filterStateKey = (column: F): FilterStateKey => `filter_${column}`; + +const filterParsers = (filterColumns: readonly F[]) => + Object.fromEntries(filterColumns.map((column) => [filterStateKey(column), optionalString])) as Record< + FilterStateKey, + OptionalStringParser + >; + +const resolveUrlKeys = ( + filterColumns: readonly F[], + keyPrefix: string, + renamed: Partial, string>>, +) => { + const stateKeys: readonly StateKey[] = [ + ...STANDARD_KEYS, + ...filterColumns.map((column) => filterStateKey(column)), + ]; + return Object.fromEntries(stateKeys.map((key) => [key, `${keyPrefix}${renamed[key] ?? key}`])) as Record< + StateKey, + 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 = (filterColumns: readonly F[], filters: ColumnFiltersState) => + Object.fromEntries(filterColumns.map((column) => [filterStateKey(column), filterValue(filters, column)])) as Record< + FilterStateKey, + string | null + >; + +const toSortOrder = (active: SortingState[number]): SortOrder => (active.desc ? "desc" : "asc"); + +export function useUrlTableState(options: UrlTableStateOptions): 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>( + () => ({ + 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(() => [{ id: sortBy, desc: sortDesc }], [sortBy, sortDesc]); + + const pagination = useMemo( + () => ({ pageIndex: values.page - 1, pageSize: values.page_size }), + [values.page, values.page_size], + ); + + const columnFilters = useMemo( + () => + 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>( + (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>( + (updaterOrValue) => { + const next = functionalUpdate(updaterOrValue, pagination); + void setValues({ page: next.pageIndex + 1, page_size: next.pageSize }); + }, + [pagination, setValues], + ); + + const onColumnFiltersChange = useCallback>( + (updaterOrValue) => { + const next = functionalUpdate(updaterOrValue, columnFilters); + void setValues({ ...filterUpdates(filterColumns, next), page: null }); + }, + [columnFilters, filterColumns, setValues], + ); + + return useMemo( + () => ({ + search: values.search, + setSearch, + sorting, + onSortingChange, + pagination, + onPaginationChange, + columnFilters, + onColumnFiltersChange, + }), + [ + values.search, + setSearch, + sorting, + onSortingChange, + pagination, + onPaginationChange, + columnFilters, + onColumnFiltersChange, + ], + ); +} diff --git a/ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx b/ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx new file mode 100644 index 00000000000..427d454ee2d --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx @@ -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 }) => ( + + {children} + + ), + }); + +const lastUrlUpdate = (onUrlUpdate: ReturnType>) => + 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(); + 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(); + 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(); + 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(); + 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(); + 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); + }); +}); diff --git a/ui/litellm-dashboard/src/hooks/useUrlTab.ts b/ui/litellm-dashboard/src/hooks/useUrlTab.ts new file mode 100644 index 00000000000..2f3d705c610 --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useUrlTab.ts @@ -0,0 +1,12 @@ +import { parseAsString, useQueryState } from "nuqs"; +import { useCallback, useEffect } from "react"; + +export function useUrlTab(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]; +} diff --git a/ui/litellm-dashboard/src/utils/tabRoutes.test.ts b/ui/litellm-dashboard/src/utils/tabRoutes.test.ts deleted file mode 100644 index 402be55c33a..00000000000 --- a/ui/litellm-dashboard/src/utils/tabRoutes.test.ts +++ /dev/null @@ -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"]); - }); -}); diff --git a/ui/litellm-dashboard/src/utils/tabRoutes.ts b/ui/litellm-dashboard/src/utils/tabRoutes.ts deleted file mode 100644 index 4af2b983cba..00000000000 --- a/ui/litellm-dashboard/src/utils/tabRoutes.ts +++ /dev/null @@ -1,26 +0,0 @@ -import { uiHref } from "@/utils/uiHref"; - -export interface TabRoutes { - baseSegment: string; - slugs: readonly Slug[]; - tabHref: (slug: string) => string; - slugFromPathname: (pathname: string) => string; -} - -export function createTabRoutes(baseSegment: string, slugs: readonly Slug[]): TabRoutes { - 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 }; -}