diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_add_mcp_server_oauth_client_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_add_mcp_server_oauth_client_table/migration.sql new file mode 100644 index 00000000000..7aa6cdb1e33 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_add_mcp_server_oauth_client_table/migration.sql @@ -0,0 +1,9 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_MCPServerOAuthClient" ( + "server_id" TEXT NOT NULL, + "credentials" JSONB, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_MCPServerOAuthClient_pkey" PRIMARY KEY ("server_id") +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index f842bf13da9..a99cec49417 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -396,6 +396,13 @@ model LiteLLM_MCPUserEnvVars { @@index([server_id]) } +model LiteLLM_MCPServerOAuthClient { + server_id String @id + credentials Json? + created_at DateTime @default(now()) @map("created_at") + updated_at DateTime @default(now()) @updatedAt @map("updated_at") +} + // Generate Tokens for Proxy model LiteLLM_VerificationToken { token String @id diff --git a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md new file mode 100644 index 00000000000..c7980b11147 --- /dev/null +++ b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md @@ -0,0 +1,53 @@ +# Provider coding standards (litellm-rust) + +Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MISTRAL_OCR_CONFIG`) is the reference; `messages` (`ANTHROPIC_MESSAGES_CONFIG`) is the next port. + +## Provider resolution + +1. Always resolve the provider/model first with `get_custom_llm_provider` (`core/src/routing_utils/provider.rs`). Nothing downstream may branch on a raw model string. +2. Model/provider is resolved once, in `prepare.rs`, and passed down as typed fields. Don't re-resolve or re-parse it in transforms or handlers. + +## Transforms and the base config + +3. Every route defines a base config trait with `transform_request` + `transform_response` (+ `complete_url`, `supported_params`), living in `core/src//transformation.rs` (e.g. `AnthropicMessagesProviderConfig`, mirroring `OcrProviderConfig`). +4. Each provider implements that trait as a `const __CONFIG` in `core/src/providers///transformation.rs`, mirroring the Python provider tree. +5. Individual configs implement only the request/response transforms. Shared behavior (param filtering, defaults) stays as trait default methods so future providers inherit existing logic instead of reimplementing it. +6. Prefer composition: a provider that extends another reuses the base trait's defaults or wraps another config; don't copy transform bodies between providers. + +## Boundaries + +7. Layers never cross: `core` = pure transforms/types (no network, env, secrets, auth, logging, global mutable state); `ai-gateway` = all I/O, auth headers, HTTP/SSE, lifecycle hooks; `python-bridge` = thin PyO3 adapter. +8. Generic/route files contain zero provider-specific branches. A provider is one module under `core/src/providers///`; a route is a module, never a new crate. +9. Route entry point stays thin: `()` -> `prepare_*` -> `CallLifecycle::run_request`, which owns the pre_call -> during_call -> provider call -> success/failure order and phase timing. Handlers validate and delegate; no business logic in them. +10. Constants (URLs, env-var names, API versions, error messages) live in a crate `constants.rs`, never inline. Env reads happen only at the host/config layer, with the `DEFAULT_*` fallback defined in `constants.rs`. + +## Types and errors + +11. Typed contracts only: no bare `serde_json::Value` / `String` / `Vec` as a transform input or output. Parse wire bytes into typed structs/enums at the host edge; a `type` discriminator is a typed field, not a raw string. +12. Model failures as values: return typed `CoreError`, don't panic. No `unwrap`/`expect`/`panic!` on user or provider input. +13. No mutation: build values in one shot (comprehensions/iterators, `collect`), prefer immutable bindings and owned typed structs over seeding-and-mutating. +14. Early returns over deep nesting; small focused files over god modules. +15. Preserve Python output shape intentionally. If a field is always serialized as `null` for parity, keep it and pin it with a test. + +## Safety and data minimization + +16. Never log request/response bodies, base64 payloads, document contents, or secrets. Truncate and bound any upstream body before it crosses a host boundary. +17. Treat empty/whitespace credentials, URLs, and config values as absent at the host resolution layer. +18. Host I/O sets connect + request timeouts (no unbounded waits), reuses a shared HTTP client, and prefers rustls TLS. + +## Tests and rollout + +19. Every provider transform ships tests for: supported-param filtering, request body shape, response normalization, missing/null fields, bad input, and `*_match_python` fixture parity. +20. Lifecycle/hook tests cover hook order, success + failure callback payloads, pre-call guardrail blocking before any provider I/O, during-call body mutation, and provider-error mapping. +21. Rust paths stay off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. + +## Checks before push + +22. Run, and keep green: + ```bash + cd litellm-rust + cargo fmt --check + cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings + cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings + cargo test --workspace + ``` diff --git a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs index a56d19b8242..656ba033b62 100644 --- a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs +++ b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs @@ -62,12 +62,17 @@ fn parse_members(manifest: &str) -> BTreeSet { members } -/// The immediate subdirectory names under `crates/`. +/// The crate subdirectory names under `crates/`. +/// +/// A directory counts as a crate only when it holds a `Cargo.toml`; non-crate +/// directories (e.g. docs like `CODING_STANDARDS/`) are ignored so they can live +/// under `crates/` without tripping the crate-proliferation guard. fn crate_dirs(root: &Path) -> BTreeSet { fs::read_dir(root.join("crates")) .expect("crates/ directory should exist") .filter_map(Result::ok) .filter(|entry| entry.file_type().map(|ty| ty.is_dir()).unwrap_or(false)) + .filter(|entry| entry.path().join("Cargo.toml").is_file()) .map(|entry| entry.file_name().to_string_lossy().into_owned()) .collect() } diff --git a/litellm/constants.py b/litellm/constants.py index e104c937a9b..6432e2176c7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1517,6 +1517,12 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ "cost_discount_config", "cost_margin_config", "budget_exceeded_throttle_percentage", + # Every field editable from the Admin UI (proxy_server._GENERAL_SETTINGS_UI_LITELLM_FIELDS) + # must be listed here so a DB write from one worker overrides the live litellm attribute on + # the others when config reloads; otherwise peer workers stay on their startup value. + # test_general_settings_ui_fields_are_db_overridable enforces that pairing. + "enable_anthropic_prompt_caching", + "anthropic_prompt_caching_ttl", ] SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) diff --git a/litellm/integrations/otel/emitter.py b/litellm/integrations/otel/emitter.py index f97f8b8394c..8651cf586cd 100644 --- a/litellm/integrations/otel/emitter.py +++ b/litellm/integrations/otel/emitter.py @@ -72,6 +72,42 @@ def _stamp_litellm_error_attributes(span: Span, error: SpanError) -> None: span.set_attribute(LiteLLMError.LLM_PROVIDER, error.llm_provider) +def stamp_error( + span: Span, + error: SpanError, + *, + record_event: bool = True, + set_status: bool = True, +) -> tuple[str, str] | None: + """Stamp the full v2 error attribute set on ``span`` and return the resolved + ``(error_type, message)`` pair, or ``None`` when the error carries neither a + type nor a message. + + Shared by the LLM-call span (``finish_span``) and the proxy-level failure + spans (the FastAPI SERVER span and the ``auth`` phase span) so every v2 error + span carries identical keys. The semconv ``exception`` event rides alongside + the attributes so backends that map unknown string attrs to a truncated + ``keyword`` (e.g. Elasticsearch's 1024-char ``ignore_above``) still see the + full untruncated message on the recognized event field. ``record_event`` and + ``set_status`` are opt-outs for callers whose span lifecycle (``use_span``) or + owner (the FastAPI instrumentor) already records the event or the status. + """ + if not (error.error_type or error.message): + return None + error_type = error.error_type or "error" + message = error.message or error.error_type or "error" + _stamp_otel_error_attributes(span, error_type, message) + _stamp_litellm_error_attributes(span, error) + if set_status: + span.set_status(Status(StatusCode.ERROR, message)) + if record_event: + span.add_event( + ExceptionEvent.NAME, + {ExceptionEvent.TYPE: error_type, ExceptionEvent.MESSAGE: message}, + ) + return error_type, message + + class SpanEmitter: def __init__( self, @@ -212,21 +248,10 @@ class SpanEmitter: ) else None ) - if error and (error.error_type or error.message): - error_type = error.error_type or "error" - message = error.message or error.error_type or "error" - _stamp_otel_error_attributes(span, error_type, message) - _stamp_litellm_error_attributes(span, error) - span.set_status(Status(StatusCode.ERROR, message)) - # Also emit the semconv ``exception`` event so backends that - # dynamic-map unknown string span attrs to ``keyword`` (e.g. - # Elasticsearch with a 1024-char ``ignore_above``) still see the - # full untruncated message on the recognized event field. - span.add_event( - ExceptionEvent.NAME, - {ExceptionEvent.TYPE: error_type, ExceptionEvent.MESSAGE: message}, - ) - if self._event_recorder is not None and role is SpanRole.LLM_CALL: + if error: + stamped = stamp_error(span, error) + if stamped is not None and self._event_recorder is not None and role is SpanRole.LLM_CALL: + error_type, message = stamped self._event_recorder.record_operation_exception( span_context=span.get_span_context(), error_type=error_type, diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index be72fabd387..778f5342e90 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -24,7 +24,7 @@ from litellm.integrations.otel.plumbing.context import ( set_request_baggage, set_request_root_span, ) -from litellm.integrations.otel.emitter import SpanEmitter +from litellm.integrations.otel.emitter import SpanEmitter, stamp_error from litellm.integrations.otel.mappers import resolve_mappers from litellm.integrations.otel.model.metadata import ( LLMCallEvent, @@ -59,6 +59,7 @@ from litellm.integrations.otel.model.spans import SpanRole, span_role_for_servic from litellm.integrations.otel.model.utils import to_ns if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import ( StandardLoggingGuardrailInformation, StandardLoggingPayload, @@ -66,6 +67,33 @@ if TYPE_CHECKING: LITELLM_TRACER_NAME = "litellm" + +def _span_error_from_exception( + exception: "Exception | None", + *, + status_code: int | None = None, + traceback_str: str | None = None, +) -> SpanError: + """A ``SpanError`` for a proxy-level failure that never produced a + ``StandardLoggingPayload`` (auth / validation / malformed-body rejections), + mirroring ``_parse_error``'s field mapping so it stamps the same v2 keys a + failed LLM call does. ``status_code`` pins ``error.code`` to the real response + status, matching v1's SERVER-span behavior.""" + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + info = StandardLoggingPayloadSetup.get_error_information( + original_exception=exception, + traceback_str=traceback_str, + ) + return SpanError( + error_type=info.get("error_class") or info.get("error_code") or None, + message=info.get("error_message") or None, + code=str(status_code) if status_code is not None else (info.get("error_code") or None), + stack_trace=info.get("traceback") or None, + llm_provider=info.get("llm_provider") or None, + ) + + # Any callback whose class belongs to one of these modules is "the OTel # callback" for proxy-global-registration purposes. _OTEL_MODULES = ( @@ -558,7 +586,12 @@ class OpenTelemetryV2(CustomLogger): def start_phase_span(self, name: str) -> "Iterator[Span]": span = self._emitter.start_span(SpanRole.SERVICE, name) with use_span(span, end_on_exit=True): - yield span + try: + yield span + except Exception as exc: + if is_recordable_span(span): + stamp_error(span, _span_error_from_exception(exc), record_event=False, set_status=False) + raise async def async_pre_call_hook( self, @@ -573,6 +606,48 @@ class OpenTelemetryV2(CustomLogger): ) return data + def record_error_attributes_on_span( + self, + span: "Span | None", + exception: "Exception | None", + status_code: int, + ) -> None: + """Stamp the v2 error.* attributes on the FastAPI-owned SERVER span for a + failure that dies before any LLM-call span exists (malformed body, auth / + validation rejection). Called from the proxy's global exception handler via + ``_close_dangling_otel_server_span``. The instrumentor still owns the span's + status and lifecycle, so this only decorates it — never sets status, never + ends it — and emits no exception event, matching v1's SERVER-span behavior + and avoiding a duplicate of the event ``async_post_call_failure_hook`` or + the ``auth`` phase span already records.""" + if span is None or not is_recordable_span(span): + return + stamp_error( + span, + _span_error_from_exception(exception, status_code=status_code), + record_event=False, + set_status=False, + ) + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: "UserAPIKeyAuth", + traceback_str: "str | None" = None, + ) -> None: + """Stamp error.* on the request's root SERVER span for a proxy-level + failure that never reached an LLM call (empty body rejected in the + endpoint, auth failure), so the failed request carries the same error keys + a failed LLM call does. v1's ``OpenTelemetry`` implemented this same hook; + v2 lost it when it stopped subclassing ``OpenTelemetry``, which is the + LIT-4179 regression for pre-call failures.""" + span = request_root_span() or user_api_key_dict.parent_otel_span + if span is None or not is_recordable_span(span): + return None + stamp_error(span, _span_error_from_exception(original_exception, traceback_str=traceback_str)) + return None + def emit_guardrail_span(self, entry: "StandardLoggingGuardrailInformation") -> None: # Emitted by the guardrail-recording code the moment a guardrail finishes, # not from a post-call hook — that hook does not fire on every path (a diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index d4258557fe7..319f03fea89 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -48,7 +48,7 @@ from ...openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, ) -from ..common_utils import FireworksAIException +from ..common_utils import FireworksAIMixin, FireworksAIException def _extract_fireworks_hidden_params(payload: dict) -> dict: @@ -70,7 +70,7 @@ def _extract_fireworks_hidden_params(payload: dict) -> dict: return {**top_level, **per_choice} -class FireworksAIConfig(OpenAIGPTConfig): +class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): """ Reference: https://docs.fireworks.ai/api-reference/post-chatcompletions @@ -114,6 +114,16 @@ class FireworksAIConfig(OpenAIGPTConfig): prompt_truncate_len: Optional[int] = None, context_length_exceeded_behavior: Optional[Literal["error", "truncate"]] = None, ) -> None: + OpenAIGPTConfig.__init__( + self, + frequency_penalty=frequency_penalty, + max_tokens=max_tokens, + n=n, + stop=stop, + temperature=temperature, + top_p=top_p, + response_format=response_format, + ) locals_ = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index a1b6309d1e0..4e22445bcc0 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -12,6 +12,23 @@ class FireworksAIException(BaseLLMException): pass +def get_fireworks_session_id(litellm_params: dict) -> str | None: + params = litellm_params + for key in ("litellm_session_id", "session_id"): + value = params.get(key) + if value: + return str(value) + metadata = params.get("metadata") + if isinstance(metadata, dict): + value = metadata.get("session_id") + if value: + return str(value) + value = params.get("litellm_trace_id") + if value: + return str(value) + return None + + class FireworksAIMixin: """ Common Base Config functions across Fireworks AI Endpoints @@ -47,4 +64,9 @@ class FireworksAIMixin: if api_key is None: raise ValueError("FIREWORKS_API_KEY is not set") - return {"Authorization": "Bearer {}".format(api_key), **headers} + validated_headers = {"Authorization": "Bearer {}".format(api_key), **headers} + if not any(key.lower() == "x-session-affinity" for key in validated_headers): + session_id = get_fireworks_session_id(litellm_params) + if session_id: + validated_headers["x-session-affinity"] = session_id + return validated_headers diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 3193b72a7d9..624190a0b61 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1744,6 +1744,30 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) return non_thinking_tokens == usage_metadata.get("totalTokenCount", 0) + @staticmethod + def _response_has_search_grounding( + completion_response: Union[GenerateContentResponseBody, BidiGenerateContentServerMessage], + ) -> bool: + """ + Whether the response used Grounding with Google Search, detected via + groundingMetadata.webSearchQueries (an actual web search was performed). + + Google bills grounding-with-Google-Search retrieved tokens separately (a per-request / + per-query search fee) and excludes them from input token billing, unlike URL context / + File Search / code execution whose tool-use tokens are charged at the input token rate. + URL context also emits groundingMetadata (with groundingChunks but no webSearchQueries), + so presence of groundingMetadata alone is not a sufficient signal. + See https://ai.google.dev/gemini-api/docs/pricing and + https://github.com/BerriAI/litellm/discussions/33198 + """ + if "candidates" not in completion_response: + return False + for candidate in completion_response["candidates"] or []: + grounding_metadata, _, _, _ = VertexGeminiConfig._extract_candidate_metadata(candidate) + if VertexGeminiConfig._calculate_web_search_requests(grounding_metadata): + return True + return False + @staticmethod def _calculate_usage( completion_response: Union[GenerateContentResponseBody, BidiGenerateContentServerMessage], @@ -1899,12 +1923,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): tool_use_tokens=tool_use_prompt_tokens, ) + billable_tool_use_prompt_tokens = ( + 0 + if VertexGeminiConfig._response_has_search_grounding(completion_response) + else (tool_use_prompt_tokens or 0) + ) + completion_tokens = response_tokens or completion_response["usageMetadata"].get("candidatesTokenCount", 0) if not VertexGeminiConfig.is_candidate_token_count_inclusive(usage_metadata) and reasoning_tokens: completion_tokens = reasoning_tokens + completion_tokens ## GET USAGE ## usage = Usage( - prompt_tokens=usage_metadata.get("promptTokenCount", 0) + (tool_use_prompt_tokens or 0), + prompt_tokens=usage_metadata.get("promptTokenCount", 0) + billable_tool_use_prompt_tokens, completion_tokens=completion_tokens, total_tokens=usage_metadata.get("totalTokenCount", 0), prompt_tokens_details=prompt_tokens_details, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ee996198b28..e5afc81b641 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -16272,7 +16272,7 @@ "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/glm-5p2": { - "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -16686,7 +16686,7 @@ "supports_vision": false }, "fireworks_ai/glm-5p2": { - "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index d55eb3ac014..9fe970f7fa9 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -33,6 +33,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( from litellm.proxy.utils import PrismaClient from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import ( + MCPServerOAuthClientRepository, MCPServerRepository, MCPUserCredentialsRepository, ) @@ -374,6 +375,12 @@ def encrypt_credentials(credentials: MCPCredentials, encryption_key: Optional[st value=client_secret, new_encryption_key=encryption_key, ) + client_private_key = credentials.get("client_private_key") + if client_private_key is not None: + credentials["client_private_key"] = encrypt_value_helper( + value=client_private_key, + new_encryption_key=encryption_key, + ) # AWS SigV4 credential fields aws_access_key_id = credentials.get("aws_access_key_id") if aws_access_key_id is not None: @@ -405,6 +412,7 @@ def decrypt_credentials( "auth_value", "client_id", "client_secret", + "client_private_key", "aws_access_key_id", "aws_secret_access_key", "aws_session_token", @@ -639,6 +647,7 @@ async def delete_mcp_server( for model, label in ( (prisma_client.db.litellm_mcpusercredentials, "credential"), (prisma_client.db.litellm_mcpuserenvvars, "env var"), + (prisma_client.db.litellm_mcpserveroauthclient, "OAuth client"), ): try: await model.delete_many(where={"server_id": server_id}) @@ -823,26 +832,66 @@ async def update_mcp_server( return updated_mcp_server -async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, touched_by: str, new_master_key: str): +async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, server_id: str) -> object | None: + """Read the persisted (encrypted) DCR OAuth client blob for a server from the + server-scoped store, or None. Config.yaml-declared servers have no + LiteLLM_MCPServerTable row, so their dynamically registered client lives here keyed + by server_id. The returned value is the raw credentials blob for + ``_get_persisted_dcr_credentials`` to parse.""" + row = await MCPServerOAuthClientRepository(prisma_client).table.find_unique(where={"server_id": server_id}) + if row is None: + return None + return row.credentials + + +async def upsert_mcp_server_oauth_client_credentials( + prisma_client: PrismaClient, server_id: str, credentials: MCPCredentials +) -> None: + """Persist a server's dynamically registered OAuth client (RFC 7591 DCR) in the + server-scoped store keyed by server_id, independent of any LiteLLM_MCPServerTable row. + client_id/client_secret are encrypted at rest with the same salt key used for the + server row's credentials blob, so ``_apply_persisted_dcr_credentials`` decrypts them the + same way regardless of which store a server's client came from.""" from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + encrypted = encrypt_credentials(credentials=dict(credentials), encryption_key=_get_salt_key()) + blob = safe_dumps(encrypted) + await MCPServerOAuthClientRepository(prisma_client).table.upsert( + where={"server_id": server_id}, + data={ + "create": {"server_id": server_id, "credentials": blob}, + "update": {"credentials": blob}, + }, + ) + + +def _reencrypt_mcp_credentials_blob(credentials: object, new_master_key: str) -> str | None: + """Decrypt an at-rest MCP credentials blob with the current key and re-encrypt it under + new_master_key, returning the serialized blob or None when there is nothing to rotate. Shared by + every table that stores an encrypted MCP credentials blob so a master-key rotation covers them + uniformly and cannot silently skip one.""" + if not credentials: + return None + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # noqa: PLC0415 # avoids circular import + + creds_dict = json.loads(credentials) if isinstance(credentials, str) else dict(credentials) + decrypted = decrypt_credentials(credentials=cast(MCPCredentials, creds_dict)) + encrypted = encrypt_credentials(credentials=decrypted, encryption_key=new_master_key) + return safe_dumps(encrypted) + + +async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, touched_by: str, new_master_key: str): + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # noqa: PLC0415 # avoids circular import + mcp_servers = await MCPServerRepository(prisma_client).table.find_many() updated = 0 for mcp_server in mcp_servers: update_data: Dict[str, Any] = {} - credentials = mcp_server.credentials - if credentials: - # Decrypt with current key first, then re-encrypt with new key - decrypted_credentials = decrypt_credentials( - credentials=cast(MCPCredentials, dict(credentials)), - ) - encrypted_credentials = encrypt_credentials( - credentials=decrypted_credentials, - encryption_key=new_master_key, - ) - update_data["credentials"] = safe_dumps(encrypted_credentials) + rotated_credentials = _reencrypt_mcp_credentials_blob(mcp_server.credentials, new_master_key) + if rotated_credentials is not None: + update_data["credentials"] = rotated_credentials rotated_env_vars = _reencrypt_global_env_var_values(mcp_server.env_vars, new_master_key) if rotated_env_vars is not None: @@ -857,9 +906,23 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, data=update_data, ) updated += 1 + + oauth_clients = await MCPServerOAuthClientRepository(prisma_client).table.find_many() + oauth_updated = 0 + for oauth_client in oauth_clients: + rotated_credentials = _reencrypt_mcp_credentials_blob(oauth_client.credentials, new_master_key) + if rotated_credentials is None: + continue + await MCPServerOAuthClientRepository(prisma_client).table.update( + where={"server_id": oauth_client.server_id}, + data={"credentials": rotated_credentials}, + ) + oauth_updated += 1 + verbose_proxy_logger.info( - "rotate_mcp_server_credentials_master_key: rotated %d MCP server row(s)", + "rotate_mcp_server_credentials_master_key: rotated %d MCP server row(s) and %d OAuth-client row(s)", updated, + oauth_updated, ) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 54aff86aab2..1af64749304 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -971,43 +971,93 @@ def _apply_persisted_dcr_credentials(mcp_server: MCPServer, credentials: _Persis return True -async def _get_persisted_mcp_server_with_dcr_client_id( - mcp_server: MCPServer, -) -> Optional[tuple["LiteLLM_MCPServerTable", _PersistedDcrCredentials]]: - from litellm.proxy._experimental.mcp_server.db import get_mcp_server # noqa: PLC0415 - from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 +async def _load_store_dcr_credentials(mcp_server: MCPServer) -> _PersistedDcrCredentials | None: + """DCR client persisted in the server-scoped OAuth-client store for a config-declared server + (which has no LiteLLM_MCPServerTable row). Returns None when the store has no usable client_id + or the DB is unreachable.""" + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # avoids circular import + get_mcp_server_oauth_client_credentials, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 # avoids circular import try: prisma_client = get_prisma_client_or_throw("Database not connected. Cannot read MCP OAuth client registration.") - persisted_mcp_server = await get_mcp_server( - prisma_client=prisma_client, - server_id=mcp_server.server_id, + blob = await get_mcp_server_oauth_client_credentials( + prisma_client=prisma_client, server_id=mcp_server.server_id ) - except Exception as exc: # noqa: BLE001 + except Exception as exc: # noqa: BLE001 # best-effort read; DB may be unreachable verbose_logger.debug( - "register_client_with_server: failed to read persisted DCR client registration for server_id=%s: %s", + "register_client_with_server: failed to read stored DCR client for server_id=%s: %s", mcp_server.server_id, exc, ) return None - if persisted_mcp_server is None: - return None - - credentials = _get_persisted_dcr_credentials(persisted_mcp_server.credentials) + credentials = _get_persisted_dcr_credentials(blob) if credentials is None or not credentials.client_id: return None + return credentials - return persisted_mcp_server, credentials + +async def hydrate_config_server_dcr_client(mcp_server: MCPServer) -> bool: + """Overlay a config-declared server's persisted DCR client onto its in-memory object so token + refresh can authenticate. Config.yaml servers have no LiteLLM_MCPServerTable row, so their + minted client lives in the server-scoped store; without this overlay the in-memory server + carries no client_id after a restart. An explicit client_id set in config.yaml wins and is never + overwritten by a persisted store client.""" + if mcp_server.client_id: + return False + credentials = await _load_store_dcr_credentials(mcp_server) + if credentials is None: + return False + return _apply_persisted_dcr_credentials(mcp_server, credentials) + + +async def _resolve_persisted_dcr_client( + mcp_server: MCPServer, +) -> tuple[Optional["LiteLLM_MCPServerTable"], _PersistedDcrCredentials | None]: + """Resolve a server's persisted DCR client using the same two-level rule the write path uses, so + read and write always agree. First, whether the server HAS a LiteLLM_MCPServerTable row: a row is + always resolved to that row and the store is never consulted for a server that has a row, so a + caller-chosen server_id colliding with a config-declared server cannot inherit that config + server's client, and a row that exists but carries no usable client_id yields (row, None) rather + than a store fallback. Second, among rowless servers: a config-declared server keeps its client in + the server-scoped store, while a rowless non-config server is a throwaway temp/session server with + no persisted client. Returns (row_or_None, credentials_or_None); the row is only needed by the + reuse path to refresh the registry for a DB-declared server.""" + from litellm.proxy._experimental.mcp_server.db import get_mcp_server # noqa: PLC0415 # avoids circular import + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # avoids circular import + global_mcp_server_manager, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 # avoids circular import + + try: + prisma_client = get_prisma_client_or_throw("Database not connected. Cannot read MCP OAuth client registration.") + row = await get_mcp_server(prisma_client=prisma_client, server_id=mcp_server.server_id) + except Exception as exc: # noqa: BLE001 # best-effort read; DB may be unreachable + verbose_logger.debug( + "register_client_with_server: failed to read persisted DCR client for server_id=%s: %s", + mcp_server.server_id, + exc, + ) + return None, None + + if row is not None: + credentials = _get_persisted_dcr_credentials(row.credentials) + if credentials is not None and credentials.client_id: + return row, credentials + return row, None + if global_mcp_server_manager.is_config_declared_server(mcp_server.server_id): + return None, await _load_store_dcr_credentials(mcp_server) + return None, None async def _reuse_persisted_dcr_client_if_available( mcp_server: MCPServer, current_redirect_uri: Optional[str] = None ) -> bool: - persisted = await _get_persisted_mcp_server_with_dcr_client_id(mcp_server) - if persisted is None: + persisted_mcp_server, credentials = await _resolve_persisted_dcr_client(mcp_server) + if credentials is None: return False - persisted_mcp_server, credentials = persisted if current_redirect_uri is not None and _redirect_uri_not_registered(credentials, current_redirect_uri): verbose_logger.debug( "register_client_with_server: not reusing persisted DCR client for server_id=%s; its registered " @@ -1021,18 +1071,19 @@ async def _reuse_persisted_dcr_client_if_available( if not _apply_persisted_dcr_credentials(mcp_server, credentials): return False - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 - global_mcp_server_manager, - ) - - try: - await global_mcp_server_manager.update_server(persisted_mcp_server) - except Exception as exc: # noqa: BLE001 - verbose_logger.warning( - "register_client_with_server: failed to refresh persisted DCR client registration for server_id=%s: %s", - mcp_server.server_id, - exc, + if persisted_mcp_server is not None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # avoids circular import + global_mcp_server_manager, ) + + try: + await global_mcp_server_manager.update_server(persisted_mcp_server) + except Exception as exc: # noqa: BLE001 # best-effort registry refresh + verbose_logger.warning( + "register_client_with_server: failed to refresh persisted DCR client registration for server_id=%s: %s", + mcp_server.server_id, + exc, + ) return bool(mcp_server.client_id) @@ -1044,10 +1095,9 @@ async def _persisted_dcr_redirect_uri_is_stale(mcp_server: MCPServer, current_re otherwise short-circuits registration before any redirect check can run. Servers without a persisted DCR recording (admin-configured client_id, or registered before redirect_uris were recorded) are never reported stale.""" - persisted = await _get_persisted_mcp_server_with_dcr_client_id(mcp_server) - if persisted is None: + _, credentials = await _resolve_persisted_dcr_client(mcp_server) + if credentials is None: return False - _, credentials = persisted if not _redirect_uri_not_registered(credentials, current_redirect_uri): return False verbose_logger.warning( @@ -1067,7 +1117,10 @@ DcrRegistrationPersistenceResult = Literal["persisted", "reused", "skipped", "fa async def _persist_dcr_client_registration( mcp_server: MCPServer, registration_response: object, current_redirect_uri: str ) -> DcrRegistrationPersistenceResult: - """Persist the dynamically registered OAuth client (RFC 7591) onto the MCP server row. + """Persist the dynamically registered OAuth client (RFC 7591) to its single home: the server's + ``LiteLLM_MCPServerTable`` row when it has one, otherwise the server-scoped store when the server + is config-declared. A rowless server that is not config-declared is a throwaway temp/session + server, so its client is overlaid in memory only and not persisted. The interactive authorization_code flow mints a ``client_id`` via Dynamic Client Registration that discovery cannot re-derive; without persisting it the autonomous @@ -1106,16 +1159,20 @@ async def _persist_dcr_client_registration( if await _reuse_persisted_dcr_client_if_available(mcp_server, current_redirect_uri=current_redirect_uri): return "reused" + token_endpoint_auth_method = ( + "client_secret_basic" if registration.token_endpoint_auth_method == "client_secret_basic" else None + ) credentials: MCPCredentials = { "client_id": registration.client_id, "client_secret": registration.client_secret, - "token_endpoint_auth_method": ( - "client_secret_basic" if registration.token_endpoint_auth_method == "client_secret_basic" else None - ), + "token_endpoint_auth_method": token_endpoint_auth_method, "redirect_uris": [current_redirect_uri], } - from litellm.proxy._experimental.mcp_server.db import update_mcp_server # noqa: PLC0415 + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # avoids circular import + update_mcp_server, + upsert_mcp_server_oauth_client_credentials, + ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 global_mcp_server_manager, ) @@ -1136,7 +1193,18 @@ async def _persist_dcr_client_registration( ), touched_by="mcp_oauth_dcr", ) - await global_mcp_server_manager.update_server(updated_row) + if updated_row is not None: + await global_mcp_server_manager.update_server(updated_row) + return "persisted" + if global_mcp_server_manager.is_config_declared_server(mcp_server.server_id): + await upsert_mcp_server_oauth_client_credentials( + prisma_client=prisma_client, + server_id=mcp_server.server_id, + credentials=credentials, + ) + mcp_server.client_id = registration.client_id + mcp_server.client_secret = registration.client_secret + mcp_server.token_endpoint_auth_method = token_endpoint_auth_method return "persisted" except Exception as exc: # noqa: BLE001 verbose_logger.warning( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 732381edb0e..1ba608b9510 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -93,6 +93,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchange_ ) from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( AuthorizationCodeConfig, + CredError, + IdJagConfig, PassthroughConfig, ServerSpec, TokenExchangeConfig, @@ -621,6 +623,47 @@ def _consumes_caller_authorization(server: MCPServer) -> bool: ) +_REGISTRY_DUMP_SECRET_FIELDS = frozenset( + {"authentication_token", "client_secret", "client_private_key", "aws_secret_access_key", "aws_session_token"} +) + + +def _redacted_registry_dump(servers: dict[str, MCPServer]) -> dict[str, dict[str, str]]: + """A JSON-safe view of the server registry with credential fields masked, for debug logging. + + The registry holds long-lived secrets as plain strings (the static token, OAuth client secret, + the ID-JAG signing key, AWS keys); dumping them verbatim hands the gateway's client identity to + anyone who can read debug logs. + """ + dumps: dict[str, dict[str, object]] = {server_id: server.model_dump() for server_id, server in servers.items()} + return { + server_id: { + field: ("**REDACTED**" if field in _REGISTRY_DUMP_SECRET_FIELDS and value is not None else str(value)) + for field, value in dump.items() + } + for server_id, dump in dumps.items() + } + + +def _to_server_spec_fail_closed(server: MCPServer) -> Optional[ServerSpec]: + """`to_server_spec`, except a half-configured `oauth2_id_jag` server refuses instead of deferring. + + ID-JAG has no v1 arm, so deferring to v1 would let `resolve_mcp_auth` honor a caller x-mcp-* + override or fall through to the static `authentication_token`, both of which bypass the per-user + identity assertion the mode promises. That is an operator misconfiguration, not a fallback. + """ + spec = to_server_spec(server) + if spec is None and server.auth_type == MCPAuth.oauth2_id_jag: + raise_public( + CredError.of_misconfigured( + "oauth2_id_jag requires token_exchange_endpoint, id_jag_resource_token_endpoint, " + "client_id, and a client_secret or client_private_key; refusing to fall back to " + "a static credential." + ) + ) + return spec + + def _caller_authorization_fans_out( server: MCPServer, scope_servers: Optional[list[MCPServer]], @@ -1100,6 +1143,14 @@ class MCPServerManager: """ return self.config_mcp_servers | self.registry + def is_config_declared_server(self, server_id: str) -> bool: + """True when server_id was declared in config.yaml (present in the in-memory config map). + Config servers are rowless and persistent, so their DCR client belongs in the server-scoped + store; a rowless server that is NOT config-declared is a throwaway temp/session server whose + client must not be persisted. This never overrides the row-existence check: a server that has + a LiteLLM_MCPServerTable row is always resolved to that row first.""" + return server_id in self.config_mcp_servers + async def load_servers_from_config( self, mcp_servers_config: dict[str, Any], @@ -1318,6 +1369,12 @@ class MCPServerManager: "subject_token_type", DEFAULT_SUBJECT_TOKEN_TYPE, ), + # ID-JAG fields + id_jag_resource_token_endpoint=server_config.get("id_jag_resource_token_endpoint", None), + id_jag_resource=server_config.get("id_jag_resource", None), + client_private_key=server_config.get("client_private_key", None), + client_private_key_id=server_config.get("client_private_key_id", None), + client_assertion_signing_alg=server_config.get("client_assertion_signing_alg", "RS256"), token_exchange_profile=server_config.get("token_exchange_profile", "rfc8693"), allow_sampling=bool(server_config.get("allow_sampling", False)), allow_elicitation=bool(server_config.get("allow_elicitation", False)), @@ -1338,10 +1395,36 @@ class MCPServerManager: base_url=server_config.get("url", ""), ) - verbose_logger.debug(f"Loaded MCP Servers: {json.dumps(self.config_mcp_servers, indent=4, default=str)}") + verbose_logger.debug( + f"Loaded MCP Servers: {json.dumps(_redacted_registry_dump(self.config_mcp_servers), indent=4)}" + ) + + await self._hydrate_config_servers_dcr_clients() self.initialize_tool_name_to_mcp_server_name_mapping() + async def _hydrate_config_servers_dcr_clients(self) -> None: + """Overlay each config-declared server's persisted DCR client (from the server-scoped + store) onto its in-memory object so token refresh authenticates after a restart. A + best-effort no-op when the DB is unreachable at config-load time.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( # noqa: PLC0415 # circular import + hydrate_config_server_dcr_client, + ) + + for server in self.config_mcp_servers.values(): + try: + if await hydrate_config_server_dcr_client(server): + verbose_logger.debug( + "hydrated persisted DCR client onto config MCP server server_id=%s", + server.server_id, + ) + except Exception as exc: # noqa: BLE001 # best-effort hydration; never fail config load + verbose_logger.debug( + "load_servers_from_config: failed to hydrate DCR client for server_id=%s: %s", + server.server_id, + exc, + ) + async def _register_openapi_tools(self, spec_path: str, server: MCPServer, base_url: str): """ Register tools from an OpenAPI specification for a given server. @@ -1765,6 +1848,21 @@ class MCPServerManager: subject_token_type=mcp_server.subject_token_type or (credentials_dict.get("subject_token_type") if credentials_dict else None) or DEFAULT_SUBJECT_TOKEN_TYPE, + # ID-JAG fields — read from credentials JSON blob + id_jag_resource_token_endpoint=( + credentials_dict.get("id_jag_resource_token_endpoint") if credentials_dict else None + ), + id_jag_resource=(credentials_dict.get("id_jag_resource") if credentials_dict else None), + client_private_key=self._decrypt_credential_field( + credentials_dict.get("client_private_key") if credentials_dict else None, + "client_private_key", + credentials_are_encrypted, + ), + client_private_key_id=(credentials_dict.get("client_private_key_id") if credentials_dict else None), + client_assertion_signing_alg=( + credentials_dict.get("client_assertion_signing_alg") if credentials_dict else None + ) + or "RS256", token_exchange_profile=mcp_server.token_exchange_profile or (credentials_dict.get("token_exchange_profile") if credentials_dict else None) or "rfc8693", @@ -2641,9 +2739,10 @@ class MCPServerManager: ) if not conflicts: return auth, extra_headers - if isinstance(spec.config, (TokenExchangeConfig, AuthorizationCodeConfig)): + if isinstance(spec.config, (TokenExchangeConfig, AuthorizationCodeConfig, IdJagConfig)): # The resolver owns the per-user credential here (token_exchange's exchanged - # token, authorization_code's stored token). It is authoritative: a guardrail such + # token, authorization_code's stored token, id_jag's minted assertion). It is + # authoritative: a guardrail such # as MCPJWTSigner, static_headers, or any other injected Authorization must NOT # shadow it (otherwise the upstream gets e.g. the signer's JWT instead of the # exchanged token and rejects it). Drop the conflicting header so the resolved @@ -2734,20 +2833,23 @@ class MCPServerManager: Configured MCP client instance. """ transport = server.transport or MCPTransport.sse - spec = None if transport == MCPTransport.stdio else to_server_spec(server) + spec = None if transport == MCPTransport.stdio else _to_server_spec_fail_closed(server) provider = cred_provider or self._cred_provider # A caller-supplied per-request override (mcp_auth_header / x-mcp-*) defers to the v1 path # so it wins - except for the modes the v2 resolver owns per-caller (authorization_code's - # stored token, token_exchange's RFC 8693 minted token, and the passthrough modes' - # forwarded caller token). A caller must not be able to substitute another user's stored - # credential, nor silently disable the OBO exchange and forward an arbitrary bearer - # upstream, so we keep the v2 spec and ignore the override for these; the REST tools - # preview supplies its not-yet-persisted token through the resolver (cred_provider), - # never this path. + # stored token, token_exchange's RFC 8693 minted token, id_jag's minted assertion, and the + # passthrough modes' forwarded caller token). A caller must not be able to substitute another + # user's stored credential, nor silently disable the OBO / ID-JAG exchange and forward an + # arbitrary bearer upstream, so we keep the v2 spec and ignore the override for these; the + # REST tools preview supplies its not-yet-persisted token through the resolver + # (cred_provider), never this path. if ( spec is not None and mcp_auth_header - and not isinstance(spec.config, (AuthorizationCodeConfig, PassthroughConfig, TokenExchangeConfig)) + and not isinstance( + spec.config, + (AuthorizationCodeConfig, IdJagConfig, PassthroughConfig, TokenExchangeConfig), + ) ): spec = None auth_value = ( @@ -4276,10 +4378,13 @@ class MCPServerManager: if server_auth_header is None: server_auth_header = mcp_auth_header - # Extract subject token for OAuth2 Token Exchange (OBO) flow + # Extract subject token for OAuth2 Token Exchange (OBO) and ID-JAG flows subject_token: Optional[str] = None extra_headers: Optional[dict[str, str]] = None - if mcp_server.auth_type == MCPAuth.oauth2_token_exchange: + if mcp_server.auth_type in ( + MCPAuth.oauth2_token_exchange, + MCPAuth.oauth2_id_jag, + ): subject_token = self._extract_bearer_token(oauth2_headers, raw_headers) elif mcp_server.auth_type == MCPAuth.oauth2: if mcp_server.has_client_credentials: @@ -4381,10 +4486,10 @@ class MCPServerManager: arguments=arguments, ) - if mcp_server.auth_type == MCPAuth.oauth2_token_exchange and subject_token: - # OBO: the exchanged token may have been revoked/rotated upstream since it was cached, so - # an upstream 401 gets one re-mint + retry. Gated to this mode; all others keep the plain - # single call below. + if mcp_server.auth_type in (MCPAuth.oauth2_token_exchange, MCPAuth.oauth2_id_jag) and subject_token: + # OBO / ID-JAG: the exchanged token may have been revoked/rotated upstream since it was + # cached, so an upstream 401 gets one invalidate + re-mint + retry. Gated to these modes; + # all others keep the plain single call below. async def _obo_call_tool_limited(): async with self._limit_outbound_concurrency(mcp_server): return await self._obo_call_tool_with_retry( @@ -4935,6 +5040,8 @@ class MCPServerManager: verbose_logger.debug("MCP registry refreshed (%s servers in registry)", len(registered_registry)) + await self._hydrate_config_servers_dcr_clients() + def get_mcp_servers_from_ids(self, server_ids: list[str]) -> list[MCPServer]: servers = [] registry = self.get_registry() diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py index 73166a45d6e..2bdb8770e4e 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/__init__.py @@ -31,10 +31,14 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( AwsCredentialSource, AwsSigV4Config, Byok, + ClientAuth, ClientCredentialsConfig, + ClientSecretAuth, CredError, + IdJagConfig, NoneConfig, PassthroughConfig, + PrivateKeyJwtAuth, ServerSpec, SharedKey, StaticKeys, @@ -59,6 +63,10 @@ __all__ = [ "AuthorizationCodeConfig", "ClientCredentialsConfig", "TokenExchangeConfig", + "IdJagConfig", + "ClientAuth", + "PrivateKeyJwtAuth", + "ClientSecretAuth", "ApiKeyConfig", "ApiKeySource", "SharedKey", diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index e87e8081ced..6631e38f524 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -21,9 +21,13 @@ from typing_extensions import assert_never from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, + ClientAuth, + ClientSecretAuth, CredError, + IdJagConfig, NoneConfig, PassthroughConfig, + PrivateKeyJwtAuth, ServerSpec, SharedKey, Subject, @@ -35,6 +39,9 @@ if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer +_TOKEN_EXCHANGE_SUBJECT_TOKEN_DEFAULT = "urn:ietf:params:oauth:token-type:access_token" +_ID_JAG_SUBJECT_TOKEN_DEFAULT = "urn:ietf:params:oauth:token-type:id_token" + def to_subject(user_api_key_auth: Optional[UserAPIKeyAuth], subject_token: Optional[str]) -> Subject: """Map v1's authenticated principal onto the resolver's Subject. @@ -96,6 +103,8 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]: ) # client_credentials (M2M) and delegate/passthrough oauth2 stay on v1 return None + case MCPAuth.oauth2_id_jag: + return _id_jag_spec(server, resource) case MCPAuth.true_passthrough | MCPAuth.oauth_delegate: return ServerSpec(server_id=server.server_id, resource=resource, config=PassthroughConfig()) case MCPAuth.oauth2_token_exchange: @@ -167,6 +176,58 @@ def _shared_key_spec( ) +def _id_jag_spec(server: MCPServer, resource: str) -> Optional[ServerSpec]: + """Build an ID-JAG spec from the v1 server's raw fields, or defer (None) if half-configured. + + The enum already routes here, but a server missing an endpoint, ``client_id``, or any client-auth + secret would make ``IdJagConfig`` raise at construction; returning None instead defers to v1 so a + partially configured server does not 500. ``token_exchange_endpoint`` is leg 1 (the IdP org AS); + leg 2 is ``id_jag_resource_token_endpoint`` (the upstream resource AS). + """ + org_token_endpoint = server.token_exchange_endpoint + resource_token_endpoint = server.id_jag_resource_token_endpoint + client_id = server.client_id + client_auth = _id_jag_client_auth(server) + if not org_token_endpoint or not resource_token_endpoint or not client_id or client_auth is None: + return None + return ServerSpec( + server_id=server.server_id, + resource=resource, + config=IdJagConfig( + org_token_endpoint=org_token_endpoint, + resource_token_endpoint=resource_token_endpoint, + client_id=client_id, + client_auth=client_auth, + subject_token_type=_id_jag_subject_token_type(server), + audience=server.audience, + resource=server.id_jag_resource, + scopes=tuple(server.scopes or ()), + ), + ) + + +def _id_jag_client_auth(server: MCPServer) -> Optional[ClientAuth]: + """Private-key JWT when a key is configured, else client_secret, else None (defer to v1).""" + if server.client_private_key: + return PrivateKeyJwtAuth( + private_key=SecretStr(server.client_private_key), + key_id=server.client_private_key_id, + signing_alg=server.client_assertion_signing_alg, + ) + if server.client_secret: + return ClientSecretAuth(client_secret=SecretStr(server.client_secret)) + return None + + +def _id_jag_subject_token_type(server: MCPServer) -> str: + """ID-JAG asserts the user's id_token, so the token-exchange access_token default maps to id_token; + an explicitly configured value (e.g. a SAML2 assertion type) is honored verbatim.""" + configured = server.subject_token_type + if configured and configured != _TOKEN_EXCHANGE_SUBJECT_TOKEN_DEFAULT: + return configured + return _ID_JAG_SUBJECT_TOKEN_DEFAULT + + def raise_public(error: CredError) -> NoReturn: """Map a resolver CredError onto the proxy's public HTTP contract. The one edge that raises.""" match error.tag: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index ecfd471190c..7e5c073870a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -16,6 +16,8 @@ follow-up PR with their seam. Pure v2: no imports from v1. from __future__ import annotations +import hashlib + import httpx from typing_extensions import assert_never @@ -33,6 +35,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( Ok, Result, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( + ExchangedToken, + ExchangedTokenCache, + TokenEndpointClient, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( TokenExchanger, ) @@ -42,16 +49,24 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( AuthSpecKind, AwsSigV4Config, Byok, + ClientAuth, ClientCredentialsConfig, + ClientSecretAuth, CredError, + IdJagConfig, NoneConfig, PassthroughConfig, + PrivateKeyJwtAuth, ServerSpec, SharedKey, Subject, TokenExchangeConfig, ) +_TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange" +_JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer" +_ID_JAG_REQUESTED_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:id-jag" + class _NullOAuthTokenStore: """Fail-closed default: with no token store wired, every user reads as not authorized.""" @@ -87,9 +102,13 @@ class UpstreamCredentialProvider: self, oauth_token_store: OAuthTokenStore | None = None, token_exchanger: TokenExchanger | None = None, + token_endpoint: TokenEndpointClient | None = None, + exchanged_tokens: ExchangedTokenCache | None = None, ) -> None: self._oauth_token_store: OAuthTokenStore = oauth_token_store or _NullOAuthTokenStore() self._token_exchanger: TokenExchanger = token_exchanger or _NullTokenExchanger() + self._token_endpoint: TokenEndpointClient = token_endpoint or TokenEndpointClient() + self._exchanged_tokens: ExchangedTokenCache = exchanged_tokens or ExchangedTokenCache() async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: @@ -103,6 +122,8 @@ class UpstreamCredentialProvider: return _not_implemented(AuthSpecKind.client_credentials) case TokenExchangeConfig() as config: return await self._token_exchange(subject, server, config) + case IdJagConfig() as config: + return await self._id_jag(subject, server, config) case AuthorizationCodeConfig(): return await self._authorization_code(subject, server) case AwsSigV4Config(): @@ -141,6 +162,53 @@ class UpstreamCredentialProvider: return Error(CredError.of_not_implemented("api_key BYOK source not implemented yet")) assert_never(config.key_source) + async def _id_jag(self, subject: Subject, server: ServerSpec, config: IdJagConfig) -> Result[httpx.Auth, CredError]: + if subject.inbound_token is None: + return Error( + CredError.of_precondition_required( + "ID-JAG requires a caller identity token; it asserts the calling " + "user's identity upstream and cannot use a static credential." + ) + ) + token = subject.inbound_token.get_secret_value() + cache_key = _id_jag_cache_key(token, server.server_id, config) + + async def _exchange() -> Result[ExchangedToken, CredError]: + leg1_params = { + "grant_type": _TOKEN_EXCHANGE_GRANT_TYPE, + "requested_token_type": _ID_JAG_REQUESTED_TOKEN_TYPE, + "subject_token": token, + "subject_token_type": config.subject_token_type, + **({"audience": config.audience} if config.audience else {}), + **({"resource": config.resource} if config.resource else {}), + **({"scope": " ".join(config.scopes)} if config.scopes else {}), + } + match await self._token_endpoint.fetch( + config.org_token_endpoint, + config.client_id, + leg1_params, + config.client_auth, + ): + case Error(err): + return Error(err) + case Ok(id_jag): + leg2_params = { + "grant_type": _JWT_BEARER_GRANT_TYPE, + "assertion": id_jag.access_token, + } + return await self._token_endpoint.fetch( + config.resource_token_endpoint, + config.client_id, + leg2_params, + config.client_auth, + ) + + match await self._exchanged_tokens.get_or_compute(cache_key, _exchange): + case Ok(access_token): + return Ok(StaticHeaderAuth(f"Bearer {access_token}")) + case Error(err): + return Error(err) + async def _authorization_code(self, subject: Subject, server: ServerSpec) -> Result[StaticHeaderAuth, CredError]: token = await self._authz_token(subject, server) if token is None: @@ -176,13 +244,19 @@ class UpstreamCredentialProvider: """Drop any cached credential the resolver owns for this `(subject, server)`. Used after an upstream rejects the injected credential, so the next resolve re-mints rather - than serving the same rejected token until TTL. Only `token_exchange` holds a re-mintable - cached credential here; other modes are a no-op. + than serving the same rejected token until TTL. `token_exchange` and `id_jag` hold a + re-mintable cached credential here; other modes are a no-op. """ - if isinstance(server.config, TokenExchangeConfig) and subject.inbound_token is not None: + if subject.inbound_token is None: + return + if isinstance(server.config, TokenExchangeConfig): await self._token_exchanger.invalidate( subject.inbound_token.get_secret_value(), server, server.config, tenant_id=subject.tenant_id ) + if isinstance(server.config, IdJagConfig): + self._exchanged_tokens.invalidate( + _id_jag_cache_key(subject.inbound_token.get_secret_value(), server.server_id, server.config) + ) async def _authz_token(self, subject: Subject, server: ServerSpec) -> OAuthToken | None: """The user's authorization_code token, or None when absent or the store is unreachable. @@ -196,5 +270,41 @@ class UpstreamCredentialProvider: return None +def _id_jag_cache_key(subject_token: str, server_id: str, config: IdJagConfig) -> str: + """Bind the cached leg-2 bearer to the caller token, the server, AND the config that minted it. + + Every exchange parameter derives from the config (endpoints, audience, resource, scopes, client + auth), so a server update that changes any of them must change the key; otherwise the old bearer, + authorized under the old policy, keeps being served until its TTL. Everything is hashed, so no + secret is held in the key. + """ + material = "\x00".join( + ( + subject_token, + server_id, + config.org_token_endpoint, + config.resource_token_endpoint, + config.client_id, + _client_auth_fingerprint(config.client_auth), + config.subject_token_type, + config.audience or "", + config.resource or "", + " ".join(config.scopes), + ) + ) + return hashlib.sha256(material.encode()).hexdigest() + + +def _client_auth_fingerprint(client_auth: ClientAuth) -> str: + match client_auth: + case PrivateKeyJwtAuth() as auth: + return "\x00".join( + ("private_key_jwt", auth.private_key.get_secret_value(), auth.key_id or "", auth.signing_alg) + ) + case ClientSecretAuth() as auth: + return "\x00".join(("client_secret", auth.client_secret.get_secret_value())) + assert_never(client_auth) + + def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]: return Error(CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet")) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py new file mode 100644 index 00000000000..4bc5732ec0e --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py @@ -0,0 +1,225 @@ +"""An authenticated OAuth token-endpoint call plus a short-lived-token cache. + +`TokenEndpointClient.fetch` POSTs one grant to a token endpoint, authenticating the gateway as +an OAuth client via `client_auth` (RFC 7523 private-key JWT, or `client_secret_post`), and returns +the minted token or a typed `CredError`. `ExchangedTokenCache` memoizes the final token string per +opaque cache key with per-key single-flight, so concurrent callers share one round-trip and a hit +skips the endpoint entirely. + +Pure v2: no imports from the v1 MCP auth handlers. The multi-leg flows that compose these (ID-JAG, +and later token_exchange / client_credentials) live in the resolver arms; this collaborator owns +only the single authenticated call and the cache. +""" + +from __future__ import annotations + +import asyncio +import json +import time +import uuid +import weakref +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass + +import httpx +import jwt +from pydantic import BaseModel, ValidationError +from typing_extensions import assert_never + +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import ( + MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, + MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, + MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, +) +from litellm.exceptions import Timeout +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # litellm http handler is untyped +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientAuth, + ClientSecretAuth, + CredError, + PrivateKeyJwtAuth, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + +CLIENT_ASSERTION_TYPE = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer" +CLIENT_ASSERTION_LIFETIME_SECONDS = 60 + + +@dataclass(frozen=True, slots=True) +class ExchangedToken: + access_token: str + expires_in: int | None + + +class _TokenEndpointResponse(BaseModel): + access_token: str + expires_in: int | None = None + + +class TokenEndpointClient: + """One authenticated POST to an OAuth token endpoint, returning the minted token as a value.""" + + async def fetch( + self, + endpoint: str, + client_id: str, + grant_params: Mapping[str, str], + client_auth: ClientAuth, + ) -> Result[ExchangedToken, CredError]: + try: + data = {**grant_params, **_client_auth_params(endpoint, client_id, client_auth)} + except (ValueError, TypeError, NotImplementedError, jwt.PyJWTError): + verbose_proxy_logger.warning("MCP token endpoint %s: could not sign the client assertion", endpoint) + return Error( + CredError.of_misconfigured( + "token exchange failed: could not sign the client assertion; " + "check client_private_key and client_assertion_signing_alg" + ) + ) + try: + raw = await _post_form(endpoint, data) + except httpx.HTTPStatusError as exc: + verbose_proxy_logger.warning( + "MCP token endpoint %s failed with status %s", endpoint, exc.response.status_code + ) + return Error( + CredError.of_upstream_unavailable(f"token exchange failed with status {exc.response.status_code}") + ) + except (httpx.RequestError, Timeout) as exc: + verbose_proxy_logger.warning("MCP token endpoint %s unreachable: %s", endpoint, type(exc).__name__) + return Error( + CredError.of_upstream_unavailable( + f"token exchange failed: token endpoint unreachable ({type(exc).__name__})" + ) + ) + except json.JSONDecodeError: + verbose_proxy_logger.warning("MCP token endpoint %s returned a non-JSON response", endpoint) + return Error( + CredError.of_upstream_unavailable("token exchange failed: token endpoint returned a non-JSON response") + ) + if raw is None: + verbose_proxy_logger.warning("MCP token endpoint %s returned no response", endpoint) + return Error(CredError.of_upstream_unavailable("token exchange failed: no response from token endpoint")) + try: + parsed = _TokenEndpointResponse.model_validate(raw) + except ValidationError: + verbose_proxy_logger.warning("MCP token endpoint %s response missing access_token", endpoint) + return Error( + CredError.of_upstream_unavailable("token exchange failed: token endpoint response missing access_token") + ) + return Ok(ExchangedToken(access_token=parsed.access_token, expires_in=parsed.expires_in)) + + +class ExchangedTokenCache: + """Memoizes the final token string per key, single-flighting concurrent misses on one lock.""" + + def __init__(self) -> None: + self._cache = InMemoryCache( + max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, + default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, + ) + self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() + + async def get_or_compute( + self, + cache_key: str, + compute: Callable[[], Awaitable[Result[ExchangedToken, CredError]]], + ) -> Result[str, CredError]: + cached = self._get(cache_key) + if cached is not None: + return Ok(cached) + async with self._lock(cache_key): + cached = self._get(cache_key) + if cached is not None: + return Ok(cached) + match await compute(): + case Ok(token): + self._cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + cache_key, + token.access_token, + ttl=_cache_ttl_seconds(token.expires_in), + ) + return Ok(token.access_token) + case Error(err): + return Error(err) + + def invalidate(self, cache_key: str) -> None: + """Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401).""" + self._cache.delete_cache(cache_key) # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + + def _get(self, cache_key: str) -> str | None: + value = self._cache.get_cache(cache_key) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # InMemoryCache is untyped; narrowed by isinstance below + return value if isinstance(value, str) else None + + def _lock(self, cache_key: str) -> asyncio.Lock: + lock = self._locks.get(cache_key) + if lock is None: + lock = asyncio.Lock() + self._locks[cache_key] = lock + return lock + + +def _cache_ttl_seconds(expires_in: int | None) -> int: + lifetime = expires_in if expires_in is not None else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL + return max( + lifetime - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, + ) + + +async def _post_form(endpoint: str, data: dict[str, str]) -> object | None: + # litellm's httpx handler and httpx.Response are only partially typed; the token endpoint + # returns a JSON object that `_TokenEndpointResponse` validates, so the untyped boundary is + # contained here. A non-2xx raises `httpx.HTTPStatusError`, an unreachable endpoint raises + # `httpx.RequestError` (or litellm's `Timeout`, which the handler substitutes for + # `httpx.TimeoutException`), and a non-JSON body raises `json.JSONDecodeError`; `fetch` maps + # each to a CredError. + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) # pyright: ignore[reportUnknownVariableType] # litellm http handler is untyped + response = await client.post(endpoint, data=data) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # litellm http handler is untyped + if response is None: + return None + response.raise_for_status() + return response.json() # pyright: ignore[reportAny] # untyped JSON; validated by _TokenEndpointResponse in fetch + + +def _client_auth_params(endpoint: str, client_id: str, client_auth: ClientAuth) -> dict[str, str]: + match client_auth: + case PrivateKeyJwtAuth() as auth: + return { + "client_id": client_id, + "client_assertion_type": CLIENT_ASSERTION_TYPE, + "client_assertion": _client_assertion(endpoint, client_id, auth), + } + case ClientSecretAuth() as auth: + return { + "client_id": client_id, + "client_secret": auth.client_secret.get_secret_value(), + } + assert_never(client_auth) + + +def _client_assertion(endpoint: str, client_id: str, auth: PrivateKeyJwtAuth) -> str: + now = int(time.time()) + return jwt.encode( + { + "iss": client_id, + "sub": client_id, + "aud": endpoint, + "jti": uuid.uuid4().hex, + "iat": now, + "exp": now + CLIENT_ASSERTION_LIFETIME_SECONDS, + }, + auth.private_key.get_secret_value(), + algorithm=auth.signing_alg, + headers={"kid": auth.key_id} if auth.key_id else None, + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 7e04be4f045..64a20255ab2 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -56,6 +56,7 @@ class AuthSpecKind(str, Enum): authorization_code = "authorization_code" # per-user 3LO; gateway-stored token client_credentials = "client_credentials" # gateway service account (M2M) token_exchange = "token_exchange" # RFC 8693: token endpoint + subject_token (OBO) + id_jag = "id_jag" # draft-ietf-oauth-identity-assertion-authz-grant: two-leg exchange then jwt-bearer api_key = "api_key" # static header, any scheme (BYOK = per-user-seeded source) passthrough = "passthrough" # client forwards an upstream-audience token none = "none" # no upstream credential; resolve yields a no-op auth, never an error @@ -225,6 +226,49 @@ class TokenExchangeConfig(BaseModel): scopes: tuple[str, ...] = () +class PrivateKeyJwtAuth(BaseModel): + """RFC 7523 private-key-JWT client authentication: the gateway signs a `client_assertion`.""" + + model_config = ConfigDict(frozen=True) + source: Literal["private_key_jwt"] = "private_key_jwt" + private_key: SecretStr + key_id: str | None = None + signing_alg: str = "RS256" + + +class ClientSecretAuth(BaseModel): + """`client_secret_post` client authentication: the gateway posts `client_id` + `client_secret`.""" + + model_config = ConfigDict(frozen=True) + source: Literal["client_secret"] = "client_secret" + client_secret: SecretStr + + +ClientAuth = Annotated[PrivateKeyJwtAuth | ClientSecretAuth, Field(discriminator="source")] + + +class IdJagConfig(BaseModel): + """draft-ietf-oauth-identity-assertion-authz-grant (Okta "AI agent token exchange"). + + Two legs: leg 1 is an RFC 8693 token exchange at the IdP org AS (`org_token_endpoint`) that + swaps the caller's identity token for an ID-JAG assertion; leg 2 is an RFC 7523 jwt-bearer at + the upstream resource AS (`resource_token_endpoint`) that swaps the assertion for the access + token. The gateway authenticates to both endpoints as `client_id` via `client_auth`. Required + fields are enforced at construction so a half-configured server cannot reach the arm. + """ + + model_config = ConfigDict(frozen=True) + kind: Literal[AuthSpecKind.id_jag] = AuthSpecKind.id_jag + org_token_endpoint: str + resource_token_endpoint: str + client_id: str + client_auth: ClientAuth + subject_token_type: str = "urn:ietf:params:oauth:token-type:id_token" + audience: str | None = None + resource: str | None = None + scopes: tuple[str, ...] = () + + class SharedKey(BaseModel): """A fixed key configured on the server, identical for every caller.""" @@ -323,6 +367,7 @@ AuthConfig = Annotated[ AuthorizationCodeConfig | ClientCredentialsConfig | TokenExchangeConfig + | IdJagConfig | ApiKeyConfig | PassthroughConfig | NoneConfig diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d102c1d1e37..b47b43411c5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1011,10 +1011,10 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): mcp_tool_search_enabled: Optional[bool] = None +from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402 from litellm.types.object_permission import ( # noqa: E402 ObjectPermissionDict as ObjectPermissionDict, ) -from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402 class GenerateRequestBase(LiteLLMPydanticObjectBase): @@ -2122,6 +2122,8 @@ class ConfigList(LiteLLMPydanticObjectBase): field_default_value: Any premium_field: bool = False nested_fields: Optional[List[FieldDetail]] = None # For nested dictionary or Pydantic fields + field_options: Optional[list[str]] = None # Allowed values, for field_type == "Select" + field_tab: Optional[str] = None # Admin UI sub-tab this field renders under; None groups it with the rest class UserHeaderMapping(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c7c9397d850..1dc0ee3f947 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -925,9 +925,12 @@ class ProxyBaseLLMRequestProcessing: # If conversion fails, use original spend pass + model_name = ProxyBaseLLMRequestProcessing._get_deployment_model_name(litellm_logging_obj) + headers = { "x-litellm-call-id": call_id, "x-litellm-model-id": model_id, + "x-litellm-model-name": model_name, "x-litellm-cache-key": cache_key, "x-litellm-model-api-base": ( api_base.split("?")[0] if api_base else None @@ -1396,6 +1399,27 @@ class ProxyBaseLLMRequestProcessing: model_id = model_info.get("id", "") or "" return model_id + @staticmethod + def _get_deployment_model_name( + litellm_logging_obj: LiteLLMLoggingObj | None, + ) -> str | None: + """Extract the underlying deployment model string (e.g. ``azure/gpt-4o``). + + The router rewrites the response ``model`` field to the model-group alias + the client requested, so neither the response body nor the existing + headers expose the concrete deployment model. The router records it under + ``litellm_params`` metadata as ``deployment``, so read it back from there. + """ + litellm_params = getattr(litellm_logging_obj, "litellm_params", None) + if not isinstance(litellm_params, dict): + return None + for key in ("litellm_metadata", "metadata"): + metadata = litellm_params.get(key, {}) or {} + deployment = metadata.get("deployment") + if deployment: + return deployment + return None + @staticmethod def _response_cost_from_logging_obj( *, diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/__init__.py new file mode 100644 index 00000000000..ba4c712764e --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/straiker/__init__.py @@ -0,0 +1,71 @@ +from typing import TYPE_CHECKING + +import litellm +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .straiker import StraikerGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + +_OPTIONAL_INIT_FIELDS = ( + "timeout", + "max_retries", + "initial_backoff", + "max_backoff", + "unreachable_fallback", + "fail_on_error", + "max_payload_bytes", + "custom_headers", + "metadata", + "verbose", +) + + +def _get_config_value(litellm_params: "LitellmParams", optional_params: object, attribute_name: str) -> object: + if optional_params is not None: + if isinstance(optional_params, dict): + value = optional_params.get(attribute_name) + else: + value = getattr(optional_params, attribute_name, None) + if value is not None: + return value + return getattr(litellm_params, attribute_name, None) + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + optional_params = getattr(litellm_params, "optional_params", None) + api_key = litellm_params.api_key + if not api_key: + raise ValueError("api_key is required for straiker") + + api_base = litellm_params.api_base or "https://api.prod.straiker.ai" + default_app = getattr(litellm_params, "default_app", None) or getattr(litellm_params, "source", None) + source = default_app if isinstance(default_app, str) and default_app else "LiteLLM Gateway" + kwargs: dict[str, object] = { + field: value + for field in _OPTIONAL_INIT_FIELDS + for value in [_get_config_value(litellm_params, optional_params, field)] + if value is not None + } + _callback = StraikerGuardrail( + api_key=api_key, + api_base=api_base if isinstance(api_base, str) else "https://api.prod.straiker.ai", + source=source, + guardrail_name=guardrail.get("guardrail_name", "straiker"), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + **kwargs, + ) + + litellm.logging_callback_manager.add_litellm_callback(_callback) + return _callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.STRAIKER.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.STRAIKER.value: StraikerGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py new file mode 100644 index 00000000000..5c9f93fc2cd --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py @@ -0,0 +1,541 @@ +from __future__ import annotations + +import asyncio +import json +import random +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, NoReturn +from urllib.parse import urlsplit + +import httpx +from pydantic import ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm._version import version as litellm_version +from litellm.exceptions import ( + BadRequestError, + GuardrailRaisedException, + ModifyResponseException, + Timeout, +) +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + get_session_id_from_request_data, + log_guardrail_information, +) +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.straiker import ( + STRAIKER_WEBHOOK_SCHEMA_VERSION, + StraikerGuardrailConfigModel, + StraikerWebhookApplication, + StraikerWebhookContent, + StraikerWebhookContext, + StraikerWebhookEvent, + StraikerWebhookIdentity, + StraikerWebhookRequest, + StraikerWebhookResponse, + StraikerWebhookStream, + StraikerWebhookUsage, +) +from litellm.types.utils import GenericGuardrailAPIInputs, Usage + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + +GUARDRAIL_NAME = "straiker" +DEFAULT_BLOCK_MESSAGE = "Content violates policy" +DEFAULT_API_BASE = "https://api.prod.straiker.ai" +DEFAULT_MAX_PAYLOAD_BYTES = 524288 +WEBHOOK_PATH = "/api/v1/detect/webhook" +RETRY_STATUS = frozenset({408, 429, 500, 502, 503, 504}) +UNREACHABLE_STATUS = frozenset({502, 503, 504}) +_APPLICATION_METADATA_KEYS = frozenset({"agent_id", "app_name"}) +_OPAQUE_METADATA_SCALAR_TYPES = (str, int, float, bool) + + +@dataclass(frozen=True, slots=True) +class _WebhookFailure: + message: str + is_unreachable: bool + + +def _as_dict(value: object) -> dict: + return value if isinstance(value, dict) else {} + + +def _merged_metadata(request_data: dict) -> dict: + return { + **_as_dict(request_data.get("metadata")), + **_as_dict(request_data.get("litellm_metadata")), + } + + +def _as_optional_str(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _build_webhook_metadata(request_data: dict, default_metadata: dict[str, str]) -> dict[str, object] | None: + out: dict[str, object] = {} + for key, value in _as_dict(request_data.get("metadata")).items(): + if key in _APPLICATION_METADATA_KEYS or key.startswith("user_api"): + continue + if key == "session_id": + continue + if isinstance(value, _OPAQUE_METADATA_SCALAR_TYPES): + out[key] = value + out.update(default_metadata) + return out or None + + +def _extract_identity(request_data: dict) -> StraikerWebhookIdentity: + meta = _merged_metadata(request_data) + return StraikerWebhookIdentity( + litellm_key=_as_optional_str(meta.get("user_api_key_alias")) + or _as_optional_str(meta.get("user_api_key_hash")) + or _as_optional_str(meta.get("user_api_key_token")), + litellm_team=_as_optional_str(meta.get("user_api_key_team_alias")) + or _as_optional_str(meta.get("user_api_key_team_id")), + litellm_user_id=_as_optional_str(meta.get("user_api_key_user_id")), + litellm_user_email=_as_optional_str(meta.get("user_api_key_user_email")), + litellm_org_id=_as_optional_str(meta.get("user_api_key_org_id")), + end_user_id=_as_optional_str(meta.get("user_api_key_end_user_id")), + ) + + +def _resolve_provider(request_data: dict, model: str | None) -> str | None: + litellm_params = _as_dict(request_data.get("litellm_params")) + custom_llm_provider = request_data.get("custom_llm_provider") or litellm_params.get("custom_llm_provider") + if custom_llm_provider: + return custom_llm_provider + if not model: + return None + try: + _, provider, _, _ = get_llm_provider( + model=model, + api_base=request_data.get("api_base") or litellm_params.get("api_base"), + api_key=request_data.get("api_key") or litellm_params.get("api_key"), + ) + except BadRequestError: + return None + return provider or None + + +def _resolve_destination(request_data: dict) -> str | None: + litellm_params = _as_dict(request_data.get("litellm_params")) + api_base = request_data.get("api_base") or litellm_params.get("api_base") + if not isinstance(api_base, str): + return None + try: + return urlsplit(api_base).hostname + except ValueError: + return None + + +def _resolve_call_surface(logging_obj: LiteLLMLoggingObj | None, request_data: dict) -> str: + call_type = ( + (getattr(logging_obj, "call_type", None) if logging_obj is not None else None) + or request_data.get("call_type") + or request_data.get("litellm_call_type") + ) + return call_type if isinstance(call_type, str) and call_type else "unknown" + + +def _response_finish_reason(response: Any) -> str | None: + choices = getattr(response, "choices", None) + if not isinstance(choices, list): + return None + for choice in choices: + reason = getattr(choice, "finish_reason", None) + if isinstance(reason, str) and reason: + return reason + return None + + +def _build_usage(response: object) -> StraikerWebhookUsage | None: + usage = getattr(response, "usage", None) + if not isinstance(usage, Usage): + return None + input_tokens = usage.prompt_tokens + output_tokens = usage.completion_tokens + if input_tokens is None and output_tokens is None: + return None + return StraikerWebhookUsage(input_tokens=input_tokens, output_tokens=output_tokens) + + +def _is_streamed_request(request_data: dict) -> bool: + if request_data.get("stream") is True: + return True + body = _as_dict(_as_dict(request_data.get("proxy_server_request")).get("body")) + return body.get("stream") is True + + +class StraikerGuardrail(CustomGuardrail): + @staticmethod + def get_config_model() -> type[GuardrailConfigModel]: + return StraikerGuardrailConfigModel + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: + return [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + + def __init__( + self, + api_key: str, + api_base: str = DEFAULT_API_BASE, + source: str = "LiteLLM Gateway", + timeout: float = 5.0, + max_retries: int = 2, + initial_backoff: float = 0.1, + max_backoff: float = 2.0, + unreachable_fallback: Literal["fail_open", "fail_closed"] = "fail_closed", + fail_on_error: bool = True, + max_payload_bytes: int = DEFAULT_MAX_PAYLOAD_BYTES, + custom_headers: dict[str, str] | None = None, + metadata: dict[str, str] | None = None, + verbose: bool = False, + async_handler: httpx.AsyncClient | None = None, + **kwargs: object, + ) -> None: + if not api_key: + raise ValueError("api_key must be non-empty") + if unreachable_fallback not in ("fail_open", "fail_closed"): + raise ValueError(f"unreachable_fallback must be 'fail_open' or 'fail_closed'; got {unreachable_fallback!r}") + + self.api_key = api_key + self.api_base = api_base.rstrip("/") + self.source = source + self.timeout = float(timeout) + self.max_retries = max(0, int(max_retries)) + self.initial_backoff = max(0.0, float(initial_backoff)) + self.max_backoff = max(self.initial_backoff, float(max_backoff)) + self.unreachable_fallback = unreachable_fallback + self.fail_on_error = fail_on_error + self.max_payload_bytes = int(max_payload_bytes) + self.custom_headers = dict(custom_headers) if custom_headers else {} + self.default_metadata = dict(metadata) if metadata else {} + self.verbose = bool(verbose) + + self.streaming_end_of_stream_only = True + self.streaming_buffer_until_moderated = True + + self.async_handler = async_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + ) + + kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) + super().__init__(**kwargs) + + def _webhook_url(self) -> str: + return f"{self.api_base}{WEBHOOK_PATH}" + + def _headers(self) -> dict[str, str]: + reserved = {"authorization", "content-type", "x-straiker-webhook-format"} + extra = {k: v for k, v in self.custom_headers.items() if k.lower() not in reserved} + return { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + "X-Straiker-Webhook-Format": "litellm", + **extra, + } + + def _build_application(self, request_data: dict) -> StraikerWebhookApplication: + meta = _merged_metadata(request_data) + agent_id = _as_optional_str(meta.get("agent_id")) + return StraikerWebhookApplication( + source=agent_id or self.source, + name=_as_optional_str(meta.get("app_name")), + ) + + def _build_context( + self, + request_data: dict, + model: str | None, + logging_obj: LiteLLMLoggingObj | None, + ) -> StraikerWebhookContext: + return StraikerWebhookContext( + call_surface=_resolve_call_surface(logging_obj, request_data), + model=model, + model_provider=_resolve_provider(request_data, model), + destination=_resolve_destination(request_data), + session_id=get_session_id_from_request_data(request_data), + litellm_call_id=getattr(logging_obj, "litellm_call_id", None) if logging_obj else None, + litellm_trace_id=getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None, + litellm_version=litellm_version, + ) + + def _build_envelope( + self, + *, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None, + ) -> StraikerWebhookRequest: + model = inputs.get("model") or request_data.get("model") + call_id = getattr(logging_obj, "litellm_call_id", None) if logging_obj else None + event_id = f"{call_id or 'litellm'}:{input_type}" + + content = StraikerWebhookContent( + texts=list(inputs.get("texts") or []), + images=list(inputs.get("images") or []), + structured_messages=inputs.get("structured_messages"), + tools=inputs.get("tools"), + tool_calls=inputs.get("tool_calls"), + ) + + if input_type == "request": + event = StraikerWebhookEvent(type="pre_call", id=event_id) + return StraikerWebhookRequest( + event=event, + request=content, + context=self._build_context(request_data, model, logging_obj), + identity=_extract_identity(request_data), + application=self._build_application(request_data), + metadata=_build_webhook_metadata(request_data, self.default_metadata), + ) + + response_obj = request_data.get("response") + content.finish_reason = _response_finish_reason(response_obj) + original_messages = request_data.get("messages") + request_content = StraikerWebhookContent( + structured_messages=original_messages if isinstance(original_messages, list) else None, + ) + phase: Literal["none", "assembled"] = "assembled" if _is_streamed_request(request_data) else "none" + event = StraikerWebhookEvent(type="post_call", id=event_id, stream=StraikerWebhookStream(phase=phase)) + return StraikerWebhookRequest( + event=event, + request=request_content, + response=content, + context=self._build_context(request_data, model, logging_obj), + identity=_extract_identity(request_data), + application=self._build_application(request_data), + usage=_build_usage(response_obj), + metadata=_build_webhook_metadata(request_data, self.default_metadata), + ) + + async def _post_webhook(self, payload: dict) -> tuple[StraikerWebhookResponse | None, _WebhookFailure | None]: + try: + body = json.dumps(payload).encode("utf-8") + except (TypeError, ValueError, OverflowError) as error: + return None, _WebhookFailure(f"request serialization failed: {error}", is_unreachable=False) + body_bytes = len(body) + if body_bytes > self.max_payload_bytes: + return None, _WebhookFailure( + f"payload {body_bytes}B exceeds max_payload_bytes {self.max_payload_bytes}", + is_unreachable=False, + ) + + url = self._webhook_url() + headers = self._headers() + attempts = self.max_retries + 1 + last_failure: _WebhookFailure | None = None + + if self.verbose: + verbose_proxy_logger.info( + json.dumps( + { + "event": "straiker.webhook_request", + "url": url, + "bytes": body_bytes, + "payload": payload, + }, + default=str, + ) + ) + + for attempt in range(attempts): + try: + resp = await self.async_handler.post(url, content=body, headers=headers, timeout=self.timeout) + if resp.status_code == 200: + try: + body = resp.json() + parsed = StraikerWebhookResponse.model_validate(body) + except (ValidationError, json.JSONDecodeError) as ve: + return None, _WebhookFailure(f"invalid response schema: {ve}", is_unreachable=False) + if self.verbose: + verbose_proxy_logger.info( + json.dumps( + { + "event": "straiker.webhook_response", + "status_code": resp.status_code, + "body": body, + }, + default=str, + ) + ) + return parsed, None + last_failure = _WebhookFailure( + f"HTTP {resp.status_code}: {resp.text[:200]}", + is_unreachable=resp.status_code in UNREACHABLE_STATUS, + ) + if resp.status_code not in RETRY_STATUS: + return None, last_failure + except (httpx.RequestError, asyncio.TimeoutError, Timeout) as e: + last_failure = _WebhookFailure(f"{type(e).__name__}: {e}", is_unreachable=True) + except (json.JSONDecodeError, TypeError, ValueError) as e: + return None, _WebhookFailure(f"{type(e).__name__}: {e}", is_unreachable=False) + + if attempt < attempts - 1: + backoff = min(self.initial_backoff * (2**attempt), self.max_backoff) + await asyncio.sleep(random.uniform(0, backoff)) + + return None, last_failure or _WebhookFailure("unknown error", is_unreachable=True) + + def _record( + self, + *, + request_data: dict, + logging_obj: LiteLLMLoggingObj | None, + parsed: StraikerWebhookResponse, + ) -> None: + if not self.verbose: + return + response_obj = request_data.get("response") + hidden = getattr(response_obj, "_hidden_params", None) + if isinstance(hidden, dict): + straiker_hidden = hidden.setdefault("straiker", {}) + if isinstance(straiker_hidden, dict): + straiker_hidden.update({"action": parsed.action, "turn_id": parsed.turn_id}) + + def _fail( + self, + *, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + error: str, + is_unreachable: bool, + ) -> GenericGuardrailAPIInputs: + fail_open = (is_unreachable and self.unreachable_fallback == "fail_open") or not self.fail_on_error + verbose_proxy_logger.error( + json.dumps( + { + "event": "straiker.error", + "input_type": input_type, + "error": error, + "fail_open": fail_open, + }, + default=str, + ) + ) + if fail_open: + return inputs + self._block( + request_data=request_data, + input_type=input_type, + message=f"Straiker detection unavailable: {error}", + ) + + def _block( + self, + *, + request_data: dict, + input_type: Literal["request", "response"], + message: str, + ) -> NoReturn: + if input_type == "request": + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name or GUARDRAIL_NAME, + message=message, + should_wrap_with_default_message=False, + ) + raise ModifyResponseException( + message=message, + model=request_data.get("model", "unknown") or "unknown", + request_data=request_data, + guardrail_name=self.guardrail_name or GUARDRAIL_NAME, + original_response=request_data.get("response"), + ) + + @staticmethod + def _intervened_inputs( + inputs: GenericGuardrailAPIInputs, + parsed: StraikerWebhookResponse, + ) -> GenericGuardrailAPIInputs: + return_inputs: GenericGuardrailAPIInputs = {} + return_inputs.update(inputs) + if parsed.texts is not None: + return_inputs["texts"] = parsed.texts + return return_inputs + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + try: + envelope = self._build_envelope( + inputs=inputs, + request_data=request_data, + input_type=input_type, + logging_obj=logging_obj, + ) + payload = envelope.model_dump(mode="json", exclude_none=True) + except (ValidationError, TypeError, ValueError) as error: + return self._fail( + inputs=inputs, + request_data=request_data, + input_type=input_type, + error=str(error), + is_unreachable=False, + ) + + parsed, failure = await self._post_webhook(payload) + if failure is not None: + return self._fail( + inputs=inputs, + request_data=request_data, + input_type=input_type, + error=failure.message, + is_unreachable=failure.is_unreachable, + ) + + if parsed is None: + return self._fail( + inputs=inputs, + request_data=request_data, + input_type=input_type, + error="empty response from Straiker", + is_unreachable=False, + ) + self._record(request_data=request_data, logging_obj=logging_obj, parsed=parsed) + + if parsed.schema_version is not None and parsed.schema_version != STRAIKER_WEBHOOK_SCHEMA_VERSION: + verbose_proxy_logger.warning( + json.dumps( + { + "event": "straiker.schema_drift", + "expected": STRAIKER_WEBHOOK_SCHEMA_VERSION, + "received": parsed.schema_version, + } + ) + ) + + if parsed.action == "BLOCKED": + self._block( + request_data=request_data, + input_type=input_type, + message=parsed.blocked_reason or DEFAULT_BLOCK_MESSAGE, + ) + if parsed.action == "GUARDRAIL_INTERVENED": + is_streamed_response = input_type == "response" and _is_streamed_request(request_data) + if parsed.texts is None or is_streamed_response: + self._block( + request_data=request_data, + input_type=input_type, + message=parsed.blocked_reason or DEFAULT_BLOCK_MESSAGE, + ) + return self._intervened_inputs(inputs, parsed) + return inputs diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8936f6e9ca9..aed345c5db4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -28,6 +28,7 @@ from typing import ( Optional, Set, Tuple, + TypedDict, Union, cast, get_args, @@ -39,6 +40,7 @@ import anyio import websockets import websockets.exceptions from pydantic import BaseModel, Json, JsonValue +from typing_extensions import NotRequired, assert_never from litellm._uuid import uuid from litellm.constants import ( @@ -363,15 +365,15 @@ from litellm.proxy.management_endpoints.cache_settings_endpoints import ( from litellm.proxy.management_endpoints.callback_management_endpoints import ( router as callback_management_endpoints_router, ) -from litellm.proxy.management_endpoints.coordination_redis_endpoints import ( - get_persisted_coordination_redis_settings, - router as coordination_redis_settings_router, -) from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_privileges, _user_has_admin_view, admin_can_invite_user, ) +from litellm.proxy.management_endpoints.coordination_redis_endpoints import ( + get_persisted_coordination_redis_settings, + router as coordination_redis_settings_router, +) from litellm.proxy.management_endpoints.cost_tracking_settings import ( router as cost_tracking_settings_router, ) @@ -1393,19 +1395,25 @@ def _close_dangling_otel_server_span(request: Request, status_code: int, exc: Op if open_telemetry_logger is None: return # Under OTel V2 the FastAPI instrumentor owns the server span (parent_otel_span - # is that same span), and it records the error + ends it itself. Ending it here - # would end it early — losing the http.* attributes the instrumentor stamps on - # completion — and double-end it. Leave it to the instrumentor. + # is that same span) and ends it itself with the http.* attributes stamped on + # completion. The instrumentor only records an error when the exception reaches + # it uncaught, but these handlers swallow it into a JSONResponse, so it never + # does; stamp the error.* attributes here (without ending or re-statusing the + # span, which the instrumentor still owns) so pre-call failures carry the error + # like v1 did. Otherwise close and annotate the dangling span ourselves. try: from litellm.integrations.otel.model.config import is_otel_v2_enabled - if is_otel_v2_enabled(): - return + v2_enabled = is_otel_v2_enabled() except Exception: - pass + v2_enabled = False try: from opentelemetry.trace import Status, StatusCode + if v2_enabled: + if status_code >= 400: + open_telemetry_logger.record_error_attributes_on_span(parent_otel_span, exc, status_code) + return open_telemetry_logger.set_response_status_code_attribute(parent_otel_span, status_code) if status_code >= 400: open_telemetry_logger.record_error_attributes_on_span(parent_otel_span, exc, status_code) @@ -1414,7 +1422,8 @@ def _close_dangling_otel_server_span(request: Request, status_code: int, exc: Op except Exception as e: verbose_proxy_logger.debug("Error closing dangling OTEL SERVER span: %s", str(e)) finally: - request.state.parent_otel_span = None + if not v2_enabled: + request.state.parent_otel_span = None @app.exception_handler(RequestValidationError) @@ -14828,7 +14837,17 @@ async def get_config_general_settings( ) -_GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, dict[str, str]] = { +GeneralSettingsUILiteLLMValue = Union[float, bool, str, None] + + +class GeneralSettingsUILiteLLMFieldSpec(TypedDict): + type: Literal["Float", "Boolean", "Select"] + description: str + options: NotRequired[tuple[str, ...]] + tab: NotRequired[str] # Admin UI sub-tab this field renders under; None groups it with the rest + + +_GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, GeneralSettingsUILiteLLMFieldSpec] = { "budget_exceeded_throttle_percentage": { "type": "Float", "description": ( @@ -14837,18 +14856,60 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: dict[str, dict[str, str]] = { "over-budget keys." ), }, + "enable_anthropic_prompt_caching": { + "type": "Boolean", + "tab": "prompt_caching", + "description": ( + "Auto-adds cache_control to the system prompt and trailing turn for supported Anthropic " + "and Bedrock Claude models. The cache is shared across callers on the same upstream credentials." + ), + }, + "anthropic_prompt_caching_ttl": { + "type": "Select", + "options": ("5m", "1h"), + "tab": "prompt_caching", + "description": "Empty uses Anthropic's 5m default. 1h suits long sessions but doubles the cache write cost.", + }, } -def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> Optional[float]: +def _general_settings_ui_litellm_default( + field_type: Literal["Float", "Boolean", "Select"], +) -> GeneralSettingsUILiteLLMValue: + """The value a field falls back to when it is cleared or reset.""" + return False if field_type == "Boolean" else None + + +def _validate_general_settings_ui_litellm_value(field_name: str, value: Any) -> GeneralSettingsUILiteLLMValue: + spec = _GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name] + field_type = spec["type"] if value is None or value == "": - return None - if isinstance(value, bool) or not isinstance(value, (int, float)) or not (0 < float(value) <= 1): - raise HTTPException( - status_code=400, - detail={"error": f"{field_name} must be a number in (0, 1] or empty"}, - ) - return float(value) + return _general_settings_ui_litellm_default(field_type) + match field_type: + case "Boolean": + if not isinstance(value, bool): + raise HTTPException( + status_code=400, + detail={"error": f"{field_name} must be true or false"}, + ) + return value + case "Select": + options = spec.get("options", ()) + if value not in options: + raise HTTPException( + status_code=400, + detail={"error": f"{field_name} must be one of: {', '.join(options)}, or empty"}, + ) + return cast(str, value) # cast-ok: membership in options proves it is one of the option strings + case "Float": + if isinstance(value, bool) or not isinstance(value, (int, float)) or not (0 < float(value) <= 1): + raise HTTPException( + status_code=400, + detail={"error": f"{field_name} must be a number in (0, 1] or empty"}, + ) + return float(value) + case _: + assert_never(field_type) async def _persist_general_settings_ui_litellm_field( @@ -14869,11 +14930,12 @@ async def _persist_general_settings_ui_litellm_field( async def _reset_general_settings_ui_litellm_field(field_name: str, user_api_key_dict: UserAPIKeyAuth) -> dict: config = await proxy_config.get_config() before_value = config.get("litellm_settings", {}).get(field_name) - setattr(litellm, field_name, None) + default_value = _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS[field_name]["type"]) + setattr(litellm, field_name, default_value) if "litellm_settings" in config: config["litellm_settings"].pop(field_name, None) await proxy_config.save_config(new_config=config) - asyncio.create_task(create_config_audit_log(field_name, "deleted", before_value, None, user_api_key_dict)) + asyncio.create_task(create_config_audit_log(field_name, "deleted", before_value, default_value, user_api_key_dict)) return {"message": f"Field {field_name} reset", "status": "success"} @@ -15041,11 +15103,12 @@ async def get_config_list( else {} ) for litellm_field_name, spec in _GENERAL_SETTINGS_UI_LITELLM_FIELDS.items(): - current_value: Optional[float] = getattr(litellm, litellm_field_name, None) + current_value: GeneralSettingsUILiteLLMValue = getattr(litellm, litellm_field_name, None) + default_value = _general_settings_ui_litellm_default(spec["type"]) stored_in_db_litellm: Optional[bool] if litellm_field_name in db_litellm_settings: stored_in_db_litellm = True - elif current_value is not None: + elif current_value != default_value: stored_in_db_litellm = False else: stored_in_db_litellm = None @@ -15056,7 +15119,9 @@ async def get_config_list( field_description=spec["description"], field_value=current_value, stored_in_db=stored_in_db_litellm, - field_default_value=None, + field_default_value=default_value, + field_options=list(spec.get("options", ())) or None, + field_tab=spec.get("tab"), nested_fields=None, ) ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f842bf13da9..a99cec49417 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -396,6 +396,13 @@ model LiteLLM_MCPUserEnvVars { @@index([server_id]) } +model LiteLLM_MCPServerOAuthClient { + server_id String @id + credentials Json? + created_at DateTime @default(now()) @map("created_at") + updated_at DateTime @default(now()) @updatedAt @map("updated_at") +} + // Generate Tokens for Proxy model LiteLLM_VerificationToken { token String @id diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index b38d5e39800..23e7711b223 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -373,6 +373,12 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs if isinstance(v, BaseModel): v = v.model_dump() additional_usage_values.update({k: v}) + if "cache_read_input_tokens" not in additional_usage_values: + prompt_tokens_details = additional_usage_values.get("prompt_tokens_details") + if isinstance(prompt_tokens_details, dict): + cached_tokens = prompt_tokens_details.get("cached_tokens") + if isinstance(cached_tokens, int) and cached_tokens > 0: + additional_usage_values["cache_read_input_tokens"] = cached_tokens clean_metadata["additional_usage_values"] = additional_usage_values if litellm.cache is not None: diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 7ce4607e1ca..dc2a7d25259 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -77,6 +77,10 @@ class MCPUserCredentialsRepository(PrismaTableRepository): table_name = "litellm_mcpusercredentials" +class MCPServerOAuthClientRepository(PrismaTableRepository): + table_name = "litellm_mcpserveroauthclient" + + class PromptRepository(PrismaTableRepository): table_name = "litellm_prompttable" diff --git a/litellm/router.py b/litellm/router.py index b1a5405ebf1..0b1471dc527 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4461,6 +4461,7 @@ class Router: model=model, request_kwargs=kwargs, messages=kwargs.get("messages", None), + input=kwargs.get("input", None), specific_deployment=kwargs.pop("specific_deployment", None), ) except Exception as e: @@ -4608,6 +4609,7 @@ class Router: deployment = self.get_available_deployment( model=model, messages=kwargs.get("messages", None), + input=kwargs.get("input", None), specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, ) @@ -10002,11 +10004,44 @@ class Router: client = self.cache.get_cache(key=cache_key, parent_otel_span=parent_otel_span) return client + def _count_pre_call_check_tokens( + self, + messages: list[dict[str, str]] | None, + input: str | list | None, + instructions: str | None = None, + ) -> int: + """ + Count input tokens for context-window pre-call checks. + + Chat Completions send `messages`; the Responses API sends `input` (a string or + a list of Responses input items) plus an optional `instructions` system prompt. + The Responses payload is normalized to chat messages via the shared + LiteLLMCompletionResponsesConfig transform so the same token_counter path covers + both API surfaces and `instructions` tokens are included in the count. + """ + if messages is not None: + return litellm.token_counter(messages=messages) + if input is not None: + from openai.types.responses.response_create_params import ResponseInputParam + + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + typed_input = cast(str | ResponseInputParam, input) # cast-ok: str | list matches transform input + input_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=typed_input, + responses_api_request={"instructions": instructions} if instructions is not None else {}, + ) + return litellm.token_counter(messages=cast(list, input_messages)) # cast-ok: transformed chat messages + raise ValueError("Either messages or input must be provided to count tokens") + def _pre_call_checks( self, model: str, healthy_deployments: List, - messages: List[Dict[str, str]], + messages: list[dict[str, str]] | None = None, + input: str | list | None = None, request_kwargs: Optional[dict] = None, ): """ @@ -10036,6 +10071,10 @@ class Router: _rate_limit_error = False parent_otel_span = _get_parent_otel_span_from_kwargs(request_kwargs) + raw_instructions = request_kwargs.get("instructions") if request_kwargs else None + instructions = raw_instructions if isinstance(raw_instructions, str) else None + has_countable_input = messages is not None or input is not None + ## get model group RPM ## dt = get_utc_datetime() current_minute = dt.strftime("%H-%M") @@ -10058,10 +10097,12 @@ class Router: _deployment_model = base_model or _litellm_params.get("model", None) max_input_tokens = model_info.get("max_input_tokens") if isinstance(model_info, dict) else None - if isinstance(max_input_tokens, int): + if isinstance(max_input_tokens, int) and has_countable_input: if input_tokens is None: try: - input_tokens = litellm.token_counter(messages=messages) + input_tokens = self._count_pre_call_check_tokens( + messages=messages, input=input, instructions=instructions + ) except Exception as e: verbose_router_logger.error( "litellm.router.py::_pre_call_checks: failed to count tokens. Returning initial list of deployments. Got - {}".format( @@ -10526,11 +10567,12 @@ class Router: parent_otel_span=parent_otel_span, ) - if self.enable_pre_call_checks and messages is not None: + if self.enable_pre_call_checks and (messages is not None or input is not None): healthy_deployments = self._pre_call_checks( model=model, healthy_deployments=cast(List[Dict], healthy_deployments), messages=messages, + input=input, request_kwargs=request_kwargs, ) # check if user wants to do tag based routing @@ -11041,11 +11083,12 @@ class Router: healthy_deployments = self._filter_blocked_deployments(healthy_deployments) # filter pre-call checks - if self.enable_pre_call_checks and messages is not None: + if self.enable_pre_call_checks and (messages is not None or input is not None): healthy_deployments = self._pre_call_checks( model=model, healthy_deployments=healthy_deployments, messages=messages, + input=input, request_kwargs=request_kwargs, ) @@ -11195,11 +11238,12 @@ class Router: pass_through_deployments = self._filter_blocked_deployments(pass_through_deployments) # 5. Apply pre-call checks (if enabled) - if self.enable_pre_call_checks and messages is not None: + if self.enable_pre_call_checks and (messages is not None or input is not None): pass_through_deployments = self._pre_call_checks( model=model, healthy_deployments=pass_through_deployments, messages=messages, + input=input, request_kwargs=request_kwargs, ) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 86e69467dbf..5b611971154 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -131,6 +131,7 @@ class SupportedGuardrailIntegrations(Enum): SINGULR = "singulr" HEADROOM = "headroom" COMPRESR = "compresr" + STRAIKER = "straiker" class Role(Enum): diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index ac411ad9d9a..377ba669082 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -38,6 +38,7 @@ class MCPAuth(str, enum.Enum): aws_sigv4 = "aws_sigv4" token = "token" oauth2_token_exchange = "oauth2_token_exchange" + oauth2_id_jag = "oauth2_id_jag" true_passthrough = "true_passthrough" oauth_delegate = "oauth_delegate" @@ -62,6 +63,7 @@ MCPAuthType = Optional[ MCPAuth.aws_sigv4, MCPAuth.token, MCPAuth.oauth2_token_exchange, + MCPAuth.oauth2_id_jag, MCPAuth.true_passthrough, MCPAuth.oauth_delegate, ] @@ -159,6 +161,31 @@ class MCPCredentials(TypedDict, total=False): the top-level request field. """ + id_jag_resource_token_endpoint: Optional[str] + """ + Resource authorization server JWT-bearer (RFC 7523) endpoint for ID-JAG leg 2 + """ + + id_jag_resource: Optional[str] + """ + Optional RFC 8707 resource indicator sent on ID-JAG leg 1 + """ + + client_private_key: Optional[str] + """ + PEM private key used to sign the private-key-JWT client_assertion (RFC 7523) + """ + + client_private_key_id: Optional[str] + """ + Key id (kid) advertised in the client_assertion JWT header + """ + + client_assertion_signing_alg: Optional[str] + """ + Signing algorithm for the client_assertion JWT. Default: RS256 + """ + token_endpoint_auth_method: Optional[MCPTokenEndpointAuthMethod] """ How the gateway authenticates to the upstream token endpoint. "client_secret_basic" diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index d0d8cc4cb28..8ae974b19a6 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -87,6 +87,15 @@ class MCPServer(BaseModel): token_exchange_endpoint: Optional[str] = None audience: Optional[str] = None subject_token_type: str = DEFAULT_SUBJECT_TOKEN_TYPE + # ID-JAG fields (draft-ietf-oauth-identity-assertion-authz-grant). + # Leg 1 reuses token_exchange_endpoint (IdP org-AS), audience (resource-AS + # identifier), scopes, subject_token_type, client_id/client_secret. Leg 2 + # posts the ID-JAG assertion to id_jag_resource_token_endpoint. + id_jag_resource_token_endpoint: Optional[str] = None + id_jag_resource: Optional[str] = None + client_private_key: Optional[str] = None + client_private_key_id: Optional[str] = None + client_assertion_signing_alg: str = "RS256" # Wire dialect: "rfc8693" (standard token-exchange grant) or "entra_obo" (Microsoft Entra # On-Behalf-Of, the RFC 7523 jwt-bearer grant + requested_token_use extension) token_exchange_profile: str = "rfc8693" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/straiker.py b/litellm/types/proxy/guardrails/guardrail_hooks/straiker.py new file mode 100644 index 00000000000..b4375237917 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/straiker.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field + +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk +from litellm.types.utils import ChatCompletionMessageToolCall + +from .base import GuardrailConfigModel + +StraikerWebhookEventType = Literal["pre_call", "post_call"] +StraikerWebhookStreamPhase = Literal["none", "assembled"] +StraikerWebhookAction = Literal["NONE", "BLOCKED", "GUARDRAIL_INTERVENED"] + +STRAIKER_WEBHOOK_SCHEMA_VERSION = "1" + + +class StraikerWebhookStream(BaseModel): + phase: StraikerWebhookStreamPhase = "none" + index: int | None = None + + +class StraikerWebhookEvent(BaseModel): + type: StraikerWebhookEventType + id: str + stream: StraikerWebhookStream = Field(default_factory=StraikerWebhookStream) + + +class StraikerWebhookContent(BaseModel): + model_config = ConfigDict(extra="ignore") + + texts: list[str] = Field(default_factory=list) + images: list[str] = Field(default_factory=list) + structured_messages: list[AllMessageValues] | None = None + tools: list[dict[str, object]] | None = None + tool_calls: list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None = None + finish_reason: str | None = None + + +class StraikerWebhookUsage(BaseModel): + input_tokens: int | None = None + output_tokens: int | None = None + + +class StraikerWebhookContext(BaseModel): + call_surface: str + model: str | None = None + model_provider: str | None = None + destination: str | None = None + session_id: str | None = None + litellm_call_id: str | None = None + litellm_trace_id: str | None = None + litellm_version: str | None = None + + +class StraikerWebhookIdentity(BaseModel): + litellm_key: str | None = None + litellm_team: str | None = None + litellm_user_id: str | None = None + litellm_user_email: str | None = None + litellm_org_id: str | None = None + end_user_id: str | None = None + + +class StraikerWebhookApplication(BaseModel): + source: str + name: str | None = None + + +class StraikerWebhookRequest(BaseModel): + schema_version: str = STRAIKER_WEBHOOK_SCHEMA_VERSION + event: StraikerWebhookEvent + request: StraikerWebhookContent + response: StraikerWebhookContent | None = None + context: StraikerWebhookContext + identity: StraikerWebhookIdentity + application: StraikerWebhookApplication + usage: StraikerWebhookUsage | None = None + metadata: dict[str, object] | None = None + + +class StraikerWebhookResponse(BaseModel): + model_config = ConfigDict(extra="allow") + + action: StraikerWebhookAction = "NONE" + blocked_reason: str | None = None + texts: list[str] | None = None + schema_version: str | None = None + turn_id: str | None = Field(default=None, alias="turnId") + + +class StraikerGuardrailConfigModelOptionalParams(BaseModel): + timeout: float | None = Field( + default=5.0, + gt=0.0, + description="Per-attempt HTTP timeout in seconds.", + ) + max_retries: int | None = Field( + default=2, + ge=0, + description="Retries on transient HTTP (408/429/5xx) and network errors.", + ) + initial_backoff: float | None = Field( + default=0.1, + ge=0.0, + description="Initial retry backoff in seconds.", + ) + max_backoff: float | None = Field( + default=2.0, + ge=0.0, + description="Maximum retry backoff in seconds.", + ) + unreachable_fallback: Literal["fail_open", "fail_closed"] | None = Field( + default="fail_closed", + description="Behavior when Straiker is unreachable after retries.", + ) + fail_on_error: bool | None = Field( + default=True, + description=( + "Behavior on any guardrail error, not just unreachability. True (default) blocks " + "the request on error; False logs and allows the request to proceed." + ), + ) + max_payload_bytes: int | None = Field( + default=524288, + gt=0, + description="Maximum serialized webhook payload size sent to Straiker.", + ) + custom_headers: dict[str, str] | None = Field( + default=None, + description="Additional HTTP headers sent to Straiker, excluding Authorization and the webhook-format header.", + ) + metadata: dict[str, str] | None = Field( + default=None, + description=( + "Default metadata key/values added to the webhook metadata bag on every request. " + "On key conflict with request-derived metadata, these configured values win." + ), + ) + verbose: bool | None = Field( + default=False, + description="Log webhook request/response payloads and record action/turn_id in response hidden params.", + ) + + +class StraikerGuardrailConfigModel(GuardrailConfigModel[StraikerGuardrailConfigModelOptionalParams]): + api_key: str = Field( + min_length=1, + description="Straiker DefendAI environment API key (Bearer token). Env: STRAIKER_API_KEY.", + json_schema_extra={"secret": True}, + ) + + api_base: str | None = Field( + default="https://api.prod.straiker.ai", + description="Straiker API base URL. Use the regional variant for non-US tenants.", + ) + + default_app: str | None = Field( + default="LiteLLM Gateway", + description=( + "Default application registered in the Straiker Defend Console. " + "Overridden per-request by metadata.agent_id when present." + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Straiker" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b1a87c444c8..5cf99ba8bac 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -16272,7 +16272,7 @@ "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/glm-5p2": { - "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -16686,7 +16686,7 @@ "supports_vision": false }, "fireworks_ai/glm-5p2": { - "cache_read_input_token_cost": 2.6e-07, + "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, diff --git a/pyproject.toml b/pyproject.toml index 108f28cb124..769a1dea469 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -189,6 +189,7 @@ dev = [ e2e-dev = [ "playwright==1.61.0", "websockets>=15.0.1,<16.0", + "locust==2.45.0", ] proxy-dev = [ "prisma==0.11.0", diff --git a/schema.prisma b/schema.prisma index f842bf13da9..a99cec49417 100644 --- a/schema.prisma +++ b/schema.prisma @@ -396,6 +396,13 @@ model LiteLLM_MCPUserEnvVars { @@index([server_id]) } +model LiteLLM_MCPServerOAuthClient { + server_id String @id + credentials Json? + created_at DateTime @default(now()) @map("created_at") + updated_at DateTime @default(now()) @updatedAt @map("updated_at") +} + // Generate Tokens for Proxy model LiteLLM_VerificationToken { token String @id diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 67e9f4f78a7..58a330775c0 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -17,6 +17,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) +- `load/` - throughput/performance under concurrency: drives real concurrent traffic through the whole stack with Locust and asserts a throughput SLO; marked `load` so the parent conftest collects it last and it never perturbs latency-sensitive suites - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees, and does not use the shared transport harness @@ -33,7 +34,7 @@ class TestPromptCompression: def test_prompt_compression_accumulate_spend(self, key_id, user_id): for _ in range(10): - response = self.resources.gateway.post("gemini-2.5-flash", key_id, user_id) + response = self.resources.proxy.post("gemini-2.5-flash", key_id, user_id) compressed_value = ... assert response.cost == compressed_value # the cost was actually reduced ``` @@ -48,9 +49,9 @@ The shape is layered so tests stay declarative `transport.py` exposes a `Transport` Protocol with `post`, `get`, `delete`, `send`, `stream`, `probe`, plus `bearer(key)` and the `master` header. `HttpTransport` fulfils it, and `SplitTransport` routes each call by path to the data plane or the control plane so a split control-plane/data-plane deployment works without any change in the test -`e2e_gateway.py` holds `Gateway`, a frozen dataclass that wraps a `Transport` and adds the operations tests reuse: `generate_key` / `delete_key` / `key_info`, `model_info`, the LLM calls `chat` / `chat_stream` / `embed` / `ocr`, the spend read-back `spend_logs`, and the poll helpers `poll_logs_for_key` / `poll_logs_for_request_id` that loop to `poll_timeout` instead of sleeping once. Add a new route as a method here so other suites get it for free +`proxy_client.py` holds `ProxyClient`, a frozen dataclass that wraps a `Transport` and adds the operations tests reuse: `generate_key` / `delete_key` / `key_info`, `model_info`, the LLM calls `chat` / `chat_stream` / `embed` / `ocr`, the spend read-back `spend_logs`, and the poll helpers `poll_logs_for_key` / `poll_logs_for_request_id` that loop to `poll_timeout` instead of sleeping once. It is exposed as the session-scoped `proxy` fixture (see tests/e2e/conftest.py), which each suite's `client` fixture depends on and injects. Add a new route as a method here so other suites get it for free -Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `Gateway` and adds suite-specific routes. Cleanup runs through that same `Gateway`, so whatever keys or customers your test creates get torn down by the `resources` fixture +Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `ProxyClient` (as `.proxy`) and adds suite-specific routes. Cleanup runs through that same `ProxyClient`, so whatever keys or customers your test creates get torn down by the `resources` fixture Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 555ac0482e2..49d776cd64b 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -113,7 +113,7 @@ class TestPromptCompression: def test_prompt_compression_accumulate_spend(self, key_id, user_id): for _ in range(10): - response = self.resources.gateway.post("gemini-2.5-flash", key_id, user_id) + response = self.resources.proxy.post("gemini-2.5-flash", key_id, user_id) compressed_value = ... assert response.cost == compressed_value # the cost was actually reduced ``` @@ -128,9 +128,9 @@ The shape is layered so tests stay declarative `transport.py` exposes a `Transport` Protocol with `post`, `get`, `delete`, `send`, `stream`, `probe`, plus `bearer(key)` and the `master` header. `HttpTransport` fulfils it, and `SplitTransport` routes each call by path to the data plane or the control plane so a split control-plane/data-plane deployment works without any change in the test -`e2e_gateway.py` holds `Gateway`, a frozen dataclass that wraps a `Transport` and adds the operations tests reuse: `generate_key` / `delete_key` / `key_info`, `model_info`, the LLM calls `chat` / `chat_stream` / `embed` / `ocr`, the spend read-back `spend_logs`, and the poll helpers `poll_logs_for_key` / `poll_logs_for_request_id` that loop to `poll_timeout` instead of sleeping once. Add a new route as a method here so other suites get it for free +`proxy_client.py` holds `ProxyClient`, a frozen dataclass that wraps a `Transport` and adds the operations tests reuse: `generate_key` / `delete_key` / `key_info`, `model_info`, the LLM calls `chat` / `chat_stream` / `embed` / `ocr`, the spend read-back `spend_logs`, and the poll helpers `poll_logs_for_key` / `poll_logs_for_request_id` that loop to `poll_timeout` instead of sleeping once. It is exposed as the session-scoped `proxy` fixture (see tests/e2e/conftest.py), which each suite's `client` fixture depends on and injects. Add a new route as a method here so other suites get it for free -Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `Gateway` and adds suite-specific routes. Cleanup runs through that same `Gateway`, so whatever keys or customers your test creates get torn down by the `resources` fixture +Each suite provides its own `client` fixture (see `llm_translation/passthrough_client.py`), a frozen dataclass that holds the shared `ProxyClient` (as `.proxy`) and adds suite-specific routes. Cleanup runs through that same `ProxyClient`, so whatever keys or customers your test creates get torn down by the `resources` fixture Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass diff --git a/tests/e2e/access_control/access_control_client.py b/tests/e2e/access_control/access_control_client.py index d7bc9c280aa..e95ad1f57ce 100644 --- a/tests/e2e/access_control/access_control_client.py +++ b/tests/e2e/access_control/access_control_client.py @@ -4,7 +4,7 @@ from __future__ import annotations from dataclasses import dataclass -from e2e_gateway import Gateway, build_gateway +from proxy_client import ProxyClient from e2e_http import StreamingResponse from models import ( ChatBody, @@ -21,29 +21,29 @@ ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route" @dataclass(frozen=True, slots=True) class AccessControlClient: - gateway: Gateway + proxy: ProxyClient def llm_only_key(self) -> str: - return self.gateway.generate_key( + return self.proxy.generate_key( KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"]) ) def delete_key(self, key: str) -> None: - self.gateway.delete_key(key) + self.proxy.delete_key(key) def chat_status(self, key: str, model: str, content: str) -> StreamingResponse: - return self.gateway.transport.send( + return self.proxy.transport.send( "/chat/completions", - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), json=ChatBody( model=model, messages=[ChatMessage(role="user", content=content)] ), ) def create_model_status(self, key: str, model_name: str) -> StreamingResponse: - return self.gateway.transport.send( + return self.proxy.transport.send( "/model/new", - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), json=ModelNewBody( model_name=model_name, litellm_params=LiteLLMParamsBody(model="openai/gpt-4o-mini"), @@ -52,5 +52,5 @@ class AccessControlClient: ) -def build_client() -> AccessControlClient: - return AccessControlClient(gateway=build_gateway()) +def build_client(proxy: ProxyClient) -> AccessControlClient: + return AccessControlClient(proxy=proxy) diff --git a/tests/e2e/access_control/conftest.py b/tests/e2e/access_control/conftest.py index b5681ff76ad..d7299014d08 100644 --- a/tests/e2e/access_control/conftest.py +++ b/tests/e2e/access_control/conftest.py @@ -3,8 +3,9 @@ import pytest from access_control_client import AccessControlClient, build_client +from proxy_client import ProxyClient @pytest.fixture(scope="session") -def client() -> AccessControlClient: - return build_client() +def client(proxy: ProxyClient) -> AccessControlClient: + return build_client(proxy) diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 4debf50bd6c..ca48204962a 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -68,7 +68,7 @@ File delete asserts `object=="file"` and `deleted==True`. | File | Covers | |------|--------| -| `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared Gateway; runtime batch model registration via /model/new; denial helpers | +| `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared ProxyClient; runtime batch model registration via /model/new; denial helpers | | `capabilities.py` | the provider x scenario matrix + per-provider /model/new params + id-shape classifiers + per-provider raw-id assertion | | `conftest.py` | session-scoped batch deployment registration and teardown | | `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial | diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index fda4e87e478..7db5d0b6beb 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -1,5 +1,5 @@ """Client for the batches e2e suite: file upload/download and the batch -operations (create / retrieve / cancel / list) over the shared Gateway. +operations (create / retrieve / cancel / list) over the shared ProxyClient. Batch deployments are registered at runtime via /model/new (see conftest.py), not baked into the proxy config. `create_batch` returns the raw HTTP outcome @@ -16,7 +16,7 @@ from dataclasses import dataclass from pydantic import BaseModel -from e2e_gateway import Gateway, build_gateway +from proxy_client import ProxyClient from e2e_http import ( FileUploadForm, NoBody, @@ -85,13 +85,13 @@ def is_result_access_denied[R: BaseModel](result: Result[R]) -> bool: @dataclass(frozen=True, slots=True) class BatchClient: - gateway: Gateway + proxy: ProxyClient def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str: - return self.gateway.create_model(model_name, litellm_params, mode="batch") + return self.proxy.create_model(model_name, litellm_params, mode="batch") def delete_model(self, model_id: str) -> None: - self.gateway.delete_model(model_id) + self.proxy.delete_model(model_id) def upload_file( self, @@ -102,9 +102,9 @@ class BatchClient: model: str | None = None, provider: str | None = None, ) -> Result[FileObject]: - return self.gateway.transport.upload( + return self.proxy.transport.upload( _files_path(provider), - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), form=form, filename="batch_input.jsonl", content=content, @@ -115,18 +115,18 @@ class BatchClient: def create_batch( self, *, body: BatchCreateBody, key: str, provider: str | None = None ) -> StreamingResponse: - return self.gateway.transport.send( + return self.proxy.transport.send( _batches_path(provider), - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), json=body, ) def retrieve_batch( self, batch_id: str, *, key: str, provider: str | None = None ) -> Result[BatchObject]: - return self.gateway.transport.get( + return self.proxy.transport.get( f"{_batches_path(provider)}/{batch_id}", - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), params=NoBody(), response_type=BatchObject, ) @@ -134,9 +134,9 @@ class BatchClient: def cancel_batch( self, batch_id: str, *, key: str, provider: str | None = None ) -> Result[BatchObject]: - return self.gateway.transport.post( + return self.proxy.transport.post( f"{_batches_path(provider)}/{batch_id}/cancel", - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), json=NoBody(), response_type=BatchObject, ) @@ -144,9 +144,9 @@ class BatchClient: def list_batches( self, *, key: str, provider: str | None = None ) -> Result[BatchList]: - return self.gateway.transport.get( + return self.proxy.transport.get( _batches_path(provider), - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), params=NoBody(), response_type=BatchList, ) @@ -154,9 +154,9 @@ class BatchClient: def delete_file( self, file_id: str, *, key: str, provider: str | None = None ) -> Result[FileDeleteResponse]: - return self.gateway.transport.delete( + return self.proxy.transport.delete( f"{_files_path(provider)}/{file_id}", - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), json=NoBody(), response_type=FileDeleteResponse, ) @@ -170,5 +170,5 @@ def _batches_path(provider: str | None) -> str: return f"/{provider}/v1/batches" if provider else "/v1/batches" -def build_client() -> BatchClient: - return BatchClient(gateway=build_gateway()) +def build_client(proxy: ProxyClient) -> BatchClient: + return BatchClient(proxy=proxy) diff --git a/tests/e2e/batches/conftest.py b/tests/e2e/batches/conftest.py index d3b6d42bc24..73e8918e2ee 100644 --- a/tests/e2e/batches/conftest.py +++ b/tests/e2e/batches/conftest.py @@ -1,7 +1,7 @@ """Batches suite's `client` fixture. The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker -live in the parent tests/e2e/conftest.py. BatchClient holds the shared Gateway, so +live in the parent tests/e2e/conftest.py. BatchClient holds the shared ProxyClient, so the `resources` fixture cleans up keys through it; tests register file deletes and batch cancels via `resources.defer(...)`. @@ -19,6 +19,7 @@ import pytest from batch_client import BatchClient, build_client from capabilities import PROVIDERS from e2e_http import NoBody +from proxy_client import ProxyClient def pytest_configure(config: pytest.Config) -> None: @@ -29,13 +30,13 @@ def pytest_configure(config: pytest.Config) -> None: @pytest.fixture(scope="session") -def client() -> BatchClient: - return build_client() +def client(proxy: ProxyClient) -> BatchClient: + return build_client(proxy) @pytest.fixture(scope="session") def batch_deployments(client: BatchClient) -> Iterator[None]: - probe = client.gateway.probe("/health/liveliness", params=NoBody()) + probe = client.proxy.probe("/health/liveliness", params=NoBody()) if not probe.healthy: yield return diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 2ee7eb36a41..b483b420c84 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -372,17 +372,17 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( environment and OOMed the e2e runner on stage. """ user_id = f"e2e-batch-rl-{unique_marker()}" - key = client.gateway.generate_key( + key = client.proxy.generate_key( KeyGenerateBody(models=[], tpm_limit=1_000_000, rpm_limit=1_000, user_id=user_id) ) - resources.defer(lambda: client.gateway.delete_key(key)) + resources.defer(lambda: client.proxy.delete_key(key)) window_start = datetime.now(timezone.utc) - timedelta(hours=1) window_end = window_start + timedelta(hours=2) before = frozenset( row.request_id for row in unattributed_rows( - client.gateway.spend_logs_window(start=window_start, end=window_end) + client.proxy.spend_logs_window(start=window_start, end=window_end) ) ) @@ -401,12 +401,12 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( batch = BatchObject.model_validate_json(created.body) resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) - _ = client.gateway.poll_logs_for_key(key, min_rows=1) + _ = client.proxy.poll_logs_for_key(key, min_rows=1) new_orphans = [ row for row in unattributed_rows( - client.gateway.spend_logs_window(start=window_start, end=window_end) + client.proxy.spend_logs_window(start=window_start, end=window_end) ) if row.request_id not in before ] diff --git a/tests/e2e/claude_code/conftest.py b/tests/e2e/claude_code/conftest.py index 8f5c09fa4e2..aaad7667936 100644 --- a/tests/e2e/claude_code/conftest.py +++ b/tests/e2e/claude_code/conftest.py @@ -577,30 +577,30 @@ from claude_code._compat_models import ( # noqa: E402 ) -def _build_control_gateway(proxy: ProxyConfig): +def _build_control_plane_client(proxy_config: ProxyConfig): """Local import of the shared harness so the pure-unit-test tree under ``_driver_unit_tests/`` etc. never has to pull it in. The control plane transport is what /model/new lives on; SplitTransport routes it correctly for both monolithic and split deployments. - The endpoints come from the *resolved* proxy, not from a second + The endpoints come from the *resolved* proxy config, not from a second independent env read, so registration and the cells always hit the same host and key. Both planes get the one URL the cells use; the deployment is fronted by a single address that routes management and LLM paths itself.""" - from e2e_gateway import build_gateway + from proxy_client import build_proxy_client - return build_gateway( - base_url=proxy.base_url, - master_key=proxy.api_key, - control_plane_base_url=proxy.base_url, + return build_proxy_client( + base_url=proxy_config.base_url, + master_key=proxy_config.api_key, + control_plane_base_url=proxy_config.base_url, ) -def _register_deployment(gateway, deployment: CompatDeployment) -> str: +def _register_deployment(proxy, deployment: CompatDeployment) -> str: """Register one deployment and return its proxy-assigned model_id once it is servable on the data plane.""" - return gateway.create_model( + return proxy.create_model( deployment.model_name, deployment.litellm_params, ) @@ -624,20 +624,20 @@ def _compat_models_registered() -> Any: but do not abort the session: the cells that need that specific deployment will 400 with "Invalid model name" and fail loudly, which is the right signal (missing cred on the proxy side).""" - proxy = resolve_proxy() - if proxy is None: + proxy_config = resolve_proxy() + if proxy_config is None: yield return from requests import RequestException - gateway = _build_control_gateway(proxy) + proxy = _build_control_plane_client(proxy_config) registered_ids: list[str] = [] failures: list[tuple[str, str]] = [] try: for deployment in load_all_deployments(): try: - model_id = _register_deployment(gateway, deployment) + model_id = _register_deployment(proxy, deployment) registered_ids.append(model_id) except (AssertionError, RequestException) as exc: failures.append((deployment.model_name, str(exc))) @@ -656,7 +656,7 @@ def _compat_models_registered() -> Any: finally: for model_id in registered_ids: try: - gateway.delete_model(model_id) + proxy.delete_model(model_id) except (AssertionError, RequestException): # Best-effort — teardown surfaces via warnings inside # ``delete_model`` already; swallowing here so one flaky diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 88a9deecb7e..5347fffca4d 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -23,7 +23,8 @@ import requests from e2e_config import CONTROL_PLANE_BASE_URL, PROXY_BASE_URL from junit_properties import attach_result_properties -from lifecycle import GatewayProvider, ResourceManager +from lifecycle import ProxyClientProvider, ResourceManager +from proxy_client import ProxyClient, build_proxy_client _E2E_TEST_RAN = pytest.StashKey[bool]() @@ -38,6 +39,10 @@ def pytest_configure(config: pytest.Config) -> None: "markers", "covers(cell_id, *, exercised_on=()): coverage-registry cell(s) this test covers", ) + config.addinivalue_line( + "markers", + "load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites", + ) def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: @@ -46,9 +51,13 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: as `` entries, on every outcome including skips and setup errors. Downstream (Loki/Grafana) reads outcome and duration from the standard report and these properties for package rollups and coverage drill-down. See - junit_properties.py.""" + junit_properties.py. + + Also sort `load`-marked items last so a whole-tree run drives heavy throughput + traffic only after the latency-sensitive suites have finished.""" for item in items: attach_result_properties(item) + items.sort(key=lambda item: item.get_closest_marker("load") is not None) def _liveness_reason(label: str, base_url: str) -> str | None: @@ -120,11 +129,18 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None: sys.path.remove(spend_dir) +@pytest.fixture(scope="session") +def proxy() -> ProxyClient: + """The shared ProxyClient every suite's client is built from. Suite `client` + fixtures depend on this and inject it, so the proxy wiring lives in one place.""" + return build_proxy_client() + + @pytest.fixture -def resources(client: GatewayProvider) -> Iterator[ResourceManager]: +def resources(client: ProxyClientProvider) -> Iterator[ResourceManager]: """init -> run -> teardown: create a manager, run the test, release resources. - Cleanup goes through the shared Gateway, whatever the suite's client adds.""" - manager = ResourceManager(client=client.gateway) + Cleanup goes through the shared ProxyClient, whatever the suite's client adds.""" + manager = ResourceManager(client=client.proxy) manager.init() yield manager manager.teardown() diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 529744d5a2c..2d0ad93e53d 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -61,6 +61,12 @@ POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120")) POLL_INTERVAL = float(os.environ.get("E2E_POLL_INTERVAL", "5")) REQUEST_TIMEOUT = float(os.environ.get("E2E_REQUEST_TIMEOUT", "60")) +LOAD_USERS = int(os.environ.get("E2E_LOAD_USERS", "750")) +LOAD_SPAWN_RATE = float(os.environ.get("E2E_LOAD_SPAWN_RATE", "50")) +LOAD_DURATION_SECONDS = float(os.environ.get("E2E_LOAD_DURATION_SECONDS", "60")) +LOAD_MIN_RPS = float(os.environ.get("E2E_LOAD_MIN_RPS", "355")) +LOAD_MAX_FAILURE_RATIO = float(os.environ.get("E2E_LOAD_MAX_FAILURE_RATIO", "0.01")) + def unique_marker() -> str: """A short unique token per call/run, so concurrent runs and the shared diff --git a/tests/e2e/lifecycle.py b/tests/e2e/lifecycle.py index b15986987a7..c0de074f0d2 100644 --- a/tests/e2e/lifecycle.py +++ b/tests/e2e/lifecycle.py @@ -13,7 +13,7 @@ the test body is run(), and the fixture's teardown is teardown(). from dataclasses import dataclass, field from typing import Callable, List, Protocol, runtime_checkable -from e2e_gateway import Gateway +from proxy_client import ProxyClient from models import KeyGenerateBody @@ -52,7 +52,7 @@ def run_case(case: E2ECase) -> None: @runtime_checkable class ResourceClient(Protocol): """Proxy operations the convenience creators use. Resource types without a - creator here are handled generically via ResourceManager.defer(). The Gateway + creator here are handled generically via ResourceManager.defer(). The ProxyClient satisfies this.""" def generate_key(self, body: KeyGenerateBody) -> str: ... @@ -63,12 +63,12 @@ class ResourceClient(Protocol): @runtime_checkable -class GatewayProvider(Protocol): - """Every suite's client exposes the shared Gateway, which the resources fixture +class ProxyClientProvider(Protocol): + """Every suite's client exposes the shared ProxyClient, which the resources fixture uses for cleanup. The client adds its own route methods on top.""" @property - def gateway(self) -> Gateway: ... + def proxy(self) -> ProxyClient: ... @dataclass diff --git a/tests/e2e/llm_translation/conftest.py b/tests/e2e/llm_translation/conftest.py index 5258b751a8c..f35ecf0760d 100644 --- a/tests/e2e/llm_translation/conftest.py +++ b/tests/e2e/llm_translation/conftest.py @@ -2,13 +2,14 @@ The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. PassthroughClient holds the shared -Gateway, so the `resources` fixture cleans up keys this suite creates. +ProxyClient, so the `resources` fixture cleans up keys this suite creates. """ import pytest from endpoints_client import EndpointsClient, build_endpoints_client from passthrough_client import PassthroughClient, build_client +from proxy_client import ProxyClient def pytest_configure(config: pytest.Config) -> None: @@ -19,10 +20,10 @@ def pytest_configure(config: pytest.Config) -> None: @pytest.fixture(scope="session") -def client() -> PassthroughClient: - return build_client() +def client(proxy: ProxyClient) -> PassthroughClient: + return build_client(proxy) @pytest.fixture(scope="session") -def endpoints_client() -> EndpointsClient: - return build_endpoints_client() +def endpoints_client(proxy: ProxyClient) -> EndpointsClient: + return build_endpoints_client(proxy) diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index c3816928564..e339922b4d1 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from pydantic import BaseModel -from e2e_gateway import Gateway, build_gateway +from proxy_client import ProxyClient from e2e_http import StreamingResponse from models import ChatMessage, LiteLLMParamsBody @@ -156,17 +156,17 @@ class ImagesResult(BaseModel): @dataclass(frozen=True, slots=True) class EndpointsClient: - gateway: Gateway + proxy: ProxyClient def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str: - return self.gateway.create_model(model_name, litellm_params) + return self.proxy.create_model(model_name, litellm_params) def delete_model(self, model_id: str) -> None: - self.gateway.delete_model(model_id) + self.proxy.delete_model(model_id) def _send(self, path: str, key: str, body: BaseModel) -> StreamingResponse: - return self.gateway.transport.send( - path, headers=self.gateway.transport.bearer(key), json=body + return self.proxy.transport.send( + path, headers=self.proxy.transport.bearer(key), json=body ) def responses(self, key: str, model: str, text: str) -> StreamingResponse: @@ -216,5 +216,5 @@ class EndpointsClient: ) -def build_endpoints_client() -> EndpointsClient: - return EndpointsClient(gateway=build_gateway()) +def build_endpoints_client(proxy: ProxyClient) -> EndpointsClient: + return EndpointsClient(proxy=proxy) diff --git a/tests/e2e/llm_translation/passthrough_client.py b/tests/e2e/llm_translation/passthrough_client.py index 77dcf68a1e3..0576321ede1 100644 --- a/tests/e2e/llm_translation/passthrough_client.py +++ b/tests/e2e/llm_translation/passthrough_client.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from pydantic import BaseModel, Field -from e2e_gateway import Gateway, build_gateway +from proxy_client import ProxyClient from e2e_http import Headers, StreamingResponse from models import ChatMessage @@ -108,7 +108,7 @@ def _tags_header(tags: list[str] | None) -> str | None: @dataclass(frozen=True, slots=True) class PassthroughClient: - gateway: Gateway + proxy: ProxyClient # ---- Gemini native passthrough (/gemini/v1beta/...) ----------------- @@ -121,7 +121,7 @@ class PassthroughClient: tools: list[GeminiTool] | None = None, tags: list[str] | None = None, ) -> StreamingResponse: - return self.gateway.transport.send( + return self.proxy.transport.send( f"/gemini/v1beta/models/{model}:generateContent", headers=GeminiHeaders(x_goog_api_key=key, tags=_tags_header(tags)), json=GeminiGenerateBody( @@ -132,7 +132,7 @@ class PassthroughClient: def gemini_stream( self, key: str, model: str, text: str, *, tags: list[str] | None = None ) -> StreamingResponse: - return self.gateway.transport.send( + return self.proxy.transport.send( f"/gemini/v1beta/models/{model}:streamGenerateContent", headers=GeminiHeaders(x_goog_api_key=key, tags=_tags_header(tags)), json=GeminiGenerateBody( @@ -151,7 +151,7 @@ class PassthroughClient: f"/vertex_ai/v1/projects/{project}/locations/{location}" f"/publishers/google/models/{model}:generateContent" ) - return self.gateway.transport.send( + return self.proxy.transport.send( path, headers=VertexHeaders(x_litellm_api_key=key), json=GeminiGenerateBody( @@ -172,7 +172,7 @@ class PassthroughClient: stream: bool = False, tags: list[str] | None = None, ) -> StreamingResponse: - return self.gateway.transport.send( + return self.proxy.transport.send( "/anthropic/v1/messages", headers=AnthropicHeaders(x_api_key=key, tags=_tags_header(tags)), json=AnthropicMessageBody( @@ -186,5 +186,5 @@ class PassthroughClient: ) -def build_client() -> PassthroughClient: - return PassthroughClient(gateway=build_gateway()) +def build_client(proxy: ProxyClient) -> PassthroughClient: + return PassthroughClient(proxy=proxy) diff --git a/tests/e2e/llm_translation/realtime/conftest.py b/tests/e2e/llm_translation/realtime/conftest.py index 8e6e596bcd3..752737e830e 100644 --- a/tests/e2e/llm_translation/realtime/conftest.py +++ b/tests/e2e/llm_translation/realtime/conftest.py @@ -1,7 +1,7 @@ """Realtime suite's `client` and `realtime_models` fixtures. The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker -live in the parent tests/e2e/conftest.py. RealtimeClient holds the shared Gateway, +live in the parent tests/e2e/conftest.py. RealtimeClient holds the shared ProxyClient, so the `resources` fixture cleans up keys this suite creates. `realtime_models` registers every provider's realtime deployment through /model/new @@ -15,11 +15,12 @@ from collections.abc import Iterator import pytest from realtime_client import PROVIDERS, RealtimeClient, build_client +from proxy_client import ProxyClient @pytest.fixture(scope="session") -def client() -> RealtimeClient: - return build_client() +def client(proxy: ProxyClient) -> RealtimeClient: + return build_client(proxy) @pytest.fixture(scope="session") @@ -34,4 +35,4 @@ def realtime_models(client: RealtimeClient) -> Iterator[dict[str, str]]: yield {provider_id: model_name for provider_id, model_name, _ in records} finally: for _, _, model_id in records: - client.gateway.delete_model(model_id) + client.proxy.delete_model(model_id) diff --git a/tests/e2e/llm_translation/realtime/realtime_client.py b/tests/e2e/llm_translation/realtime/realtime_client.py index b50160e3538..e6c5c19cbd1 100644 --- a/tests/e2e/llm_translation/realtime/realtime_client.py +++ b/tests/e2e/llm_translation/realtime/realtime_client.py @@ -22,7 +22,7 @@ from websockets.sync.client import connect from websockets.sync.connection import Connection from e2e_config import PROXY_BASE_URL, unique_marker -from e2e_gateway import Gateway, build_gateway +from proxy_client import ProxyClient from models import LiteLLMParamsBody _M = TypeVar("_M", bound=BaseModel) @@ -329,7 +329,7 @@ class RealtimeSession: @dataclass(frozen=True, slots=True) class RealtimeClient: - gateway: Gateway + proxy: ProxyClient def provision(self, provider: RealtimeProvider) -> tuple[str, str]: """Register this provider's realtime deployment through /model/new and return @@ -338,7 +338,7 @@ class RealtimeClient: show up as a realtime model on /model/info. add_deployment runs synchronously, so the deployment is connectable as soon as this returns.""" model_name = f"{provider.alias}-{unique_marker()}" - model_id = self.gateway.create_model( + model_id = self.proxy.create_model( model_name, provider.litellm_params, mode="realtime" ) return model_name, model_id @@ -355,5 +355,5 @@ class RealtimeClient: yield RealtimeSession(connection=connection) -def build_client() -> RealtimeClient: - return RealtimeClient(gateway=build_gateway()) +def build_client(proxy: ProxyClient) -> RealtimeClient: + return RealtimeClient(proxy=proxy) diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index ba701f4869c..a2c17b0fb66 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -81,9 +81,9 @@ def _cache_chat( RichMessage(role="user", content=[TextBlock(text="Reply with one word.")]), ], ) - return client.gateway.transport.post( + return client.proxy.transport.post( "/chat/completions", - headers=client.gateway.transport.bearer(key), + headers=client.proxy.transport.bearer(key), json=body, response_type=ChatResponse, ) @@ -120,11 +120,11 @@ class TestCacheControl: self, client: PassthroughClient, resources: ResourceManager ) -> None: model = f"e2e-bedrock-cache-{unique_marker()}" - model_id = client.gateway.create_model( + model_id = client.proxy.create_model( model, LiteLLMParamsBody(model=BEDROCK_MODEL, aws_region_name="us-east-1"), ) - resources.defer(lambda: client.gateway.delete_model(model_id)) + resources.defer(lambda: client.proxy.delete_model(model_id)) _assert_cache_read_on_second_call(client, resources.key(), model) @pytest.mark.covers( @@ -135,7 +135,7 @@ class TestCacheControl: self, client: PassthroughClient, resources: ResourceManager ) -> None: model = f"e2e-vertex-cache-{unique_marker()}" - model_id = client.gateway.create_model( + model_id = client.proxy.create_model( model, LiteLLMParamsBody( model=VERTEX_MODEL, @@ -144,5 +144,5 @@ class TestCacheControl: vertex_credentials=os.environ.get("VERTEXAI_CREDENTIALS"), ), ) - resources.defer(lambda: client.gateway.delete_model(model_id)) + resources.defer(lambda: client.proxy.delete_model(model_id)) _assert_cache_read_on_second_call(client, resources.key(), model) diff --git a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py index 5cc4ff308fa..f882bc5b4e4 100644 --- a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py +++ b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py @@ -43,7 +43,7 @@ class TestChatCompletionsRegression: self, client: PassthroughClient, scoped_key: str, model: str, route: str ) -> None: response = unwrap( - client.gateway.chat( + client.proxy.chat( scoped_key, ChatBody( model=model, diff --git a/tests/e2e/llm_translation/test_custom_pricing_e2e.py b/tests/e2e/llm_translation/test_custom_pricing_e2e.py index 7894b447be9..b4ff631a56b 100644 --- a/tests/e2e/llm_translation/test_custom_pricing_e2e.py +++ b/tests/e2e/llm_translation/test_custom_pricing_e2e.py @@ -23,7 +23,7 @@ import pytest from pydantic import BaseModel, RootModel from e2e_config import unique_marker -from e2e_gateway import Gateway +from proxy_client import ProxyClient from e2e_http import Success, unwrap from endpoints_client import EndpointsClient from lifecycle import ResourceManager @@ -116,14 +116,14 @@ def _model_info_entry(entries: list[ModelInfoEntry], model_name: str) -> ModelIn pytest.fail(f"{model_name} absent from /model/info; the override did not load") -def _poll_breakdown_row(gateway: Gateway, key: str, response_id: str | None) -> _SpendRow: +def _poll_breakdown_row(proxy: ProxyClient, key: str, response_id: str | None) -> _SpendRow: """Poll /spend/logs until the call's row lands with a cost breakdown (rows flush ~60s behind the call via proxy_batch_write_at).""" - deadline = time.monotonic() + gateway.poll_timeout + deadline = time.monotonic() + proxy.poll_timeout while time.monotonic() < deadline: - result = gateway.transport.get( + result = proxy.transport.get( "/spend/logs", - headers=gateway.transport.master, + headers=proxy.transport.master, params=SpendLogsParams(api_key=key), response_type=_SpendRows, ) @@ -144,7 +144,7 @@ def _poll_breakdown_row(gateway: Gateway, key: str, response_id: str | None) -> return row if priced and response_id is None: return priced[0] - time.sleep(gateway.poll_interval) + time.sleep(proxy.poll_interval) pytest.fail("no spend row with a cost breakdown landed before the deadline") @@ -158,7 +158,7 @@ class TestCustomPricing: model = _provision_custom_priced(endpoints_client, resources) chat = unwrap( - endpoints_client.gateway.chat( + endpoints_client.proxy.chat( scoped_key, ChatBody( model=model, @@ -172,7 +172,7 @@ class TestCustomPricing: ) ) - row = _poll_breakdown_row(endpoints_client.gateway, scoped_key, chat.id) + row = _poll_breakdown_row(endpoints_client.proxy, scoped_key, chat.id) assert row.metadata and row.metadata.cost_breakdown # guaranteed by the poll breakdown = row.metadata.cost_breakdown @@ -198,7 +198,7 @@ class TestCustomPricing: self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: model = _provision_custom_priced(endpoints_client, resources) - entry = _model_info_entry(endpoints_client.gateway.model_info(), model) + entry = _model_info_entry(endpoints_client.proxy.model_info(), model) assert entry.litellm_params.input_cost_per_token == CUSTOM_INPUT_RATE, ( f"/model/info litellm_params input rate " @@ -223,7 +223,7 @@ class TestCustomPricing: output_cost_per_token=None, ) - entries = {entry.model_name: entry for entry in endpoints_client.gateway.model_info()} + entries = {entry.model_name: entry for entry in endpoints_client.proxy.model_info()} custom_entry = entries.get(custom) sibling_entry = entries.get(sibling) assert custom_entry is not None, f"{custom} absent from /model/info" diff --git a/tests/e2e/llm_translation/test_deepseek_reasoning_e2e.py b/tests/e2e/llm_translation/test_deepseek_reasoning_e2e.py index 5adb8c24f9f..b06b241c0b5 100644 --- a/tests/e2e/llm_translation/test_deepseek_reasoning_e2e.py +++ b/tests/e2e/llm_translation/test_deepseek_reasoning_e2e.py @@ -33,11 +33,11 @@ PROMPT = "What is 17 + 26? Answer with just the number." def _register_reasoner(client: PassthroughClient, resources: ResourceManager) -> str: model = f"e2e-deepseek-reasoner-{unique_marker()}" - model_id = client.gateway.create_model( + model_id = client.proxy.create_model( model, LiteLLMParamsBody(model=REASONER, api_key="os.environ/DEEPSEEK_API_KEY"), ) - resources.defer(lambda: client.gateway.delete_model(model_id)) + resources.defer(lambda: client.proxy.delete_model(model_id)) return model @@ -56,7 +56,7 @@ class TestDeepSeekReasoningDisable: key = resources.key() response = unwrap( - client.gateway.chat( + client.proxy.chat( key, ChatBody( model=model, @@ -78,7 +78,7 @@ class TestDeepSeekReasoningDisable: key = resources.key() response = unwrap( - client.gateway.chat( + client.proxy.chat( key, ChatBody( model=model, @@ -100,7 +100,7 @@ class TestDeepSeekReasoningDisable: key = resources.key() response = unwrap( - client.gateway.chat( + client.proxy.chat( key, ChatBody( model=model, diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py index 1025d603fca..4b3191e60bb 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py @@ -76,9 +76,9 @@ def _system_reminder_turn() -> RichMessage: def _post_messages( client: EndpointsClient, key: str, body: RichMessagesRequest ) -> Result[MessagesResult]: - return client.gateway.transport.post( + return client.proxy.transport.post( "/v1/messages", - headers=client.gateway.transport.bearer(key), + headers=client.proxy.transport.bearer(key), json=body, response_type=MessagesResult, ) diff --git a/tests/e2e/llm_translation/test_ocr_rust_e2e.py b/tests/e2e/llm_translation/test_ocr_rust_e2e.py index e735d9c01b5..cdbf1883314 100644 --- a/tests/e2e/llm_translation/test_ocr_rust_e2e.py +++ b/tests/e2e/llm_translation/test_ocr_rust_e2e.py @@ -150,7 +150,7 @@ class TestRustOcrGateway: resources.defer(lambda: endpoints_client.delete_model(model_id)) key = resources.key() - response = unwrap(endpoints_client.gateway.ocr(key, OcrBody(model=model, document=case.document))) + response = unwrap(endpoints_client.proxy.ocr(key, OcrBody(model=model, document=case.document))) _assert_ocr_document(response) diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index 37d55c665b3..c8806faf3ea 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -35,7 +35,7 @@ def _fetch_cost_breakdown(client: PassthroughClient, result: StreamingResponse) whole point of passthrough spend tracking. """ assert result.call_id, "passthrough response had no x-litellm-call-id header" - rows = client.gateway.poll_logs_for_request_id( + rows = client.proxy.poll_logs_for_request_id( result.call_id, predicate=lambda rs: (rs[0].spend or 0) > 0, ) diff --git a/tests/e2e/llm_translation/test_provider_features_e2e.py b/tests/e2e/llm_translation/test_provider_features_e2e.py index 822a1d8d4d5..2ea1d28748d 100644 --- a/tests/e2e/llm_translation/test_provider_features_e2e.py +++ b/tests/e2e/llm_translation/test_provider_features_e2e.py @@ -37,17 +37,17 @@ class TestServiceTier: self, client: PassthroughClient, resources: ResourceManager ) -> None: model = f"e2e-service-tier-{unique_marker()}" - model_id = client.gateway.create_model( + model_id = client.proxy.create_model( model, LiteLLMParamsBody( model="openai/gpt-5.5", api_key="os.environ/OPENAI_API_KEY" ), ) - resources.defer(lambda: client.gateway.delete_model(model_id)) + resources.defer(lambda: client.proxy.delete_model(model_id)) key = resources.key() response = unwrap( - client.gateway.chat( + client.proxy.chat( key, ChatBody( model=model, diff --git a/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py b/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py index 78d2bb358d2..b6d90b7f6a2 100644 --- a/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py @@ -91,9 +91,9 @@ def _add_vertex_passthrough_model( client: PassthroughClient, model_name: str, project: str, credentials: str ) -> str: return unwrap( - client.gateway.transport.post( + client.proxy.transport.post( "/model/new", - headers=client.gateway.transport.master, + headers=client.proxy.transport.master, json=_ModelNewBody( model_name=model_name, litellm_params=_VertexDeploymentParams( @@ -111,9 +111,9 @@ def _add_vertex_passthrough_model( def _delete_model(client: PassthroughClient, model_id: str) -> None: - _ = client.gateway.transport.post( + _ = client.proxy.transport.post( "/model/delete", - headers=client.gateway.transport.master, + headers=client.proxy.transport.master, json=_ModelDeleteBody(id=model_id), response_type=NoBody, ) @@ -126,7 +126,7 @@ def _costed_row(client: PassthroughClient, call_id: str | None) -> SpendLogRow: a billed Vertex call that LiteLLM did not track is the exact regression #31689 guards against.""" assert call_id, "vertex passthrough response had no x-litellm-call-id header" - rows = client.gateway.poll_logs_for_request_id( + rows = client.proxy.poll_logs_for_request_id( call_id, predicate=lambda rs: (rs[0].spend or 0) > 0, ) diff --git a/tests/e2e/load/conftest.py b/tests/e2e/load/conftest.py new file mode 100644 index 00000000000..e2d135092fb --- /dev/null +++ b/tests/e2e/load/conftest.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from collections.abc import Iterator + +import pytest +from requests import RequestException + +from e2e_gateway import Gateway +from e2e_http import NoBody, Success +from load_client import LoadClient, build_client +from load_constants import LOAD_MODEL +from models import KeyGenerateBody, LiteLLMParamsBody, ModelsListResponse +from lifecycle import ResourceManager + +LOAD_MODEL_PARAMS = LiteLLMParamsBody( + model="openai/load-mock", + mock_response="This is a mock response for the throughput load test.", +) + + +@pytest.fixture(scope="session") +def client() -> LoadClient: + return build_client() + + +def _model_is_servable(gateway: Gateway, model_name: str) -> bool: + result = gateway.transport.get( + "/v1/models", + headers=gateway.transport.master, + params=NoBody(), + response_type=ModelsListResponse, + ) + return isinstance(result, Success) and any(entry.id == model_name for entry in result.data.data) + + +@pytest.fixture(scope="session", autouse=True) +def _ensure_load_model( # pyright: ignore[reportUnusedFunction] # pytest autouse session fixture, wired by name + client: LoadClient, +) -> Iterator[None]: + gateway = client.gateway + if _model_is_servable(gateway, LOAD_MODEL): + yield + return + + try: + model_id = gateway.create_model(LOAD_MODEL, LOAD_MODEL_PARAMS) + except (AssertionError, RequestException) as exc: + if _model_is_servable(gateway, LOAD_MODEL): + yield + return + raise AssertionError( + f"failed to register {LOAD_MODEL!r} for the throughput load test " + f"(not listed on the data plane and /model/new failed): {exc}" + ) from exc + + try: + yield + finally: + gateway.delete_model(model_id) + + +@pytest.fixture +def load_key(resources: ResourceManager, client: LoadClient) -> str: + key = client.gateway.generate_key(KeyGenerateBody(models=[LOAD_MODEL], user_id="e2e-load")) + resources.defer(lambda: client.gateway.delete_key(key)) + return key diff --git a/tests/e2e/load/load_client.py b/tests/e2e/load/load_client.py new file mode 100644 index 00000000000..df7c91fadf9 --- /dev/null +++ b/tests/e2e/load/load_client.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from e2e_gateway import Gateway, build_gateway + + +@dataclass(frozen=True, slots=True) +class LoadClient: + gateway: Gateway + + +def build_client() -> LoadClient: + return LoadClient(gateway=build_gateway()) diff --git a/tests/e2e/load/load_constants.py b/tests/e2e/load/load_constants.py new file mode 100644 index 00000000000..fd97f1398f4 --- /dev/null +++ b/tests/e2e/load/load_constants.py @@ -0,0 +1,3 @@ +from __future__ import annotations + +LOAD_MODEL = "load-mock" diff --git a/tests/e2e/load/locust_load.py b/tests/e2e/load/locust_load.py new file mode 100644 index 00000000000..990e24ee064 --- /dev/null +++ b/tests/e2e/load/locust_load.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import os +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + +from pydantic import BaseModel, TypeAdapter + +_LOCUSTFILE = Path(__file__).with_name("locustfile.py") + + +class _LocustStatEntry(BaseModel): + num_requests: int + num_failures: int + start_time: float + last_request_timestamp: float + + +_STATS_ADAPTER: TypeAdapter[list[_LocustStatEntry]] = TypeAdapter(list[_LocustStatEntry]) + + +@dataclass(frozen=True, slots=True) +class LoadResult: + requests: int + failures: int + requests_per_second: float + + @property + def failure_ratio(self) -> float: + return self.failures / self.requests if self.requests else 1.0 + + +def _aggregate(entries: list[_LocustStatEntry]) -> LoadResult: + requests = sum(entry.num_requests for entry in entries) + failures = sum(entry.num_failures for entry in entries) + if not entries or requests == 0: + return LoadResult(requests=requests, failures=failures, requests_per_second=0.0) + elapsed = max(entry.last_request_timestamp for entry in entries) - min(entry.start_time for entry in entries) + rps = requests / elapsed if elapsed > 0 else 0.0 + return LoadResult(requests=requests, failures=failures, requests_per_second=rps) + + +def run_chat_load( + *, + base_url: str, + api_key: str, + model: str, + users: int, + spawn_rate: float, + duration_seconds: float, +) -> LoadResult: + completed = subprocess.run( + [ + sys.executable, + "-m", + "locust", + "--headless", + "--json", + "--locustfile", + str(_LOCUSTFILE), + "--host", + base_url, + "--users", + str(users), + "--spawn-rate", + str(spawn_rate), + "--run-time", + f"{int(duration_seconds)}s", + "--exit-code-on-error", + "0", + ], + env={**os.environ, "LOAD_API_KEY": api_key, "LOAD_MODEL": model}, + capture_output=True, + text=True, + timeout=duration_seconds + 120, + check=False, + ) + if completed.returncode != 0: + raise RuntimeError( + f"locust exited {completed.returncode} before it could report throughput " + f"(a startup failure, not request failures, which are folded into the JSON summary via " + f"--exit-code-on-error 0):\n{completed.stderr}" + ) + try: + entries = _STATS_ADAPTER.validate_json(completed.stdout) + except ValueError as exc: + raise RuntimeError( + f"locust exited 0 but did not print a parseable --json throughput summary on stdout; " + f"got stdout={completed.stdout!r}, stderr={completed.stderr!r}" + ) from exc + return _aggregate(entries) diff --git a/tests/e2e/load/locustfile.py b/tests/e2e/load/locustfile.py new file mode 100644 index 00000000000..4aa7517ca0b --- /dev/null +++ b/tests/e2e/load/locustfile.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import os + +from locust import FastHttpUser, constant, task + +_MODEL = os.environ["LOAD_MODEL"] +_HEADERS = {"Authorization": f"Bearer {os.environ['LOAD_API_KEY']}"} +_PAYLOAD = { + "model": _MODEL, + "messages": [{"role": "user", "content": "load test ping"}], + "temperature": 0, + "max_tokens": 16, +} + + +class ChatUser(FastHttpUser): + wait_time = constant(0) + + @task + def chat(self) -> None: + self.client.post( # pyright: ignore[reportUnknownMemberType] # locust FastHttpSession.post types json/**kwargs as Any + "/chat/completions", + json=_PAYLOAD, + headers=_HEADERS, + name="/chat/completions", + ) diff --git a/tests/e2e/load/test_chat_completions_throughput_e2e.py b/tests/e2e/load/test_chat_completions_throughput_e2e.py new file mode 100644 index 00000000000..6dd5fff971a --- /dev/null +++ b/tests/e2e/load/test_chat_completions_throughput_e2e.py @@ -0,0 +1,42 @@ +import pytest + +from e2e_config import ( + LOAD_DURATION_SECONDS, + LOAD_MAX_FAILURE_RATIO, + LOAD_MIN_RPS, + LOAD_SPAWN_RATE, + LOAD_USERS, + PROXY_BASE_URL, +) +from load_client import LoadClient +from load_constants import LOAD_MODEL +from locust_load import run_chat_load + +pytestmark = [pytest.mark.e2e, pytest.mark.load] + + +class TestChatCompletionsThroughput: + @pytest.mark.covers("reliability.perf.throughput.under_slo") + def test_sustains_throughput_slo_under_load(self, client: LoadClient, load_key: str) -> None: + result = run_chat_load( + base_url=PROXY_BASE_URL, + api_key=load_key, + model=LOAD_MODEL, + users=LOAD_USERS, + spawn_rate=LOAD_SPAWN_RATE, + duration_seconds=LOAD_DURATION_SECONDS, + ) + + assert result.requests > 0, ( + f"no requests completed against {PROXY_BASE_URL} in {LOAD_DURATION_SECONDS}s; " + f"the load generator never drove traffic (proxy unreachable or model unservable)" + ) + assert result.failure_ratio <= LOAD_MAX_FAILURE_RATIO, ( + f"{result.failures}/{result.requests} requests failed " + f"({result.failure_ratio:.1%} > {LOAD_MAX_FAILURE_RATIO:.1%} allowed); " + f"throughput of {result.requests_per_second:.1f} RPS is not a clean read under this error rate" + ) + assert result.requests_per_second >= LOAD_MIN_RPS, ( + f"sustained {result.requests_per_second:.1f} RPS over {LOAD_DURATION_SECONDS}s with " + f"{LOAD_USERS} users, below the {LOAD_MIN_RPS} RPS SLO; the proxy request path regressed under load" + ) diff --git a/tests/e2e/logging/conftest.py b/tests/e2e/logging/conftest.py index 65be753154e..2285eb8d695 100644 --- a/tests/e2e/logging/conftest.py +++ b/tests/e2e/logging/conftest.py @@ -13,6 +13,7 @@ import pytest from logging_client import LangfuseCreds, LoggingClient, build_logging_client, load_langfuse_creds from datadog_reader import DdLogsReader, build_dd_logs_reader from otel_client import OtelReader, build_otel_reader +from proxy_client import ProxyClient def pytest_configure(config: pytest.Config) -> None: @@ -23,11 +24,11 @@ def pytest_configure(config: pytest.Config) -> None: @pytest.fixture(scope="session") -def client() -> LoggingClient: - """The logging suite's client: holds the shared Gateway so `resources` / +def client(proxy: ProxyClient) -> LoggingClient: + """The logging suite's client: holds the shared ProxyClient so `resources` / `scoped_key` clean up keys and teams, and adds `/metrics` scraping plus Langfuse read-back.""" - return build_logging_client() + return build_logging_client(proxy) @pytest.fixture(scope="session") diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index 8be573d72a9..7053bd1dfd2 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -1,7 +1,7 @@ """Client for the logging e2e suite: team/key/org-scoped Langfuse OTEL callbacks, chat (including tools), Prometheus scrape, and Langfuse observation read-back. -Holds the shared Gateway so the ``resources`` fixture cleans up keys, teams, +Holds the shared ProxyClient so the ``resources`` fixture cleans up keys, teams, users, orgs, and models it creates. External Langfuse reads go through ``e2e_http`` (the only module allowed to call ``requests.*``). @@ -24,7 +24,7 @@ import pytest from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError from e2e_config import POLL_INTERVAL, POLL_TIMEOUT -from e2e_gateway import Gateway, build_gateway +from proxy_client import ProxyClient from e2e_http import ( URL, AuthHeaders, @@ -262,7 +262,7 @@ def observation_has_guardrail(obs: LangfuseObservation, *, guardrail_name: str) @dataclass(frozen=True, slots=True) class LoggingClient: - gateway: Gateway + proxy: ProxyClient def key_with_alias( self, @@ -274,7 +274,7 @@ class LoggingClient: organization_id: str | None = None, metadata: KeyMetadata | None = None, ) -> str: - return self.gateway.generate_key( + return self.proxy.generate_key( KeyGenerateBody( key_alias=alias, models=models, @@ -286,7 +286,7 @@ class LoggingClient: ) def delete_key(self, key: str) -> None: - self.gateway.delete_key(key) + self.proxy.delete_key(key) def create_team( self, @@ -296,9 +296,9 @@ class LoggingClient: organization_id: str | None = None, ) -> str: return unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/team/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TeamNewBody( team_alias=alias, models=models, @@ -309,18 +309,18 @@ class LoggingClient: ).team_id def delete_team(self, team_id: str) -> None: - _ = self.gateway.transport.post( + _ = self.proxy.transport.post( "/team/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TeamDeleteBody(team_ids=[team_id]), response_type=NoBody, ) def create_user(self, *, user_email: str, user_id: str | None = None) -> str: return unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/user/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=UserNewBody( user_email=user_email, user_role="internal_user", @@ -331,27 +331,27 @@ class LoggingClient: ).user_id def delete_user(self, user_id: str) -> None: - _ = self.gateway.transport.post( + _ = self.proxy.transport.post( "/user/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=UserDeleteBody(user_ids=[user_id]), response_type=NoBody, ) def create_org(self, alias: str, *, models: list[str]) -> str: return unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/organization/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=OrgNewBody(organization_alias=alias, models=models), response_type=OrgNewResponse, ) ).organization_id def delete_org(self, organization_id: str) -> None: - _ = self.gateway.transport.delete( + _ = self.proxy.transport.delete( "/organization/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=OrgDeleteBody(organization_ids=[organization_id]), response_type=NoBody, ) @@ -364,9 +364,9 @@ class LoggingClient: callback_type: Literal["success", "failure", "success_and_failure"] = "success_and_failure", ) -> None: response = unwrap( - self.gateway.transport.post( + self.proxy.transport.post( f"/team/{team_id}/callback", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TeamCallbackBody( callback_name="langfuse_otel", callback_type=callback_type, @@ -382,9 +382,9 @@ class LoggingClient: def create_tool_permission_guardrail(self, name: str, *, allowed_tool: str) -> str: """Register a tool_permission guardrail that allows one tool and denies the rest.""" response = unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/guardrails", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=CreateGuardrailBody( guardrail=GuardrailSpec( guardrail_name=name, @@ -412,22 +412,22 @@ class LoggingClient: return guardrail_id def delete_guardrail(self, guardrail_id: str) -> None: - _ = self.gateway.transport.delete( + _ = self.proxy.transport.delete( f"/guardrails/{guardrail_id}", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=NoBody(), response_type=NoBody, ) def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str: - return self.gateway.create_model(model_name, litellm_params) + return self.proxy.create_model(model_name, litellm_params) def delete_model(self, model_id: str) -> None: - self.gateway.delete_model(model_id) + self.proxy.delete_model(model_id) def chat(self, key: str, model: str, text: str) -> ChatResponse: return unwrap( - self.gateway.chat( + self.proxy.chat( key, ChatBody( model=model, @@ -459,10 +459,10 @@ class LoggingClient: guardrails=guardrails, ) if stream: - return self.gateway.chat_stream(key, body) - return self.gateway.transport.send( + return self.proxy.chat_stream(key, body) + return self.proxy.transport.send( "/chat/completions", - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), json=body, ) @@ -479,11 +479,11 @@ class LoggingClient: stream=True if stream else None, ) if stream: - return self.gateway.transport.stream( - "/v1/messages", headers=self.gateway.transport.bearer(key), json=body + return self.proxy.transport.stream( + "/v1/messages", headers=self.proxy.transport.bearer(key), json=body ) - return self.gateway.transport.send( - "/v1/messages", headers=self.gateway.transport.bearer(key), json=body + return self.proxy.transport.send( + "/v1/messages", headers=self.proxy.transport.bearer(key), json=body ) def responses_raw( @@ -498,15 +498,15 @@ class LoggingClient: model=model, input=text, max_output_tokens=max_output_tokens, stream=True if stream else None ) if stream: - return self.gateway.transport.stream( - "/v1/responses", headers=self.gateway.transport.bearer(key), json=body + return self.proxy.transport.stream( + "/v1/responses", headers=self.proxy.transport.bearer(key), json=body ) - return self.gateway.transport.send( - "/v1/responses", headers=self.gateway.transport.bearer(key), json=body + return self.proxy.transport.send( + "/v1/responses", headers=self.proxy.transport.bearer(key), json=body ) def scrape_metrics(self) -> str: - return self.gateway.probe("/metrics", params=NoBody()).body + return self.proxy.probe("/metrics", params=NoBody()).body def poll_proxy_spend_for_key( self, @@ -529,7 +529,7 @@ class LoggingClient: return False return True - rows = self.gateway.poll_logs_for_key( + rows = self.proxy.poll_logs_for_key( key, min_rows=1, predicate=lambda rs: any(_matches(r) for r in rs) ) for row in rows: @@ -623,15 +623,15 @@ def first_ok(client: LoggingClient, send: Callable[[], StreamingResponse]) -> St the data plane's auth cache picks it up, so retry on 401 to a deadline; a 401 is rejected before the LLM call, so it cannot contaminate delivery or trace assertions. Any other failure is behavior under test and fails hard.""" - deadline = time.monotonic() + client.gateway.poll_timeout + deadline = time.monotonic() + client.proxy.poll_timeout while True: outcome = send() if outcome.ok: return outcome if outcome.status_code != 401 or time.monotonic() >= deadline: require_successful_call(outcome) - time.sleep(client.gateway.poll_interval) + time.sleep(client.proxy.poll_interval) -def build_logging_client() -> LoggingClient: - return LoggingClient(gateway=build_gateway()) +def build_logging_client(proxy: ProxyClient) -> LoggingClient: + return LoggingClient(proxy=proxy) diff --git a/tests/e2e/logging/test_datadog_log_e2e.py b/tests/e2e/logging/test_datadog_log_e2e.py index 48c111b6467..94811c6217e 100644 --- a/tests/e2e/logging/test_datadog_log_e2e.py +++ b/tests/e2e/logging/test_datadog_log_e2e.py @@ -52,7 +52,7 @@ def _assert_datadog_configured(client: LoggingClient) -> None: """Recorded state: the proxy reports the DataDog callback among its active callbacks, so a missing destination config fails here, before any delivery-based assertion can time out confusingly.""" - result = client.gateway.probe("/health/readiness/details", params=NoBody()) + result = client.proxy.probe("/health/readiness/details", params=NoBody()) assert result.status_code == 200, ( f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}" ) diff --git a/tests/e2e/logging/test_otel_trace_e2e.py b/tests/e2e/logging/test_otel_trace_e2e.py index 4db8813b1f9..b0a4988594e 100644 --- a/tests/e2e/logging/test_otel_trace_e2e.py +++ b/tests/e2e/logging/test_otel_trace_e2e.py @@ -48,7 +48,7 @@ def _assert_otel_destination_configured(client: LoggingClient) -> None: """Recorded state: the proxy reports the OTEL v2 logger among its active callbacks, so a missing/failed destination config fails here, before any traffic-based assertion can time out confusingly.""" - result = client.gateway.probe("/health/readiness/details", params=NoBody()) + result = client.proxy.probe("/health/readiness/details", params=NoBody()) assert result.status_code == 200, ( f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}" ) @@ -698,13 +698,13 @@ class TestOtelTraceCompleteness: key = client.key_with_alias(f"otel-err-{unique_marker()}", models=[model_name]) resources.defer(lambda: client.delete_key(key)) - deadline = time.monotonic() + client.gateway.poll_timeout + deadline = time.monotonic() + client.proxy.poll_timeout while True: outcome = client.chat_raw(key, model_name, "trigger an upstream auth failure", max_tokens=16) assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid" if "AnthropicException" in outcome.body or time.monotonic() >= deadline: break - time.sleep(client.gateway.poll_interval) + time.sleep(client.proxy.poll_interval) assert "AnthropicException" in outcome.body, ( "never saw the upstream provider failure before the deadline; the key may still be " f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}" diff --git a/tests/e2e/logging/test_prometheus_cardinality_e2e.py b/tests/e2e/logging/test_prometheus_cardinality_e2e.py index 163293a3009..44e3d93c07b 100644 --- a/tests/e2e/logging/test_prometheus_cardinality_e2e.py +++ b/tests/e2e/logging/test_prometheus_cardinality_e2e.py @@ -55,13 +55,13 @@ class TestPrometheusPerKeyCardinality: assert response.model, f"driver call for {alias} returned no model: {response}" wanted = frozenset(aliases) - deadline = time.monotonic() + client.gateway.poll_timeout + deadline = time.monotonic() + client.proxy.poll_timeout seen: frozenset[str] = frozenset() while time.monotonic() < deadline: seen = _aliases_in_metric(client.scrape_metrics(), REQUESTS_METRIC, ALIAS_LABEL) if wanted <= seen: break - time.sleep(client.gateway.poll_interval) + time.sleep(client.proxy.poll_interval) missing = wanted - seen assert not missing, ( diff --git a/tests/e2e/management/conftest.py b/tests/e2e/management/conftest.py index 4f2dc874a33..52c618f4fcd 100644 --- a/tests/e2e/management/conftest.py +++ b/tests/e2e/management/conftest.py @@ -14,6 +14,7 @@ import pytest from e2e_config import UI_BASE_URL, UI_PASSWORD, UI_USERNAME from management_client import ManagementClient, build_client +from proxy_client import ProxyClient if TYPE_CHECKING: from playwright.sync_api import Browser, Page @@ -27,8 +28,8 @@ def pytest_configure(config: pytest.Config) -> None: @pytest.fixture(scope="session") -def client() -> ManagementClient: - return build_client() +def client(proxy: ProxyClient) -> ManagementClient: + return build_client(proxy) @pytest.fixture(scope="session") diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 144a3c0c435..e967fb7b504 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -1,4 +1,4 @@ -"""Client for the management-routes e2e suite: the shared Gateway plus the +"""Client for the management-routes e2e suite: the shared ProxyClient plus the key/team/user/organization writes, the info/list read-backs the tests assert, and the raw-status calls judged by HTTP outcome (chat under a scoped key, an llm-only key hitting a management route). @@ -9,7 +9,7 @@ from __future__ import annotations import time from dataclasses import dataclass -from e2e_gateway import Gateway, build_gateway +from proxy_client import ProxyClient from e2e_http import NoBody, ProbeResult, Result, StreamingResponse, Success, UnknownApiError, unwrap from models import ( ChatBody, @@ -50,17 +50,17 @@ _TEAM_READY_SLEEP_SECONDS = 0.4 @dataclass(frozen=True, slots=True) class ManagementClient: - gateway: Gateway + proxy: ProxyClient def llm_only_key(self) -> str: - return self.gateway.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"])) + return self.proxy.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"])) def update_key_models(self, key: str, models: list[str]) -> None: last: Result[NoBody] | None = None for attempt in range(5): - last = self.gateway.transport.post( + last = self.proxy.transport.post( "/key/update", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=KeyUpdateBody(key=key, models=models), response_type=NoBody, ) @@ -79,11 +79,11 @@ class ManagementClient: def delete_key_strict(self, key: str) -> None: """Strict delete for the act phase of a test: a failed delete is a hard - failure, unlike the warn-only Gateway.delete_key used at teardown.""" + failure, unlike the warn-only ProxyClient.delete_key used at teardown.""" _ = unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/key/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=KeyDeleteBody(keys=[key]), response_type=NoBody, ) @@ -91,9 +91,9 @@ class ManagementClient: def key_alias_count(self, key_alias: str) -> int: return unwrap( - self.gateway.transport.get( + self.proxy.transport.get( "/key/list", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=KeyListParams(key_alias=key_alias), response_type=KeyListResponse, ) @@ -101,9 +101,9 @@ class ManagementClient: def create_team(self, body: TeamNewBody) -> str: team_id = unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/team/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=body, response_type=TeamNewResponse, ) @@ -112,32 +112,32 @@ class ManagementClient: return team_id def delete_team(self, team_id: str) -> None: - _ = self.gateway.transport.post( + _ = self.proxy.transport.post( "/team/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TeamDeleteBody(team_ids=[team_id]), response_type=NoBody, ) def team_info(self, team_id: str) -> TeamData: return unwrap( - self.gateway.transport.get( + self.proxy.transport.get( "/team/info", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=TeamInfoParams(team_id=team_id), response_type=TeamInfoResponse, ) ).team_info def team_info_status(self, team_id: str) -> ProbeResult: - return self.gateway.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id)) + return self.proxy.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id)) def _wait_for_team(self, team_id: str) -> None: last: Result[TeamInfoResponse] | None = None for _ in range(_TEAM_READY_ATTEMPTS): - last = self.gateway.transport.get( + last = self.proxy.transport.get( "/team/info", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=TeamInfoParams(team_id=team_id), response_type=TeamInfoResponse, ) @@ -152,9 +152,9 @@ class ManagementClient: def add_team_member(self, team_id: str, user_id: str) -> None: last: Result[NoBody] | None = None for attempt in range(_TEAM_READY_ATTEMPTS): - last = self.gateway.transport.post( + last = self.proxy.transport.post( "/team/member_add", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TeamMemberAddBody(team_id=team_id, member=TeamMemberEntry(role="user", user_id=user_id)), response_type=NoBody, ) @@ -173,9 +173,9 @@ class ManagementClient: def delete_team_member(self, team_id: str, user_id: str) -> None: _ = unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/team/member_delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TeamMemberDeleteBody(team_id=team_id, user_id=user_id), response_type=NoBody, ) @@ -183,27 +183,27 @@ class ManagementClient: def create_user(self, body: UserNewBody) -> str: return unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/user/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=body, response_type=UserNewResponse, ) ).user_id def delete_user(self, user_id: str) -> None: - _ = self.gateway.transport.post( + _ = self.proxy.transport.post( "/user/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=UserDeleteBody(user_ids=[user_id]), response_type=NoBody, ) def user_info(self, user_id: str) -> UserInfoResponse: return unwrap( - self.gateway.transport.get( + self.proxy.transport.get( "/user/info", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=UserInfoParams(user_id=user_id), response_type=UserInfoResponse, ) @@ -211,9 +211,9 @@ class ManagementClient: def user_count(self, user_id: str) -> int: return unwrap( - self.gateway.transport.get( + self.proxy.transport.get( "/user/list", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=UserListParams(user_ids=user_id), response_type=UserListResponse, ) @@ -221,48 +221,48 @@ class ManagementClient: def create_org(self, body: OrgNewBody) -> str: return unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/organization/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=body, response_type=OrgNewResponse, ) ).organization_id def delete_org(self, organization_id: str) -> None: - _ = self.gateway.transport.delete( + _ = self.proxy.transport.delete( "/organization/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=OrgDeleteBody(organization_ids=[organization_id]), response_type=NoBody, ) def org_info(self, organization_id: str) -> OrgInfoResponse: return unwrap( - self.gateway.transport.get( + self.proxy.transport.get( "/organization/info", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=OrgInfoParams(organization_id=organization_id), response_type=OrgInfoResponse, ) ) def chat_status(self, key: str, model: str, content: str) -> StreamingResponse: - return self.gateway.transport.send( + return self.proxy.transport.send( "/chat/completions", - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), json=ChatBody(model=model, messages=[ChatMessage(role="user", content=content)], max_tokens=16), ) def key_generate_status(self, key: str, body: KeyGenerateBody) -> StreamingResponse: - return self.gateway.transport.send("/key/generate", headers=self.gateway.transport.bearer(key), json=body) + return self.proxy.transport.send("/key/generate", headers=self.proxy.transport.bearer(key), json=body) def team_new_status(self, key: str, body: TeamNewBody) -> StreamingResponse: - return self.gateway.transport.send("/team/new", headers=self.gateway.transport.bearer(key), json=body) + return self.proxy.transport.send("/team/new", headers=self.proxy.transport.bearer(key), json=body) def user_new_status(self, key: str, body: UserNewBody) -> StreamingResponse: - return self.gateway.transport.send("/user/new", headers=self.gateway.transport.bearer(key), json=body) + return self.proxy.transport.send("/user/new", headers=self.proxy.transport.bearer(key), json=body) -def build_client() -> ManagementClient: - return ManagementClient(gateway=build_gateway()) +def build_client(proxy: ProxyClient) -> ManagementClient: + return ManagementClient(proxy=proxy) diff --git a/tests/e2e/management/test_key_models_dropdown_e2e.py b/tests/e2e/management/test_key_models_dropdown_e2e.py index 78e3f9e1a7b..20bd2191ac4 100644 --- a/tests/e2e/management/test_key_models_dropdown_e2e.py +++ b/tests/e2e/management/test_key_models_dropdown_e2e.py @@ -104,8 +104,8 @@ def _provision_team(client: ManagementClient, resources: ResourceManager, alias: def _provision_key( client: ManagementClient, resources: ResourceManager, alias: str, team_id: str | None = None ) -> str: - key = client.gateway.generate_key(KeyGenerateBody(key_alias=alias, models=["gpt-5.5"], team_id=team_id)) - resources.defer(lambda: client.gateway.delete_key(key)) + key = client.proxy.generate_key(KeyGenerateBody(key_alias=alias, models=["gpt-5.5"], team_id=team_id)) + resources.defer(lambda: client.proxy.delete_key(key)) return key @@ -122,9 +122,9 @@ class TestKeyModelsDropdownUI: assert "All Team Models" not in options, f"teamless create offered 'All Team Models': {options}" key = _submit_create_modal(ui_page, sentinel_label="All Proxy Models") - resources.defer(lambda: client.gateway.delete_key(key)) + resources.defer(lambda: client.proxy.delete_key(key)) - info = client.gateway.key_info(key) + info = client.proxy.key_info(key) assert info.models == ["all-proxy-models"], f"persisted models {info.models}" assert info.team_id is None, f"teamless key persisted with team {info.team_id}" @@ -144,9 +144,9 @@ class TestKeyModelsDropdownUI: assert "all-proxy-models" not in options, f"team key create offered the raw sentinel: {options}" key = _submit_create_modal(ui_page, sentinel_label="All Team Models") - resources.defer(lambda: client.gateway.delete_key(key)) + resources.defer(lambda: client.proxy.delete_key(key)) - info = client.gateway.key_info(key) + info = client.proxy.key_info(key) assert info.models == ["all-team-models"], f"persisted models {info.models}" assert info.team_id == team_id, f"persisted team {info.team_id}, expected {team_id}" diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index 7a05a5c520f..adbf3e8b065 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -27,18 +27,18 @@ from models import KeyGenerateBody, OrgNewBody, TeamNewBody, UserNewBody pytestmark = pytest.mark.e2e def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T: - deadline = time.monotonic() + client.gateway.poll_timeout + deadline = time.monotonic() + client.proxy.poll_timeout while time.monotonic() < deadline: found = attempt() if found is not None: return found - time.sleep(client.gateway.poll_interval) + time.sleep(client.proxy.poll_interval) pytest.fail(failure) def _generate_key(client: ManagementClient, resources: ResourceManager, body: KeyGenerateBody) -> str: - key = client.gateway.generate_key(body) - resources.defer(lambda: client.gateway.delete_key(key)) + key = client.proxy.generate_key(body) + resources.defer(lambda: client.proxy.delete_key(key)) return key @@ -114,7 +114,7 @@ class TestKeyRoutes: KeyGenerateBody(models=["gemini-2.5-flash"], key_alias=alias, tpm_limit=424242, rpm_limit=424243), ) - info = client.gateway.key_info(key) + info = client.proxy.key_info(key) assert info.key_alias == alias, f"/key/info reports key_alias {info.key_alias!r}, configured {alias!r}" assert info.models == ["gemini-2.5-flash"], ( f"/key/info reports models {info.models}, configured ['gemini-2.5-flash']" @@ -143,7 +143,7 @@ class TestKeyRoutes: client.update_key_models(key, ["gpt-5.5"]) - info = client.gateway.key_info(key) + info = client.proxy.key_info(key) assert info.models == ["gpt-5.5"], ( f"/key/info reports models {info.models} after /key/update to ['gpt-5.5']" ) @@ -184,7 +184,7 @@ class TestTeamRoutes: ) key = _generate_key(client, resources, KeyGenerateBody(team_id=team_id)) - key_info = client.gateway.key_info(key) + key_info = client.proxy.key_info(key) assert key_info.team_id == team_id, ( f"key generated under team {team_id} carries team_id {key_info.team_id!r} in /key/info" ) @@ -260,7 +260,7 @@ class TestManagementRoutePermissions: self, client: ManagementClient, resources: ResourceManager ) -> None: key = client.llm_only_key() - resources.defer(lambda: client.gateway.delete_key(key)) + resources.defer(lambda: client.proxy.delete_key(key)) marker = unique_marker() alias = f"e2e-mgmt-forbidden-key-{marker}" team_id = f"e2e-mgmt-forbidden-team-{marker}" diff --git a/tests/e2e/mcp/conftest.py b/tests/e2e/mcp/conftest.py index 77fef574706..3f970f3c008 100644 --- a/tests/e2e/mcp/conftest.py +++ b/tests/e2e/mcp/conftest.py @@ -2,15 +2,16 @@ The shared lifecycle (resources/scoped_key), proxy liveness handling, and the `e2e`/`covers` markers live in the parent tests/e2e/conftest.py. McpClient holds -the shared Gateway, so the `resources` fixture tears down whatever this suite -creates (keys via the Gateway, MCP servers via the deferred cleanups). +the shared ProxyClient, so the `resources` fixture tears down whatever this suite +creates (keys via the ProxyClient, MCP servers via the deferred cleanups). """ import pytest from mcp_client import McpClient, build_client +from proxy_client import ProxyClient @pytest.fixture(scope="session") -def client() -> McpClient: - return build_client() +def client(proxy: ProxyClient) -> McpClient: + return build_client(proxy) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index a1dac3fdac4..59358305ee7 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -15,9 +15,9 @@ from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field, RootModel -from e2e_gateway import Gateway, build_gateway from e2e_http import Headers, NoBody, Result, unwrap from models import KeyGenerateBody, ObjectPermission +from proxy_client import ProxyClient class ApiKeyHeaders(Headers): @@ -92,31 +92,31 @@ class McpCallToolResponse(BaseModel): @dataclass(frozen=True, slots=True) class McpClient: - gateway: Gateway + proxy: ProxyClient def register_server(self, *, server_name: str, alias: str, url: str) -> str: return unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/v1/mcp/server", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=McpServerNewBody(server_name=server_name, alias=alias, url=url), response_type=McpServerNewResponse, ) ).server_id def delete_server(self, server_id: str) -> None: - _ = self.gateway.transport.delete( + _ = self.proxy.transport.delete( f"/v1/mcp/server/{server_id}", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=NoBody(), response_type=NoBody, ) def registered_servers(self) -> list[McpServerRow]: return unwrap( - self.gateway.transport.get( + self.proxy.transport.get( "/v1/mcp/server", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=NoBody(), response_type=McpServersListResponse, ) @@ -126,12 +126,12 @@ class McpClient: object_permission = ( ObjectPermission(mcp_servers=mcp_servers) if mcp_servers is not None else None ) - return self.gateway.generate_key( + return self.proxy.generate_key( KeyGenerateBody(models=[], user_id=user_id, object_permission=object_permission) ) def list_tools(self, key: str) -> Result[McpToolsListResponse]: - return self.gateway.transport.get( + return self.proxy.transport.get( "/mcp-rest/tools/list", headers=ApiKeyHeaders(x_litellm_api_key=key), params=NoBody(), @@ -141,7 +141,7 @@ class McpClient: def call_tool( self, key: str, *, server_id: str, name: str, arguments: dict[str, int] ) -> Result[McpCallToolResponse]: - return self.gateway.transport.post( + return self.proxy.transport.post( "/mcp-rest/tools/call", headers=ApiKeyHeaders(x_litellm_api_key=key), json=McpCallToolBody(name=name, arguments=arguments, server_id=server_id), @@ -149,5 +149,5 @@ class McpClient: ) -def build_client() -> McpClient: - return McpClient(gateway=build_gateway()) +def build_client(proxy: ProxyClient) -> McpClient: + return McpClient(proxy=proxy) diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index eaa49af5b69..ee316b44e68 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -40,7 +40,7 @@ def _register_math_server(client: McpClient, resources: ResourceManager) -> str: def _key(client: McpClient, resources: ResourceManager, *, mcp_servers: list[str] | None) -> str: label = "allowed" if mcp_servers else "denied" key = client.generate_key(user_id=f"e2e-mcp-{label}-{unique_marker()}", mcp_servers=mcp_servers) - resources.defer(lambda: client.gateway.delete_key(key)) + resources.defer(lambda: client.proxy.delete_key(key)) return key diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 39832d1a17f..daf85b7fc74 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -268,7 +268,7 @@ class SpendLogsParams(BaseModel): raise ValueError( "unfiltered /spend/logs returns the entire spend table and OOMs the " "runner on long-lived environments; filter by request_id or api_key, " - "or use Gateway.spend_logs_window for a bounded /spend/logs/v2 read" + "or use ProxyClient.spend_logs_window for a bounded /spend/logs/v2 read" ) return self @@ -448,6 +448,7 @@ class LiteLLMParamsBody(BaseModel): extra_headers: dict[str, str] | None = None use_in_pass_through: bool | None = None complexity_router_config: dict[str, object] | None = None + mock_response: str | None = None ModelMode = Literal["batch", "realtime", "image_generation"] diff --git a/tests/e2e/e2e_gateway.py b/tests/e2e/proxy_client.py similarity index 96% rename from tests/e2e/e2e_gateway.py rename to tests/e2e/proxy_client.py index ad8b2e833a8..0dbb80990c0 100644 --- a/tests/e2e/e2e_gateway.py +++ b/tests/e2e/proxy_client.py @@ -1,8 +1,8 @@ -"""Gateway: the shared proxy operations, DI'd into every client (composition). +"""ProxyClient: the shared proxy operations, DI'd into every client (composition). A frozen-slots dataclass holding a Transport plus poll config. Clients hold a -Gateway and add their own route methods; the lifecycle ResourceManager uses the -Gateway's key/customer methods for cleanup. Read-backs are eventually consistent +ProxyClient and add their own route methods; the lifecycle ResourceManager uses the +ProxyClient's key/customer methods for cleanup. Read-backs are eventually consistent (proxy_batch_write_at ~60s) so they poll to a deadline. """ @@ -69,7 +69,7 @@ RowsPredicate = Callable[[list[SpendLogRow]], bool] @dataclass(frozen=True, slots=True) -class Gateway: +class ProxyClient: transport: Transport poll_timeout: float = 120.0 poll_interval: float = 5.0 @@ -319,13 +319,13 @@ class Gateway: return self.transport.probe(path, params=params) -def build_gateway( +def build_proxy_client( *, base_url: str = PROXY_BASE_URL, master_key: str = MASTER_KEY, control_plane_base_url: str = CONTROL_PLANE_BASE_URL, -) -> Gateway: - """The Gateway every suite's client is built from: a SplitTransport that routes +) -> ProxyClient: + """The ProxyClient every suite's client is built from: a SplitTransport that routes LLM calls to the data plane (PROXY_BASE_URL) and management/admin calls to the control plane (CONTROL_PLANE_BASE_URL), with the shared poll budget. The two base URLs are the same for a monolithic proxy, so routing is then a no-op. @@ -334,7 +334,7 @@ def build_gateway( way than ``e2e_config``'s env names (see ``claude_code/_env.py``); they must pass all three together, since a caller that overrides only the data plane would leave management calls pointed at the env default.""" - return Gateway( + return ProxyClient( transport=SplitTransport( data=HttpTransport( base_url=base_url, diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index d3b2193226b..e9611df139b 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -5,3 +5,4 @@ addopts = --strict-markers --strict-config markers = e2e: live test that requires a running proxy and real provider keys + load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites diff --git a/tests/e2e/quota_management/budgets/budget_client.py b/tests/e2e/quota_management/budgets/budget_client.py index c2f5dcdd57c..53da8f68b1b 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -1,4 +1,4 @@ -"""Client for budget e2e tests: the shared Gateway plus budget-bearing entity +"""Client for budget e2e tests: the shared ProxyClient plus budget-bearing entity management (user / team / team-member / org / customer / tag / budget-table) and info reads. @@ -15,7 +15,7 @@ from dataclasses import dataclass from pydantic import AliasPath, BaseModel, Field, RootModel -from e2e_gateway import Gateway, build_gateway +from proxy_client import ProxyClient from e2e_http import NoBody, Result, StreamingResponse, Success, unwrap from models import ( AnthropicMessagesBody, @@ -185,9 +185,9 @@ def model_budget(model: str, limit: float, period: str = "30d") -> dict[str, Mod @dataclass(frozen=True, slots=True) class BudgetClient: - gateway: Gateway + proxy: ProxyClient - # ---- generic key ops (delegate to the shared Gateway) --------------- + # ---- generic key ops (delegate to the shared ProxyClient) --------------- def generate_key( self, @@ -203,7 +203,7 @@ class BudgetClient: budget_fallbacks: dict[str, list[str]] | None = None, budget_limits: list[BudgetWindow] | None = None, ) -> str: - return self.gateway.generate_key( + return self.proxy.generate_key( KeyGenerateBody( models=models or [], max_budget=max_budget, @@ -219,10 +219,10 @@ class BudgetClient: ) def delete_key(self, key: str) -> None: - self.gateway.delete_key(key) + self.proxy.delete_key(key) def delete_customers(self, user_ids: list[str]) -> None: - self.gateway.delete_customers(user_ids) + self.proxy.delete_customers(user_ids) # ---- chat (raw HTTP outcome: a budget block surfaces as a non-2xx) -- @@ -236,9 +236,9 @@ class BudgetClient: user: str | None = None, tags: list[str] | None = None, ) -> StreamingResponse: - return self.gateway.transport.send( + return self.proxy.transport.send( "/chat/completions", - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), json=ChatBody( model=model, messages=[ChatMessage(role="user", content=content)], @@ -256,9 +256,9 @@ class BudgetClient: *, max_tokens: int = 16, ) -> StreamingResponse: - return self.gateway.transport.send( + return self.proxy.transport.send( "/v1/messages", - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), json=AnthropicMessagesBody( model=model, messages=[ChatMessage(role="user", content=content)], @@ -270,26 +270,26 @@ class BudgetClient: def create_user(self, *, max_budget: float, budget_duration: str | None = None) -> str: return unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/user/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=UserNewBody(max_budget=max_budget, budget_duration=budget_duration), response_type=UserNewResponse, ) ).user_id def delete_user(self, user_id: str) -> None: - _ = self.gateway.transport.post( + _ = self.proxy.transport.post( "/user/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=UserDeleteBody(user_ids=[user_id]), response_type=NoBody, ) def user_info(self, user_id: str) -> UserInfoRow | None: - result = self.gateway.transport.get( + result = self.proxy.transport.get( "/user/info", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=UserInfoParams(user_id=user_id), response_type=UserInfoResponse, ) @@ -302,9 +302,9 @@ class BudgetClient: # ---- customer / end-user ------------------------------------------- def create_customer(self, customer_id: str, *, max_budget: float) -> str: - resp = self.gateway.transport.send( + resp = self.proxy.transport.send( "/customer/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=CustomerNewBody(user_id=customer_id, max_budget=max_budget), ) assert resp.ok, resp.body @@ -314,9 +314,9 @@ class BudgetClient: def create_org(self, *, max_budget: float, alias: str, budget_duration: str | None = None) -> str: return unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/organization/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=OrgNewBody( organization_alias=alias, max_budget=max_budget, @@ -330,9 +330,9 @@ class BudgetClient: """The id of the budget row backing an org; its budget_reset_at is read via budget_info (LIT-4570: /organization/new stores budget_duration without scheduling budget_reset_at, so the reset job's first tick schedules it).""" - result = self.gateway.transport.get( + result = self.proxy.transport.get( "/organization/info", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=OrgInfoParams(organization_id=org_id), response_type=OrgInfoResponse, ) @@ -343,9 +343,9 @@ class BudgetClient: return None def delete_org(self, org_id: str) -> None: - _ = self.gateway.transport.delete( + _ = self.proxy.transport.delete( "/organization/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=OrgDeleteBody(organization_ids=[org_id]), response_type=NoBody, ) @@ -362,9 +362,9 @@ class BudgetClient: budget_limits: list[BudgetWindow] | None = None, ) -> str: team_id = unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/team/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TeamNewBody( team_alias=alias, max_budget=max_budget, @@ -379,9 +379,9 @@ class BudgetClient: return team_id def delete_team(self, team_id: str) -> None: - _ = self.gateway.transport.post( + _ = self.proxy.transport.post( "/team/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TeamDeleteBody(team_ids=[team_id]), response_type=NoBody, ) @@ -389,9 +389,9 @@ class BudgetClient: def _wait_for_team(self, team_id: str) -> None: last: Result[TeamInfoResponse] | None = None for _ in range(_TEAM_READY_ATTEMPTS): - last = self.gateway.transport.get( + last = self.proxy.transport.get( "/team/info", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=TeamInfoParams(team_id=team_id), response_type=TeamInfoResponse, ) @@ -406,9 +406,9 @@ class BudgetClient: def add_team_member(self, team_id: str, user_id: str, *, max_budget_in_team: float | None = None) -> None: last_body = "" for attempt in range(_TEAM_READY_ATTEMPTS): - resp = self.gateway.transport.send( + resp = self.proxy.transport.send( "/team/member_add", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TeamMemberAddBody( team_id=team_id, member=TeamMember(role="user", user_id=user_id), @@ -432,9 +432,9 @@ class BudgetClient: max_budget_in_team: float | None = None, budget_duration: str | None = None, ) -> None: - resp = self.gateway.transport.send( + resp = self.proxy.transport.send( "/team/member_update", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TeamMemberUpdateBody( team_id=team_id, user_id=user_id, @@ -448,9 +448,9 @@ class BudgetClient: """The member's per-team budget_reset_at as /team/info reports it, or None if no reset is scheduled. The reset job advances this each time the window elapses; a job that skips the row leaves it pinned forever.""" - result = self.gateway.transport.get( + result = self.proxy.transport.get( "/team/info", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=TeamInfoParams(team_id=team_id), response_type=TeamInfoResponse, ) @@ -466,18 +466,18 @@ class BudgetClient: # ---- tag ------------------------------------------------------------ def create_tag(self, name: str, *, max_budget: float) -> str: - resp = self.gateway.transport.send( + resp = self.proxy.transport.send( "/tag/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TagNewBody(name=name, max_budget=max_budget), ) assert resp.ok, resp.body return name def delete_tag(self, name: str) -> None: - _ = self.gateway.transport.post( + _ = self.proxy.transport.post( "/tag/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=TagDeleteBody(name=name), response_type=NoBody, ) @@ -492,9 +492,9 @@ class BudgetClient: budget_duration: str | None = None, ) -> str: return unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/budget/new", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=BudgetNewBody( max_budget=max_budget, soft_budget=soft_budget, @@ -505,17 +505,17 @@ class BudgetClient: ).budget_id def delete_budget(self, budget_id: str) -> None: - _ = self.gateway.transport.post( + _ = self.proxy.transport.post( "/budget/delete", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=BudgetDeleteBody(id=budget_id), response_type=NoBody, ) def budget_info(self, budget_id: str) -> tuple[BudgetRow, ...]: - result = self.gateway.transport.post( + result = self.proxy.transport.post( "/budget/info", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=BudgetInfoBody(budgets=[budget_id]), response_type=BudgetInfoResponse, ) @@ -526,5 +526,5 @@ class BudgetClient: return () -def build_client() -> BudgetClient: - return BudgetClient(gateway=build_gateway()) +def build_client(proxy: ProxyClient) -> BudgetClient: + return BudgetClient(proxy=proxy) diff --git a/tests/e2e/quota_management/budgets/conftest.py b/tests/e2e/quota_management/budgets/conftest.py index 4299d2ffd49..639761cdfc8 100644 --- a/tests/e2e/quota_management/budgets/conftest.py +++ b/tests/e2e/quota_management/budgets/conftest.py @@ -1,7 +1,7 @@ """Budgets suite's `client` fixture. The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker -live in the parent tests/e2e/conftest.py. BudgetClient holds the shared Gateway, +live in the parent tests/e2e/conftest.py. BudgetClient holds the shared ProxyClient, so the `resources` fixture cleans up keys through it; tests register entity deletes via `resources.defer(...)`. """ @@ -9,8 +9,9 @@ via `resources.defer(...)`. import pytest from budget_client import BudgetClient, build_client +from proxy_client import ProxyClient @pytest.fixture(scope="session") -def client() -> BudgetClient: - return build_client() +def client(proxy: ProxyClient) -> BudgetClient: + return build_client(proxy) diff --git a/tests/e2e/quota_management/budgets/test_budget_crud_e2e.py b/tests/e2e/quota_management/budgets/test_budget_crud_e2e.py index 47473326f32..5070ec89704 100644 --- a/tests/e2e/quota_management/budgets/test_budget_crud_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_crud_e2e.py @@ -30,7 +30,7 @@ def test_budget_crud_roundtrip(client: BudgetClient, resources: ResourceManager) # Attach the budget to a key and confirm the key reflects it. key = client.generate_key(budget_id=budget_id) resources.defer(lambda: client.delete_key(key)) - info = client.gateway.key_info(key) + info = client.proxy.key_info(key) linked = info.litellm_budget_table assert info.budget_id == budget_id or (linked is not None and linked.max_budget == 12.5), ( f"key does not reflect attached budget: {info.budget_id}, {linked}" @@ -49,7 +49,7 @@ def test_budget_duration_schedules_reset_on_key(client: BudgetClient, resources: key = client.generate_key(max_budget=10.0, budget_duration="30d") resources.defer(lambda: client.delete_key(key)) - reset_at = client.gateway.key_info(key).budget_reset_at + reset_at = client.proxy.key_info(key).budget_reset_at assert reset_at, "budget_duration did not set budget_reset_at on the key" # budget_duration schedules a FUTURE reset. Don't assume now+30d exactly: the diff --git a/tests/e2e/quota_management/budgets/test_budget_fallback_e2e.py b/tests/e2e/quota_management/budgets/test_budget_fallback_e2e.py index 56197dcfee2..fe6db8f0454 100644 --- a/tests/e2e/quota_management/budgets/test_budget_fallback_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_fallback_e2e.py @@ -53,7 +53,7 @@ def test_budget_fallback_reroutes_anthropic_messages_to_openai( # The rerouted call must be recorded under the fallback model, not the # exhausted primary - proving spend tracking followed the reroute. - rows = client.gateway.poll_logs_for_key( + rows = client.proxy.poll_logs_for_key( key, predicate=lambda rows: any(FALLBACK_MODEL in (r.model or "") for r in rows) ) assert any(FALLBACK_MODEL in (r.model or "") for r in rows), ( diff --git a/tests/e2e/quota_management/budgets/test_budget_reset_advances_e2e.py b/tests/e2e/quota_management/budgets/test_budget_reset_advances_e2e.py index bc634f70b13..88cb1a045de 100644 --- a/tests/e2e/quota_management/budgets/test_budget_reset_advances_e2e.py +++ b/tests/e2e/quota_management/budgets/test_budget_reset_advances_e2e.py @@ -62,7 +62,7 @@ def test_key_with_budget_duration_schedules_reset_at_creation( key = client.generate_key(max_budget=TINY_CAP, budget_duration=f"{WINDOW_SECONDS}s") resources.defer(lambda: client.delete_key(key)) - info = client.gateway.key_info(key) + info = client.proxy.key_info(key) assert info.budget_reset_at is not None, "budget_duration set no budget_reset_at" assert _as_datetime(info.budget_reset_at) > _as_datetime("1970-01-01T00:00:00Z") @@ -99,7 +99,7 @@ def test_key_budget_reset_at_advances_after_window( key = client.generate_key(max_budget=TINY_CAP, budget_duration=f"{WINDOW_SECONDS}s") resources.defer(lambda: client.delete_key(key)) - before_raw = client.gateway.key_info(key).budget_reset_at + before_raw = client.proxy.key_info(key).budget_reset_at assert before_raw is not None, "no budget_reset_at scheduled at creation" before = _as_datetime(before_raw) @@ -112,7 +112,7 @@ def test_key_budget_reset_at_advances_after_window( if not result.ok: assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}" continue - info = client.gateway.key_info(key) + info = client.proxy.key_info(key) assert info.budget_reset_at is not None, "budget_reset_at cleared by reset" assert _as_datetime(info.budget_reset_at) > before, ( "budget_reset_at did not advance past the pre-reset value" @@ -145,7 +145,7 @@ def test_multi_window_key_resets_each_window_independently( start = time.monotonic() _drive_to_block(client, key) - spend_at_block = client.gateway.key_info(key).spend or 0.0 + spend_at_block = client.proxy.key_info(key).spend or 0.0 deadline = time.monotonic() + RESET_DEADLINE_SECONDS while time.monotonic() < deadline: @@ -156,7 +156,7 @@ def test_multi_window_key_resets_each_window_independently( assert elapsed < WINDOW_SECONDS + 90, ( f"tight window reset took {elapsed:.0f}s - too long for {WINDOW_SECONDS}s" ) - assert (client.gateway.key_info(key).spend or 0.0) >= spend_at_block, ( + assert (client.proxy.key_info(key).spend or 0.0) >= spend_at_block, ( "roomy window spend was wiped when only the tight window should reset" ) return diff --git a/tests/e2e/quota_management/budgets/test_spend_counter_reseed_e2e.py b/tests/e2e/quota_management/budgets/test_spend_counter_reseed_e2e.py index 52fe112ed29..d9f9830ee20 100644 --- a/tests/e2e/quota_management/budgets/test_spend_counter_reseed_e2e.py +++ b/tests/e2e/quota_management/budgets/test_spend_counter_reseed_e2e.py @@ -141,7 +141,7 @@ def test_cold_counter_reseed_keeps_counter_equal_to_db_spend( "the spend counter never went cold; default_redis_ttl must be short enough for it " "to expire, otherwise the burst reads a warm counter and the reseed is never exercised" ) - db_spend = client.gateway.key_info(key).spend or 0.0 + db_spend = client.proxy.key_info(key).spend or 0.0 assert db_spend > 0, f"no DB spend accumulated from real calls: {db_spend}" burst_results = [] diff --git a/tests/e2e/quota_management/budgets/test_team_member_budget_e2e.py b/tests/e2e/quota_management/budgets/test_team_member_budget_e2e.py index 2717b23d9dc..0fd0a545660 100644 --- a/tests/e2e/quota_management/budgets/test_team_member_budget_e2e.py +++ b/tests/e2e/quota_management/budgets/test_team_member_budget_e2e.py @@ -46,7 +46,7 @@ def member(client: BudgetClient) -> Iterator[_Member]: Cleanups register progressively and run LIFO best-effort through ResourceManager, so a partial-setup failure still releases what came before and one failed delete never strands the rest on the shared proxy.""" - resources = ResourceManager(client=client.gateway) + resources = ResourceManager(client=client.proxy) try: marker = unique_marker() team_id = client.create_team(alias=f"e2e-team-member-{marker}", max_budget=TEAM_BUDGET) @@ -64,7 +64,7 @@ def member(client: BudgetClient) -> Iterator[_Member]: def _send(client: BudgetClient, key: str) -> str | None: """One member call; its response id (== the spend-log request_id) if it went through, else None.""" - match client.gateway.chat( + match client.proxy.chat( key, ChatBody( model=MODEL, @@ -83,7 +83,7 @@ class TestTeamMemberBudget: sent = frozenset(rid for rid in (_send(client, member.key) for _ in range(BURST)) if rid) assert sent, "no member call went through; cannot check attribution" - rows = client.gateway.poll_logs_for_key( + rows = client.proxy.poll_logs_for_key( member.key, predicate=lambda rs: bool(sent & {r.request_id for r in rs}) ) logged = [row for row in rows if row.request_id in sent] diff --git a/tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py b/tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py index 87855a9a1c1..f03518f8a17 100644 --- a/tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py +++ b/tests/e2e/quota_management/budgets/test_team_member_budget_isolation_e2e.py @@ -43,7 +43,7 @@ def pair(client: BudgetClient) -> Iterator[_Pair]: """One team with a large budget and two members on it: a tight member capped at a tiny per-team budget and a roomy member with headroom, each with their own key. Shared across the class and torn down LIFO best-effort when it finishes.""" - resources = ResourceManager(client=client.gateway) + resources = ResourceManager(client=client.proxy) try: marker = unique_marker() team_id = client.create_team(alias=f"e2e-member-iso-{marker}", max_budget=TEAM_BUDGET) @@ -71,7 +71,7 @@ def pair(client: BudgetClient) -> Iterator[_Pair]: def _roomy_send(client: BudgetClient, key: str) -> str: """One roomy-member call that must go through; returns its request id.""" - match client.gateway.chat( + match client.proxy.chat( key, ChatBody( model=MODEL, @@ -105,7 +105,7 @@ class TestTeamMemberBudgetIsolation: client.chat(pair.tight_key, MODEL, f"tight {unique_marker()}", max_tokens=16) ), "tight member stopped being blocked once the peer spent" - rows = client.gateway.poll_logs_for_key( + rows = client.proxy.poll_logs_for_key( pair.roomy_key, predicate=lambda rs: bool(sent & {r.request_id for r in rs}) ) logged = [row for row in rows if row.request_id in sent] diff --git a/tests/e2e/quota_management/ratelimit/conftest.py b/tests/e2e/quota_management/ratelimit/conftest.py index 59dee5e65b3..0b00c87c45c 100644 --- a/tests/e2e/quota_management/ratelimit/conftest.py +++ b/tests/e2e/quota_management/ratelimit/conftest.py @@ -1,15 +1,16 @@ """Quota-management suite's `client` fixture. The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker -live in the parent tests/e2e/conftest.py. QuotaClient holds the shared Gateway, +live in the parent tests/e2e/conftest.py. QuotaClient holds the shared ProxyClient, so the `resources` fixture cleans up keys through it. """ import pytest from quota_client import QuotaClient, build_client +from proxy_client import ProxyClient @pytest.fixture(scope="session") -def client() -> QuotaClient: - return build_client() +def client(proxy: ProxyClient) -> QuotaClient: + return build_client(proxy) diff --git a/tests/e2e/quota_management/ratelimit/quota_client.py b/tests/e2e/quota_management/ratelimit/quota_client.py index 806ab1d1557..a3a467a1d71 100644 --- a/tests/e2e/quota_management/ratelimit/quota_client.py +++ b/tests/e2e/quota_management/ratelimit/quota_client.py @@ -1,4 +1,4 @@ -"""Client for the quota-management suite: the shared Gateway plus raw chat +"""Client for the quota-management suite: the shared ProxyClient plus raw chat calls judged by HTTP status, body, and headers (a rate-limit block is a 429 whose body and retry-after header carry the contract, not a typed success model).""" @@ -7,19 +7,19 @@ from __future__ import annotations from dataclasses import dataclass -from e2e_gateway import Gateway, build_gateway +from proxy_client import ProxyClient from e2e_http import StreamingResponse from models import ChatBody, ChatMessage @dataclass(frozen=True, slots=True) class QuotaClient: - gateway: Gateway + proxy: ProxyClient def chat(self, key: str, model: str, content: str, *, max_tokens: int = 16) -> StreamingResponse: - return self.gateway.transport.send( + return self.proxy.transport.send( "/chat/completions", - headers=self.gateway.transport.bearer(key), + headers=self.proxy.transport.bearer(key), json=ChatBody( model=model, messages=[ChatMessage(role="user", content=content)], @@ -28,5 +28,5 @@ class QuotaClient: ) -def build_client() -> QuotaClient: - return QuotaClient(gateway=build_gateway()) +def build_client(proxy: ProxyClient) -> QuotaClient: + return QuotaClient(proxy=proxy) diff --git a/tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py b/tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py index a6c15b79bb1..7d87686b06c 100644 --- a/tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py +++ b/tests/e2e/quota_management/ratelimit/test_rate_limit_e2e.py @@ -131,8 +131,8 @@ def _limited_key( rpm_limit: int | None = None, tpm_limit: int | None = None, ) -> str: - key = client.gateway.generate_key(KeyGenerateBody(models=[MODEL], rpm_limit=rpm_limit, tpm_limit=tpm_limit)) - resources.defer(lambda: client.gateway.delete_key(key)) + key = client.proxy.generate_key(KeyGenerateBody(models=[MODEL], rpm_limit=rpm_limit, tpm_limit=tpm_limit)) + resources.defer(lambda: client.proxy.delete_key(key)) return key @@ -147,7 +147,7 @@ def _first_ok(client: QuotaClient, key: str) -> _FirstOk: cache picks it up, so retry on 401 to a deadline; a 401 never reaches the rate limiter, so only the successful call consumes budget. Any other failure is behavior under test and fails hard.""" - deadline = time.monotonic() + client.gateway.poll_timeout + deadline = time.monotonic() + client.proxy.poll_timeout while True: sent_at = time.monotonic() outcome = _chat(client, key) @@ -155,7 +155,7 @@ def _first_ok(client: QuotaClient, key: str) -> _FirstOk: return _FirstOk(sent_at=sent_at, response=outcome) if outcome.status_code != 401 or time.monotonic() >= deadline: require_successful_call(outcome) - time.sleep(client.gateway.poll_interval) + time.sleep(client.proxy.poll_interval) def _assert_rate_limited(outcome: StreamingResponse, limit_type: str) -> None: @@ -178,7 +178,7 @@ class TestKeyRateLimits: @pytest.mark.covers("quota_management.ratelimit.rpm.blocks_over_limit") def test_rpm_limit_blocks_over_limit(self, client: QuotaClient, resources: ResourceManager) -> None: key = _limited_key(client, resources, rpm_limit=3) - info = client.gateway.key_info(key) + info = client.proxy.key_info(key) assert info.rpm_limit == 3, f"/key/info reports rpm_limit {info.rpm_limit}, configured 3" _ = _first_ok(client, key) @@ -190,7 +190,7 @@ class TestKeyRateLimits: @pytest.mark.covers("quota_management.ratelimit.tpm.blocks_over_limit") def test_tpm_limit_blocks_over_limit(self, client: QuotaClient, resources: ResourceManager) -> None: key = _limited_key(client, resources, tpm_limit=TPM_LIMIT) - info = client.gateway.key_info(key) + info = client.proxy.key_info(key) assert info.tpm_limit == TPM_LIMIT, f"/key/info reports tpm_limit {info.tpm_limit}, configured {TPM_LIMIT}" first = _first_ok(client, key) @@ -213,7 +213,7 @@ class TestKeyRateLimits: first = _first_ok(client, key) _assert_rate_limited(_chat(client, key), "requests") - deadline = time.monotonic() + client.gateway.poll_timeout + deadline = time.monotonic() + client.proxy.poll_timeout while time.monotonic() < deadline: attempt_sent_at = time.monotonic() outcome = _chat(client, key) @@ -228,7 +228,7 @@ class TestKeyRateLimits: assert outcome.status_code == 429, ( f"while the window drains only 429s are acceptable, got {outcome.status_code}: {outcome.body[:300]}" ) - time.sleep(client.gateway.poll_interval) + time.sleep(client.proxy.poll_interval) pytest.fail("a blocked key never recovered after the rate-limit window elapsed") @pytest.mark.covers("quota_management.ratelimit.rpm.headers_report_remaining") diff --git a/tests/e2e/quota_management/spend_tracking/conftest.py b/tests/e2e/quota_management/spend_tracking/conftest.py index 434af15b182..c31e6b3c090 100644 --- a/tests/e2e/quota_management/spend_tracking/conftest.py +++ b/tests/e2e/quota_management/spend_tracking/conftest.py @@ -1,8 +1,8 @@ """Spend-tracking suite's `client` fixture and driver-model registration. The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker -live in the parent tests/e2e/conftest.py. SpendClient exposes the shared Gateway -(GatewayProvider), so the `resources` fixture cleans up keys and customers this +live in the parent tests/e2e/conftest.py. SpendClient exposes the shared ProxyClient +(ProxyClientProvider), so the `resources` fixture cleans up keys and customers this suite creates. The suite drives real calls through three deployments. On the stage gateway they @@ -21,6 +21,7 @@ import pytest from models import LiteLLMParamsBody from spend_e2e_client import SpendClient, build_client +from proxy_client import ProxyClient def _driver_params(provider_model: str, env_var: str) -> LiteLLMParamsBody: @@ -38,18 +39,18 @@ DRIVER_MODELS: tuple[tuple[str, str, str], ...] = ( @pytest.fixture(scope="session") -def client() -> SpendClient: - return build_client() +def client(proxy: ProxyClient) -> SpendClient: + return build_client(proxy) @pytest.fixture(scope="session", autouse=True) def driver_models(client: SpendClient) -> Iterator[None]: - existing = frozenset(entry.model_name for entry in client.gateway.model_info()) + existing = frozenset(entry.model_name for entry in client.proxy.model_info()) created = tuple( - client.gateway.create_model(name, _driver_params(provider_model, env_var)) + client.proxy.create_model(name, _driver_params(provider_model, env_var)) for name, provider_model, env_var in DRIVER_MODELS if name not in existing ) yield for model_id in created: - client.gateway.delete_model(model_id) + client.proxy.delete_model(model_id) diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index c4991199187..29ca5eb2ce6 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -1,7 +1,7 @@ -"""Spend-tracking e2e client: a Gateway plus the spend-specific read endpoints. +"""Spend-tracking e2e client: a ProxyClient plus the spend-specific read endpoints. Generic proxy operations (keys, customers, chat/embed, route probing, SpendLogs -polling) come from the shared Gateway, DI'd in (composition, not inheritance). +polling) come from the shared ProxyClient, DI'd in (composition, not inheritance). This client adds only the spend surface: /spend/calculate, /spend/tags, key-spend polling, and the route probes the breadth test uses. @@ -27,7 +27,7 @@ from e2e_http import ( is_ok, unwrap, ) -from e2e_gateway import Gateway, build_gateway +from proxy_client import ProxyClient from models import ( ChatBody, ChatMessage, @@ -96,7 +96,7 @@ def _chat_body( @dataclass(frozen=True, slots=True) class SpendClient: - gateway: Gateway + proxy: ProxyClient def chat( self, @@ -108,19 +108,19 @@ class SpendClient: tags: list[str] | None = None, user: str | None = None, ) -> Result[ChatResponse]: - return self.gateway.chat( + return self.proxy.chat( key, _chat_body(model, content, max_tokens=max_tokens, tags=tags, user=user) ) def chat_stream( self, key: str, model: str, content: str, *, max_tokens: int | None = None ) -> StreamingResponse: - return self.gateway.chat_stream( + return self.proxy.chat_stream( key, _chat_body(model, content, max_tokens=max_tokens, stream=True) ) def embed(self, key: str, model: str, content: str) -> Result[EmbedResponse]: - return self.gateway.embed(key, EmbedBody(model=model, input=content)) + return self.proxy.embed(key, EmbedBody(model=model, input=content)) def poll_logs_for_key( self, @@ -129,15 +129,15 @@ class SpendClient: min_rows: int = 1, predicate: Callable[[list[SpendLogRow]], bool] | None = None, ) -> list[SpendLogRow]: - return self.gateway.poll_logs_for_key( + return self.proxy.poll_logs_for_key( key, min_rows=min_rows, predicate=predicate ) def calculate_spend(self, model: str, content: str) -> float: return unwrap( - self.gateway.transport.post( + self.proxy.transport.post( "/spend/calculate", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, json=SpendCalculateBody( model=model, messages=[ChatMessage(role="user", content=content)] ), @@ -146,9 +146,9 @@ class SpendClient: ).cost def spend_by_tags(self) -> list[TagSpend]: - result = self.gateway.transport.get( + result = self.proxy.transport.get( "/spend/tags", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=NoBody(), response_type=SpendTagsResponse, ) @@ -160,7 +160,7 @@ class SpendClient: def poll_tag_spend(self, tag: str, *, minimum: float = 0.0) -> TagSpend | None: """Poll /spend/tags until the tag's aggregate reaches `minimum`; last seen.""" - deadline = time.monotonic() + self.gateway.poll_timeout + deadline = time.monotonic() + self.proxy.poll_timeout entry: TagSpend | None = None while time.monotonic() < deadline: matches = [ @@ -170,17 +170,17 @@ class SpendClient: entry = matches[0] if (entry.total_spend or 0.0) >= minimum: return entry - time.sleep(self.gateway.poll_interval) + time.sleep(self.proxy.poll_interval) return entry def poll_key_spend(self, key: str, *, minimum: float = 0.0) -> float: - deadline = time.monotonic() + self.gateway.poll_timeout + deadline = time.monotonic() + self.proxy.poll_timeout spend = 0.0 while time.monotonic() < deadline: - spend = self.gateway.key_info(key).spend or 0.0 + spend = self.proxy.key_info(key).spend or 0.0 if spend > minimum: return spend - time.sleep(self.gateway.poll_interval) + time.sleep(self.proxy.poll_interval) return spend def spend_logs_page( @@ -191,9 +191,9 @@ class SpendClient: now = datetime.now(timezone.utc) fmt = "%Y-%m-%d %H:%M:%S" return unwrap( - self.gateway.transport.get( + self.proxy.transport.get( "/spend/logs/v2", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=SpendLogsPageParams( start_date=(now - timedelta(days=1)).strftime(fmt), end_date=(now + timedelta(days=1)).strftime(fmt), @@ -206,18 +206,18 @@ class SpendClient: ) def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult: - return self.gateway.transport.probe(path, params=params) + return self.proxy.transport.probe(path, params=params) def openapi(self) -> OpenAPISchema: return unwrap( - self.gateway.transport.get( + self.proxy.transport.get( "/openapi.json", - headers=self.gateway.transport.master, + headers=self.proxy.transport.master, params=NoBody(), response_type=OpenAPISchema, ) ) -def build_client() -> SpendClient: - return SpendClient(gateway=build_gateway()) +def build_client(proxy: ProxyClient) -> SpendClient: + return SpendClient(proxy=proxy) diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py index 2f0ffae44e3..465046e89af 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py @@ -475,12 +475,12 @@ def test_spend_logs_endpoint_returns_spend( ) ) - gateway = client.gateway - deadline = time.monotonic() + gateway.poll_timeout + proxy = client.proxy + deadline = time.monotonic() + proxy.poll_timeout while True: - result = gateway.transport.get( + result = proxy.transport.get( "/spend/logs", - headers=gateway.transport.master, + headers=proxy.transport.master, params=SpendLogsParams(api_key=scoped_key), response_type=SpendLogs, ) @@ -493,4 +493,4 @@ def test_spend_logs_endpoint_returns_spend( f"/spend/logs never surfaced the key's spend before the deadline; " f"saw {_summarize(rows)}" ) - time.sleep(gateway.poll_interval) + time.sleep(proxy.poll_interval) diff --git a/tests/e2e/router/complexity_router_client.py b/tests/e2e/router/complexity_router_client.py index 929acbb3461..0093fc7480d 100644 --- a/tests/e2e/router/complexity_router_client.py +++ b/tests/e2e/router/complexity_router_client.py @@ -1,20 +1,20 @@ """Client for the complexity auto-router e2e tests. -The suite drives the shared /chat/completions and spend-log reads on the Gateway, -so this client only carries the Gateway the shared lifecycle needs for cleanup. +The suite drives the shared /chat/completions and spend-log reads on the ProxyClient, +so this client only carries the ProxyClient the shared lifecycle needs for cleanup. """ from __future__ import annotations from dataclasses import dataclass -from e2e_gateway import Gateway, build_gateway +from proxy_client import ProxyClient @dataclass(frozen=True, slots=True) class ComplexityRouterClient: - gateway: Gateway + proxy: ProxyClient -def build_client() -> ComplexityRouterClient: - return ComplexityRouterClient(gateway=build_gateway()) +def build_client(proxy: ProxyClient) -> ComplexityRouterClient: + return ComplexityRouterClient(proxy=proxy) diff --git a/tests/e2e/router/conftest.py b/tests/e2e/router/conftest.py index 344d8ab5c13..8ddc19aa94f 100644 --- a/tests/e2e/router/conftest.py +++ b/tests/e2e/router/conftest.py @@ -2,7 +2,7 @@ The shared lifecycle (resources/scoped_key), proxy liveness gate, and e2e marker live in the parent tests/e2e/conftest.py. ComplexityRouterClient holds the shared -Gateway, so the `resources` fixture cleans up keys this suite creates. +ProxyClient, so the `resources` fixture cleans up keys this suite creates. Also registers `complexity-smart-router` via management /model/new when the proxy does not already list it (compose has it in static config; stage does not). @@ -16,7 +16,7 @@ import pytest from requests import RequestException from complexity_router_client import ComplexityRouterClient, build_client -from e2e_gateway import Gateway +from proxy_client import ProxyClient from e2e_http import NoBody, Success from lifecycle import ResourceManager from models import ( @@ -46,27 +46,27 @@ ROUTER_KEY_MODELS = [ROUTER_MODEL, "gpt-5.5", "claude-haiku-4-5"] @pytest.fixture(scope="session") -def client() -> ComplexityRouterClient: - return build_client() +def client(proxy: ProxyClient) -> ComplexityRouterClient: + return build_client(proxy) -def _model_is_servable(gateway: Gateway, model_name: str) -> bool: - result = gateway.transport.get( +def _model_is_servable(proxy: ProxyClient, model_name: str) -> bool: + result = proxy.transport.get( "/v1/models", - headers=gateway.transport.master, + headers=proxy.transport.master, params=NoBody(), response_type=ModelsListResponse, ) return isinstance(result, Success) and any(entry.id == model_name for entry in result.data.data) -def _router_is_callable(gateway: Gateway) -> bool: +def _router_is_callable(proxy: ProxyClient) -> bool: """True only when a short chat against the virtual router succeeds; every error (the Invalid-model-name reload race, but also 401, 5xx, and network) counts as not-callable so infra/auth blips can't be mistaken for a working router.""" - key = gateway.generate_key(KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-probe")) + key = proxy.generate_key(KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-probe")) try: - result = gateway.chat( + result = proxy.chat( key, ChatBody( model=ROUTER_MODEL, @@ -75,7 +75,7 @@ def _router_is_callable(gateway: Gateway) -> bool: ), ) finally: - gateway.delete_key(key) + proxy.delete_key(key) return isinstance(result, Success) @@ -86,18 +86,18 @@ def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] # """Ensure the complexity router virtual model exists for this session. Compose already declares it in docker-compose.yml; stage does not. Register - via Gateway.create_model (waits for data-plane /v1/models) when missing, then + via ProxyClient.create_model (waits for data-plane /v1/models) when missing, then probe a real chat so a list-only false positive cannot pass the fixture. """ - gateway = client.gateway - if _model_is_servable(gateway, ROUTER_MODEL) and _router_is_callable(gateway): + proxy = client.proxy + if _model_is_servable(proxy, ROUTER_MODEL) and _router_is_callable(proxy): yield return try: - model_id = gateway.create_model(ROUTER_MODEL, ROUTER_PARAMS) + model_id = proxy.create_model(ROUTER_MODEL, ROUTER_PARAMS) except (AssertionError, RequestException) as exc: - if _model_is_servable(gateway, ROUTER_MODEL) and _router_is_callable(gateway): + if _model_is_servable(proxy, ROUTER_MODEL) and _router_is_callable(proxy): yield return raise AssertionError( @@ -106,7 +106,7 @@ def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] # ) from exc try: - if not _router_is_callable(gateway): + if not _router_is_callable(proxy): raise AssertionError( f"{ROUTER_MODEL!r} registered as {model_id!r} and listed on " f"/v1/models but chat still returns Invalid model name; " @@ -114,14 +114,14 @@ def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] # ) yield finally: - gateway.delete_model(model_id) + proxy.delete_model(model_id) @pytest.fixture def complexity_key(resources: ResourceManager, client: ComplexityRouterClient) -> str: """Per-test key allowed to call the complexity router and its tier backends.""" - key = client.gateway.generate_key( + key = client.proxy.generate_key( KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-router") ) - resources.defer(lambda: client.gateway.delete_key(key)) + resources.defer(lambda: client.proxy.delete_key(key)) return key diff --git a/tests/e2e/router/test_complexity_router_e2e.py b/tests/e2e/router/test_complexity_router_e2e.py index e9ec020994c..a495e2fdf4d 100644 --- a/tests/e2e/router/test_complexity_router_e2e.py +++ b/tests/e2e/router/test_complexity_router_e2e.py @@ -48,7 +48,7 @@ class TestComplexityRouterLlmClassifier: self, client: ComplexityRouterClient, complexity_key: str ) -> None: chat = unwrap( - client.gateway.chat( + client.proxy.chat( complexity_key, ChatBody( model=ROUTER_MODEL, @@ -59,7 +59,7 @@ class TestComplexityRouterLlmClassifier: ) assert chat.choices, f"router returned no choices: {chat}" - rows = client.gateway.poll_logs_for_key(complexity_key, min_rows=1) + rows = client.proxy.poll_logs_for_key(complexity_key, min_rows=1) served = [row.model for row in rows] # Exactly one spend row for the routed completion (not the classifier sub-call). # Membership allows alias vs provider-prefixed forms across compose and stage. diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py new file mode 100644 index 00000000000..7d1b10e6bc3 --- /dev/null +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -0,0 +1,385 @@ +"""Tests for the optional Rust-backed Anthropic Messages path.""" + +import importlib +from typing import cast + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, +) +from litellm.types.router import GenericLiteLLMParams + +rust_messages = importlib.import_module("litellm.rust_bridge.messages") +rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") + +FAKE_MESSAGES_RESPONSE: dict[str, object] = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "hello world"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 5, "output_tokens": 3}, +} + +REQUEST_BODY: dict[str, object] = { + "model": "claude-sonnet-4-5", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], +} + + +class RecordingMessages: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> dict[str, object]: + self.calls.append( + { + "model": model, + "body": body, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "timeout_seconds": timeout_seconds, + } + ) + return dict(FAKE_MESSAGES_RESPONSE) + + +class RecordingAsyncMessages: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> dict[str, object]: + self.calls.append( + { + "model": model, + "body": body, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "timeout_seconds": timeout_seconds, + } + ) + return dict(FAKE_MESSAGES_RESPONSE) + + +class ExplodingAsyncMessages: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, **kwargs: object) -> dict[str, object]: + self.calls += 1 + raise AssertionError("bridge must not be called") + + +class RaisingAsyncMessages: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, **kwargs: object) -> dict[str, object]: + self.calls += 1 + raise RuntimeError("upstream request failed with status 400: bad request") + + +class NoneAsyncMessages: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, **kwargs: object) -> dict[str, object] | None: + self.calls += 1 + return None + + +@pytest.fixture(autouse=True) +def _reset_rust_flag(): + litellm.use_litellm_rust(False, messages=None, amessages=None) + rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL + yield + litellm.use_litellm_rust(False, messages=None, amessages=None) + rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL + + +def test_load_rust_messages_returns_injected_impl(): + bridge = RecordingMessages() + litellm.use_litellm_rust(True, messages=bridge) + assert rust_messages.load_rust_messages() is bridge + + +def test_configuring_messages_does_not_enable_ocr(): + from litellm.rust_bridge.ocr import rust_ocr_enabled + + litellm.use_litellm_rust(False) + assert rust_ocr_enabled() is False + + litellm.use_litellm_rust(True, messages=RecordingMessages()) + + assert rust_ocr_enabled() is False + + +def test_bare_use_litellm_rust_still_toggles_ocr(): + from litellm.rust_bridge.ocr import rust_ocr_enabled + + litellm.use_litellm_rust(True) + assert rust_ocr_enabled() is True + + litellm.use_litellm_rust(False) + assert rust_ocr_enabled() is False + + +def test_load_rust_amessages_returns_injected_impl(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + assert rust_messages.load_rust_amessages() is bridge + + +def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch): + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: None, + ) + litellm.use_litellm_rust(True) + assert rust_messages.load_rust_messages() is None + result = rust_messages.messages( + model="claude", + body=REQUEST_BODY, + api_key="k", + api_base="b", + custom_llm_provider="azure_ai", + extra_headers={}, + timeout=30.0, + ) + assert result is None + + +def test_messages_wrapper_forwards_args_and_converts_timeout(): + bridge = RecordingMessages() + litellm.use_litellm_rust(True, messages=bridge) + + response = rust_messages.messages( + model="claude-sonnet-4-5", + body=REQUEST_BODY, + api_key="sk-azure", + api_base="https://resource.services.ai.azure.com/anthropic", + custom_llm_provider="azure_ai", + extra_headers={"anthropic-beta": "token-efficient-tools-2025-02-19"}, + timeout=httpx.Timeout(600.0, read=42.0), + ) + + assert response == FAKE_MESSAGES_RESPONSE + assert bridge.calls[0] == { + "model": "claude-sonnet-4-5", + "body": REQUEST_BODY, + "api_key": "sk-azure", + "api_base": "https://resource.services.ai.azure.com/anthropic", + "custom_llm_provider": "azure_ai", + "extra_headers": {"anthropic-beta": "token-efficient-tools-2025-02-19"}, + "timeout_seconds": 42.0, + } + + +@pytest.mark.asyncio +async def test_amessages_wrapper_forwards_args(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await rust_messages.amessages( + model="claude-sonnet-4-5", + body=REQUEST_BODY, + api_key="sk-azure", + api_base="https://resource.services.ai.azure.com/anthropic", + custom_llm_provider="azure_ai", + extra_headers=None, + timeout=12.5, + ) + + assert response == FAKE_MESSAGES_RESPONSE + assert bridge.calls[0]["model"] == "claude-sonnet-4-5" + assert bridge.calls[0]["timeout_seconds"] == 12.5 + + +def _gate(**overrides): + kwargs = { + "custom_llm_provider": "azure_ai", + "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True), + "stream": False, + "rust_stream_eligible": False, + "model": "claude-sonnet-4-5", + "api_key": "sk-azure", + "api_base": "https://resource.services.ai.azure.com/anthropic", + "headers": {"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"}, + "request_body": dict(REQUEST_BODY), + "timeout": 30.0, + } + kwargs.update(overrides) + return BaseLLMHTTPHandler._maybe_rust_anthropic_messages(**kwargs) + + +@pytest.mark.asyncio +async def test_gate_invokes_rust_and_marks_response_header(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate() + + assert response is not None + assert response["id"] == "msg_123" + assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} + call = bridge.calls[0] + assert call["model"] == "claude-sonnet-4-5" + assert call["body"] == REQUEST_BODY + assert call["api_key"] == "sk-azure" + assert call["api_base"] == "https://resource.services.ai.azure.com/anthropic" + assert call["extra_headers"] == {"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"} + assert call["timeout_seconds"] == 30.0 + + +@pytest.mark.asyncio +async def test_gate_invokes_rust_for_native_anthropic_provider(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate( + custom_llm_provider="anthropic", + api_key="sk-ant-test", + api_base="https://api.anthropic.com", + headers={"anthropic-version": "2023-06-01"}, + litellm_params=GenericLiteLLMParams(api_key="sk-ant-test", rust=True), + ) + + assert response is not None + assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} + call = bridge.calls[0] + assert call["custom_llm_provider"] == "anthropic" + assert call["api_key"] == "sk-ant-test" + assert call["api_base"] == "https://api.anthropic.com" + assert call["extra_headers"] == {"anthropic-version": "2023-06-01"} + + +@pytest.mark.asyncio +async def test_gate_falls_back_to_python_when_bridge_raises(): + bridge = RaisingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate() + + assert response is None + assert bridge.calls == 1 + + +@pytest.mark.asyncio +async def test_gate_skips_rust_when_flag_absent(): + bridge = ExplodingAsyncMessages() + litellm.use_litellm_rust(False, amessages=bridge) + + response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure")) + + assert response is None + assert bridge.calls == 0 + + +@pytest.mark.asyncio +async def test_gate_skips_rust_when_flag_false(): + bridge = ExplodingAsyncMessages() + litellm.use_litellm_rust(False, amessages=bridge) + + response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False)) + + assert response is None + assert bridge.calls == 0 + + +@pytest.mark.asyncio +async def test_gate_skips_rust_for_non_listed_provider(): + bridge = NoneAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate(custom_llm_provider="openai") + + assert response is None + assert bridge.calls == 1 + + +@pytest.mark.asyncio +async def test_gate_skips_rust_when_streaming_but_not_eligible(): + bridge = ExplodingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate(stream=True, rust_stream_eligible=False) + + assert response is None + assert bridge.calls == 0 + + +@pytest.mark.asyncio +async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + streaming_body = {**REQUEST_BODY, "stream": True} + response = await _gate( + stream=True, + rust_stream_eligible=True, + request_body=streaming_body, + ) + + assert response is not None + assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} + assert "stream" not in bridge.calls[0]["body"] + assert bridge.calls[0]["body"] == REQUEST_BODY + + +@pytest.mark.asyncio +async def test_fake_stream_wraps_rust_response_as_anthropic_sse(): + response = cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE)) + stream = BaseLLMHTTPHandler._rust_anthropic_messages_fake_stream(response) + + assert stream._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + + chunks = [chunk async for chunk in stream] + joined = b"".join(chunks) + + assert b"event: message_start" in joined + assert b"event: content_block_delta" in joined + assert b"hello world" in joined + assert b"event: message_stop" in joined + + +@pytest.mark.asyncio +async def test_gate_falls_back_when_bridge_unavailable(monkeypatch): + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: None, + ) + litellm.use_litellm_rust(True) + + response = await _gate() + + assert response is None diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index 48190a798da..6b1da4c2952 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -16,10 +16,12 @@ from litellm.integrations.otel import ( # noqa: E402 from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402 from litellm.integrations.otel.plumbing import providers # noqa: E402 from litellm.integrations.otel.emitter import SpanEmitter # noqa: E402 +from litellm.integrations.otel.emitter import stamp_error # noqa: E402 from litellm.integrations.otel.model.payloads import ( # noqa: E402 GuardrailSpanData, LLMCallSpanData, ServiceSpanData, + SpanError, ) from litellm.integrations.otel.model.spans import SPAN_REGISTRY, SpanRole # noqa: E402 @@ -155,6 +157,46 @@ def test_error_span_sets_status_and_error_type(): assert span.attributes["error.type"] == "RateLimitError" +def test_stamp_error_writes_full_attribute_set_and_event(): + engine, exporter = _engine() + span = engine.start_span(SpanRole.PROXY_REQUEST, "POST /chat/completions") + result = stamp_error( + span, SpanError("ProxyException", "boom", code="401", stack_trace="tb", llm_provider="anthropic") + ) + span.end() + (s,) = exporter.get_finished_spans() + assert result == ("ProxyException", "boom") + assert s.attributes["error.type"] == "ProxyException" + assert s.attributes["error.message"] == "boom" + assert s.attributes["litellm.provider.error.code"] == "401" + assert s.attributes["litellm.provider.error.stack_trace"] == "tb" + assert s.attributes["litellm.provider.error.llm_provider"] == "anthropic" + assert s.status.status_code is StatusCode.ERROR + assert [e.name for e in s.events] == ["exception"] + + +def test_stamp_error_opt_outs_skip_status_and_event(): + engine, exporter = _engine() + span = engine.start_span(SpanRole.PROXY_REQUEST, "POST /chat/completions") + stamp_error(span, SpanError("ProxyException", "boom", code="401"), record_event=False, set_status=False) + span.end() + (s,) = exporter.get_finished_spans() + assert s.attributes["error.type"] == "ProxyException" + assert s.attributes["litellm.provider.error.code"] == "401" + assert s.status.status_code is StatusCode.UNSET + assert s.events == () + + +def test_stamp_error_without_type_or_message_is_noop(): + engine, exporter = _engine() + span = engine.start_span(SpanRole.PROXY_REQUEST, "POST /chat/completions") + assert stamp_error(span, SpanError()) is None + span.end() + (s,) = exporter.get_finished_spans() + assert "error.type" not in s.attributes + assert s.status.status_code is StatusCode.UNSET + + def test_hierarchy_and_kinds_match_registry(): engine, exporter = _engine() data = LLMCallSpanData.from_standard_logging_payload(_payload()) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index b5e077e3561..5f6002f4cdf 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -860,6 +860,120 @@ def test_guardrail_span_anchors_to_root_inside_active_phase_span(): assert guard.parent.span_id != auth_span.get_span_context().span_id +# --------------------------------------------------------------------------- # +# LIT-4179 — proxy-level failures that never reach an LLM call must still stamp +# the structured error.* attributes onto the request's spans, restoring the v1 +# behavior v2 dropped when it stopped subclassing ``OpenTelemetry``. +# --------------------------------------------------------------------------- # + + +def _proxy_exc(message, code): + from litellm.proxy._types import ProxyException + + return ProxyException(message=message, type="bad_request_error", param=None, code=code) + + +def test_async_post_call_failure_hook_stamps_error_on_root_span(): + """PATH B: an endpoint-level failure (empty body rejected before dispatch) + reaches ``async_post_call_failure_hook``; it must stamp error.* + an exception + event on the anchored request root span.""" + from litellm.proxy._types import UserAPIKeyAuth + + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + set_request_root_span(server) + exc = _proxy_exc("litellm.BadRequestError: messages is required", 400) + result = asyncio.run( + logger.async_post_call_failure_hook( + request_data={}, original_exception=exc, user_api_key_dict=UserAPIKeyAuth() + ) + ) + server.end() + assert result is None + (span,) = exporter.get_finished_spans() + assert span.attributes["error.type"] == "ProxyException" + assert "messages is required" in span.attributes["error.message"] + assert span.attributes["litellm.provider.error.code"] == "400" + assert span.status.status_code is StatusCode.ERROR + assert any(e.name == "exception" for e in span.events) + + +def test_async_post_call_failure_hook_falls_back_to_user_api_key_parent_span(): + """With no anchor set (a path that never captured the root), the hook must fall + back to ``user_api_key_dict.parent_otel_span`` rather than dropping the error.""" + from litellm.proxy._types import UserAPIKeyAuth + + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + asyncio.run( + logger.async_post_call_failure_hook( + request_data={}, + original_exception=_proxy_exc("boom", 401), + user_api_key_dict=UserAPIKeyAuth(parent_otel_span=server), + ) + ) + server.end() + (span,) = exporter.get_finished_spans() + assert span.attributes["error.type"] == "ProxyException" + assert span.attributes["litellm.provider.error.code"] == "401" + + +def test_record_error_attributes_on_span_decorates_without_ending(): + """PATH A: a failure that dies before any LLM-call span (malformed body, + validation) is stamped onto the instrumentor-owned SERVER span. The method must + not end the span or emit a duplicate exception event, and must pin error.code + to the real response status (not the exception's own code).""" + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + logger.record_error_attributes_on_span(server, _proxy_exc("Invalid JSON body", 400), 422) + assert server.is_recording() + server.end() + (span,) = exporter.get_finished_spans() + assert span.attributes["error.type"] == "ProxyException" + assert span.attributes["error.message"] == "Invalid JSON body" + assert span.attributes["litellm.provider.error.code"] == "422" + assert all(e.name != "exception" for e in span.events) + + +def test_record_error_attributes_on_span_ignores_below_400_and_missing_span(): + logger, _ = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + logger.record_error_attributes_on_span(None, _proxy_exc("boom", 400), 400) # no span → no-op + logger.record_error_attributes_on_span(server, None, 400) # no exception → no-op + server.end() + assert "error.type" not in (server.attributes or {}) + + +def test_start_phase_span_stamps_error_attributes_on_failure(): + """An ``auth`` phase span that dies (expired key) must carry the structured + error.* attributes, not only the exception event ``use_span`` records.""" + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + set_request_root_span(server) + exc = _proxy_exc("Authentication Error, ExpiredToken", 401) + with trace.use_span(server, end_on_exit=False): + with contextlib.suppress(Exception): + with logger.start_phase_span("auth /chat/completions"): + raise exc + server.end() + by_name = {s.name: s for s in exporter.get_finished_spans()} + auth = by_name["auth /chat/completions"] + assert auth.attributes["error.type"] == "ProxyException" + assert "ExpiredToken" in auth.attributes["error.message"] + assert auth.attributes["litellm.provider.error.code"] == "401" + assert auth.status.status_code is StatusCode.ERROR + assert any(e.name == "exception" for e in auth.events) + + +def test_start_phase_span_success_carries_no_error(): + logger, exporter = _logger() + with logger.start_phase_span("auth /chat/completions"): + pass + (span,) = exporter.get_finished_spans() + assert "error.type" not in span.attributes + assert span.status.status_code is not StatusCode.ERROR + + def test_real_logging_pre_call_opens_span_end_to_end(): """Regression guard: a real ``LiteLLMLoggingObj.pre_call`` must fire ``log_pre_api_call`` on the V2 logger (via ``litellm.input_callback``), so the diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 03e763a4161..6809799d34f 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -13,6 +13,7 @@ sys.path.insert( from litellm import get_model_info, supports_reasoning, supports_vision from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig +from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import ( ChatCompletionMessageToolCall, Function, @@ -32,6 +33,105 @@ def force_local_model_cost(monkeypatch): litellm.model_cost = get_model_cost_map(url=litellm.model_cost_map_url) +def test_validate_environment_sets_session_affinity_from_litellm_session_id(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"litellm_session_id": "session-123"}, + api_key="test-key", + ) + + assert headers["x-session-affinity"] == "session-123" + + +def test_validate_environment_sets_session_affinity_from_metadata_session_id(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"metadata": {"session_id": "metadata-session-123"}}, + api_key="test-key", + ) + + assert headers["x-session-affinity"] == "metadata-session-123" + + +def test_validate_environment_sets_session_affinity_from_session_id(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"session_id": "session-id-123"}, + api_key="test-key", + ) + + assert headers["x-session-affinity"] == "session-id-123" + + +def test_validate_environment_sets_session_affinity_from_trace_id(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"litellm_trace_id": "trace-id-123"}, + api_key="test-key", + ) + + assert headers["x-session-affinity"] == "trace-id-123" + + +def test_validate_environment_does_not_set_session_affinity_without_session_id(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test-key", + ) + + assert "x-session-affinity" not in headers + + +def test_validate_environment_preserves_explicit_session_affinity_header(): + config = FireworksAIConfig() + + headers = config.validate_environment( + headers={"x-session-affinity": "explicit-session"}, + model="accounts/fireworks/models/test-model", + messages=[], + optional_params={}, + litellm_params={"litellm_session_id": "session-123"}, + api_key="test-key", + ) + + assert headers["x-session-affinity"] == "explicit-session" + + +def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id(): + assert ( + get_fireworks_session_id( + {"litellm_session_id": "session-123", "litellm_trace_id": "trace-123"} + ) + == "session-123" + ) + + def test_handle_message_content_with_tool_calls(): config = FireworksAIConfig() message = Message( diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py index 99dcaa36c75..3297750fa6e 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py @@ -5,12 +5,15 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) +import litellm from litellm.llms.fireworks_ai.cost_calculator import cost_per_token from litellm.types.utils import PromptTokensDetailsWrapper, Usage MODEL = "accounts/fireworks/models/glm-5p2" INPUT_COST = 1.4e-06 -CACHE_READ_COST = 2.6e-07 +# Read the cached rate from the price map so this test tracks the shipped value +# (glm-5p2 is $0.14/1M) instead of hardcoding a number that breaks when it changes. +CACHE_READ_COST = litellm.get_model_info(model=MODEL, custom_llm_provider="fireworks_ai")["cache_read_input_token_cost"] OUTPUT_COST = 4.4e-06 diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 5adc5b76990..95e8e6561f1 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -547,6 +547,103 @@ def test_vertex_ai_non_grounded_usage_omits_tool_use_tokens(): assert not hasattr(usage.prompt_tokens_details, "tool_use_tokens") +def test_response_has_search_grounding_detection(): + """ + Only groundingMetadata.webSearchQueries signals an actual Google Search. URL context also + emits groundingMetadata (groundingChunks but no webSearchQueries) and must not be treated + as search grounding. + """ + assert ( + VertexGeminiConfig._response_has_search_grounding( + {"candidates": [{"groundingMetadata": {"webSearchQueries": ["latest nobel physics"]}}]} + ) + is True + ) + assert ( + VertexGeminiConfig._response_has_search_grounding( + { + "candidates": [ + { + "urlContextMetadata": {"urlMetadata": []}, + "groundingMetadata": { + "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] + }, + } + ] + } + ) + is False + ) + assert ( + VertexGeminiConfig._response_has_search_grounding({"candidates": [{"groundingMetadata": {"webSearchQueries": []}}]}) + is False + ) + assert VertexGeminiConfig._response_has_search_grounding({"candidates": []}) is False + assert VertexGeminiConfig._response_has_search_grounding({}) is False + + +def test_vertex_ai_search_grounding_tool_use_tokens_excluded_from_prompt_tokens(): + """ + Grounding with Google Search retrieved tokens are not billed at the input token rate + (Google charges a separate per-request / per-query search fee), so toolUsePromptTokenCount + must be surfaced on prompt_tokens_details.tool_use_tokens but excluded from prompt_tokens. + See https://ai.google.dev/gemini-api/docs/pricing and + https://github.com/BerriAI/litellm/discussions/33198 + """ + v = VertexGeminiConfig() + completion_response = { + "candidates": [{"groundingMetadata": {"webSearchQueries": ["latest nobel physics"]}}], + "usageMetadata": UsageMetadata( + promptTokenCount=19, + candidatesTokenCount=304, + thoughtsTokenCount=122, + toolUsePromptTokenCount=142, + totalTokenCount=587, + ), + } + + usage = v._calculate_usage(completion_response=completion_response) + + assert usage.prompt_tokens == 19 + assert usage.completion_tokens == 304 + 122 + assert usage.total_tokens == 587 + assert usage.prompt_tokens_details.tool_use_tokens == 142 + assert usage.total_tokens - usage.prompt_tokens - usage.completion_tokens == 142 + + +def test_vertex_ai_url_context_tool_use_tokens_billed_as_input_tokens(): + """ + URL context / File Search / code execution tool-use tokens are billed as input tokens, so + toolUsePromptTokenCount is folded into prompt_tokens when the response is not search grounded. + """ + v = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "urlContextMetadata": {"urlMetadata": []}, + "groundingMetadata": { + "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] + }, + } + ], + "usageMetadata": UsageMetadata( + promptTokenCount=19, + candidatesTokenCount=304, + thoughtsTokenCount=122, + toolUsePromptTokenCount=142, + totalTokenCount=587, + ), + } + + usage = v._calculate_usage(completion_response=completion_response) + + assert usage.prompt_tokens == 19 + 142 + assert usage.completion_tokens == 304 + 122 + assert usage.total_tokens == 587 + assert usage.prompt_tokens_details.tool_use_tokens == 142 + assert usage.total_tokens - usage.prompt_tokens - usage.completion_tokens == 0 + + def test_streaming_chunk_includes_reasoning_tokens(): from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 17960e917a4..707374e7061 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -21,9 +21,12 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, + ClientSecretAuth, CredError, + IdJagConfig, NoneConfig, PassthroughConfig, + PrivateKeyJwtAuth, SharedKey, TokenExchangeConfig, ) @@ -35,6 +38,21 @@ def _server(**kwargs) -> MCPServer: return MCPServer(server_id="s", name="n", transport=MCPTransport.http, **kwargs) +def _id_jag_server(**overrides) -> MCPServer: + defaults = dict( + auth_type=MCPAuth.oauth2_id_jag, + url="https://mcp.example.com/mcp", + client_id="litellm-client-id", + client_secret="litellm-client-secret", + token_exchange_endpoint="https://idp.example.com/token", + id_jag_resource_token_endpoint="https://mcp-as.example.com/token", + audience="api://mcp-server", + scopes=["mcp.read", "mcp.write"], + ) + defaults.update(overrides) + return _server(**defaults) + + def test_none_maps_to_none_config(): spec = to_server_spec(_server(auth_type=None)) assert spec is not None @@ -413,3 +431,63 @@ def test_raise_token_exchange_challenge_uses_insufficient_claims_with_claims_pre assert 'error="invalid_token"' not in www assert f'claims="{base64.b64encode(claims.encode()).decode()}"' in www assert claims not in www # raw JSON never appears; only the base64 form + + +def test_id_jag_client_secret_maps_to_config(): + spec = to_server_spec(_id_jag_server()) + assert spec is not None and isinstance(spec.config, IdJagConfig) + assert spec.config.org_token_endpoint == "https://idp.example.com/token" + assert spec.config.resource_token_endpoint == "https://mcp-as.example.com/token" + assert spec.config.client_id == "litellm-client-id" + assert spec.config.audience == "api://mcp-server" + assert spec.config.scopes == ("mcp.read", "mcp.write") + # ID-JAG asserts the user's id_token; the access_token default maps to id_token. + assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:id_token" + assert isinstance(spec.config.client_auth, ClientSecretAuth) + assert spec.config.client_auth.client_secret.get_secret_value() == ( + "litellm-client-secret" + ) + + +def test_id_jag_private_key_maps_to_private_key_jwt_auth(): + spec = to_server_spec( + _id_jag_server( + client_secret=None, + client_private_key="PEM-DATA", + client_private_key_id="kid-1", + client_assertion_signing_alg="RS384", + ) + ) + assert spec is not None and isinstance(spec.config, IdJagConfig) + assert isinstance(spec.config.client_auth, PrivateKeyJwtAuth) + assert spec.config.client_auth.private_key.get_secret_value() == "PEM-DATA" + assert spec.config.client_auth.key_id == "kid-1" + assert spec.config.client_auth.signing_alg == "RS384" + + +def test_id_jag_private_key_wins_over_client_secret(): + spec = to_server_spec(_id_jag_server(client_private_key="PEM-DATA")) + assert spec is not None and isinstance(spec.config, IdJagConfig) + assert isinstance(spec.config.client_auth, PrivateKeyJwtAuth) + + +def test_id_jag_honors_explicit_subject_token_type(): + spec = to_server_spec( + _id_jag_server(subject_token_type="urn:ietf:params:oauth:token-type:saml2") + ) + assert spec is not None and isinstance(spec.config, IdJagConfig) + assert spec.config.subject_token_type == "urn:ietf:params:oauth:token-type:saml2" + + +@pytest.mark.parametrize( + "server", + [ + _id_jag_server(token_exchange_endpoint=None), + _id_jag_server(id_jag_resource_token_endpoint=None), + _id_jag_server(client_id=None), + _id_jag_server(client_secret=None, client_private_key=None), + ], +) +def test_id_jag_half_configured_defers_to_v1(server): + # A half-configured server must defer (None) rather than 500 at IdJagConfig construction. + assert to_server_spec(server) is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index c88027abcd4..ba7720ffd51 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -16,8 +16,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( AwsSigV4Config, Byok, ClientCredentialsConfig, + ClientSecretAuth, CredError, Error, + IdJagConfig, NoneConfig, NoOpAuth, Ok, @@ -34,10 +36,40 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto OAuthToken, TokenStoreUnavailable, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( + ExchangedToken, +) _SUBJECT = Subject(tenant_id="", subject_id="") +def _id_jag_config() -> IdJagConfig: + return IdJagConfig( + org_token_endpoint="https://idp.example.com/token", + resource_token_endpoint="https://mcp-as.example.com/token", + client_id="litellm", + client_auth=ClientSecretAuth(client_secret=SecretStr("s")), + audience="api://mcp", + scopes=("mcp.read",), + ) + + +class _FakeTokenEndpoint: + """Records each fetch and returns the next canned Result, leg by leg.""" + + def __init__(self, results: list[Result[ExchangedToken, CredError]]) -> None: + self._results = list(results) + self.calls: list[tuple[str, str, dict[str, str]]] = [] + + async def fetch(self, endpoint, client_id, grant_params, client_auth): + self.calls.append((endpoint, client_id, dict(grant_params))) + return self._results.pop(0) + + +def _with_inbound(token: str) -> Subject: + return Subject(tenant_id="", subject_id="alice", inbound_token=SecretStr(token)) + + def _spec(config): return ServerSpec(server_id="s", resource="https://upstream.example.com", config=config) @@ -305,3 +337,183 @@ async def test_unbuilt_arms_fail_closed_with_not_implemented(label, config): result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, _spec(config)) assert isinstance(result, Error) assert result.error.tag == "not_implemented" + + +@pytest.mark.asyncio +async def test_id_jag_runs_both_legs_and_returns_the_leg2_bearer(): + endpoint = _FakeTokenEndpoint( + [ + Ok(ExchangedToken(access_token="the-id-jag", expires_in=300)), + Ok(ExchangedToken(access_token="final-access", expires_in=3600)), + ] + ) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + result = await provider.resolve_credentials( + _with_inbound("user-id-token"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Ok) + assert _emitted(result.ok)["Authorization"] == "Bearer final-access" + + leg1_endpoint, _, leg1_params = endpoint.calls[0] + assert leg1_endpoint == "https://idp.example.com/token" + assert ( + leg1_params["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + ) + assert ( + leg1_params["requested_token_type"] == "urn:ietf:params:oauth:token-type:id-jag" + ) + assert leg1_params["subject_token"] == "user-id-token" + + leg2_endpoint, _, leg2_params = endpoint.calls[1] + assert leg2_endpoint == "https://mcp-as.example.com/token" + assert leg2_params["grant_type"] == "urn:ietf:params:oauth:grant-type:jwt-bearer" + # The leg-1 token is forwarded verbatim as the leg-2 assertion. + assert leg2_params["assertion"] == "the-id-jag" + + +@pytest.mark.asyncio +async def test_id_jag_without_inbound_token_is_precondition_required_no_http(): + endpoint = _FakeTokenEndpoint([]) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "precondition_required" + assert endpoint.calls == [] + + +@pytest.mark.asyncio +async def test_id_jag_propagates_a_leg1_error_without_calling_leg2(): + endpoint = _FakeTokenEndpoint( + [Error(CredError.of_upstream_unavailable("leg1 down"))] + ) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + result = await provider.resolve_credentials( + _with_inbound("user-id-token"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert "leg1 down" in result.error.summary + assert len(endpoint.calls) == 1 + + +@pytest.mark.asyncio +async def test_id_jag_propagates_a_leg2_error(): + endpoint = _FakeTokenEndpoint( + [ + Ok(ExchangedToken(access_token="the-id-jag", expires_in=300)), + Error(CredError.of_upstream_unavailable("leg2 forbidden")), + ] + ) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + result = await provider.resolve_credentials( + _with_inbound("user-id-token"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert "leg2 forbidden" in result.error.summary + assert len(endpoint.calls) == 2 + + +def _two_leg_ok(bearer: str) -> list: + return [ + Ok(ExchangedToken(access_token="the-id-jag", expires_in=300)), + Ok(ExchangedToken(access_token=bearer, expires_in=3600)), + ] + + +@pytest.mark.asyncio +async def test_id_jag_reuses_the_cached_bearer_for_an_unchanged_config(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("first-bearer")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + + first = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) + second = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(second.ok)["Authorization"] == "Bearer first-bearer" + assert len(endpoint.calls) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "changed", + [ + _id_jag_config().model_copy(update={"audience": "api://other"}), + _id_jag_config().model_copy(update={"resource": "https://other.example.com/mcp"}), + _id_jag_config().model_copy(update={"scopes": ("mcp.read", "mcp.write")}), + _id_jag_config().model_copy(update={"org_token_endpoint": "https://idp.example.com/v2/token"}), + _id_jag_config().model_copy(update={"resource_token_endpoint": "https://mcp-as.example.com/v2/token"}), + _id_jag_config().model_copy(update={"client_id": "litellm-rotated"}), + _id_jag_config().model_copy(update={"client_auth": ClientSecretAuth(client_secret=SecretStr("rotated"))}), + _id_jag_config().model_copy(update={"subject_token_type": "urn:ietf:params:oauth:token-type:saml2"}), + ], + ids=[ + "audience", + "resource", + "scopes", + "org_token_endpoint", + "resource_token_endpoint", + "client_id", + "client_auth", + "subject_token_type", + ], +) +async def test_id_jag_config_change_forces_a_fresh_exchange(changed): + endpoint = _FakeTokenEndpoint(_two_leg_ok("old-policy-bearer") + _two_leg_ok("new-policy-bearer")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + + before = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(_id_jag_config())) + after = await provider.resolve_credentials(_with_inbound("user-id-token"), _spec(changed)) + + assert isinstance(before, Ok) and isinstance(after, Ok) + assert _emitted(after.ok)["Authorization"] == "Bearer new-policy-bearer" + assert len(endpoint.calls) == 4 + + +@pytest.mark.asyncio +async def test_id_jag_does_not_share_the_cached_bearer_across_caller_tokens(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("alice-bearer") + _two_leg_ok("bob-bearer")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + + alice = await provider.resolve_credentials(_with_inbound("alice-id-token"), _spec(_id_jag_config())) + bob = await provider.resolve_credentials(_with_inbound("bob-id-token"), _spec(_id_jag_config())) + + assert isinstance(alice, Ok) and isinstance(bob, Ok) + assert _emitted(bob.ok)["Authorization"] == "Bearer bob-bearer" + assert len(endpoint.calls) == 4 + + +@pytest.mark.asyncio +async def test_invalidate_credentials_evicts_the_id_jag_bearer_so_the_next_resolve_re_exchanges(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("rejected-bearer") + _two_leg_ok("fresh-bearer")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + subject = _with_inbound("user-id-token") + + first = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + await provider.invalidate_credentials(subject, _spec(_id_jag_config())) + second = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(second.ok)["Authorization"] == "Bearer fresh-bearer" + assert len(endpoint.calls) == 4 + + +@pytest.mark.asyncio +async def test_invalidate_credentials_for_id_jag_is_a_noop_without_a_caller_token(): + endpoint = _FakeTokenEndpoint(_two_leg_ok("cached-bearer")) + provider = UpstreamCredentialProvider(token_endpoint=endpoint) + subject = _with_inbound("user-id-token") + + first = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + await provider.invalidate_credentials(Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config())) + second = await provider.resolve_credentials(subject, _spec(_id_jag_config())) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert _emitted(second.ok)["Authorization"] == "Bearer cached-bearer" + assert len(endpoint.calls) == 2 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py new file mode 100644 index 00000000000..f100bd56f8f --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py @@ -0,0 +1,408 @@ +"""Tests for the v2 token-endpoint collaborator. + +`TokenEndpointClient.fetch` makes one authenticated POST and returns the minted token as a value; +`ExchangedTokenCache` memoizes it with per-key single-flight. These pin the grant/client-auth wire +shape, the private-key-JWT vs client_secret authentication, the error-as-value mapping, and the +cache's hit/single-flight behavior. Each assertion fails under a real mutation of the feature. +""" + +import asyncio +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import jwt +import litellm +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa + +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( + CLIENT_ASSERTION_TYPE, + ExchangedToken, + ExchangedTokenCache, + TokenEndpointClient, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + ClientSecretAuth, + CredError, + PrivateKeyJwtAuth, +) +from pydantic import SecretStr + +_PATCH_TARGET = ( + "litellm.proxy._experimental.mcp_server.outbound_credentials." + "token_endpoint.get_async_httpx_client" +) + +_ENDPOINT = "https://idp.example.com/oauth2/token" +_CLIENT_ID = "litellm-client-id" + +_RSA_KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048) +_PRIVATE_PEM = _RSA_KEY.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), +).decode() +_PUBLIC_PEM = ( + _RSA_KEY.public_key() + .public_bytes( + serialization.Encoding.PEM, + serialization.PublicFormat.SubjectPublicKeyInfo, + ) + .decode() +) + + +def _resp(token="access", expires_in=3600): + resp = MagicMock() + resp.json.return_value = {"access_token": token, "expires_in": expires_in} + resp.raise_for_status = MagicMock() + return resp + + +def _client(response): + client = AsyncMock() + client.post.return_value = response + return client + + +def _posted_data(client): + return client.post.call_args.kwargs["data"] + + +@pytest.mark.asyncio +async def test_fetch_forwards_grant_params_and_client_secret(): + client = _client(_resp("the-token", expires_in=1200)) + with patch(_PATCH_TARGET, return_value=client): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g", "subject_token": "user-jwt"}, + ClientSecretAuth(client_secret=SecretStr("shhh")), + ) + + assert isinstance(result, Ok) + assert result.ok == ExchangedToken(access_token="the-token", expires_in=1200) + assert client.post.call_args.args[0] == _ENDPOINT + data = _posted_data(client) + assert data["grant_type"] == "g" + assert data["subject_token"] == "user-jwt" + assert data["client_id"] == _CLIENT_ID + assert data["client_secret"] == "shhh" + assert "client_assertion" not in data + + +@pytest.mark.asyncio +async def test_fetch_private_key_jwt_client_assertion(): + client = _client(_resp()) + with patch(_PATCH_TARGET, return_value=client): + await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + PrivateKeyJwtAuth( + private_key=SecretStr(_PRIVATE_PEM), + key_id="kid-1", + signing_alg="RS256", + ), + ) + + data = _posted_data(client) + assert data["client_assertion_type"] == CLIENT_ASSERTION_TYPE + assert "client_secret" not in data + decoded = jwt.decode( + data["client_assertion"], + _PUBLIC_PEM, + algorithms=["RS256"], + audience=_ENDPOINT, + ) + assert decoded["iss"] == _CLIENT_ID + assert decoded["sub"] == _CLIENT_ID + assert decoded["aud"] == _ENDPOINT + assert "exp" in decoded + assert jwt.get_unverified_header(data["client_assertion"])["kid"] == "kid-1" + + +@pytest.mark.asyncio +async def test_fetch_http_error_maps_to_upstream_unavailable_with_status(): + error_resp = MagicMock() + error_resp.status_code = 403 + error_resp.raise_for_status.side_effect = httpx.HTTPStatusError( + "Forbidden", request=MagicMock(), response=error_resp + ) + with patch(_PATCH_TARGET, return_value=_client(error_resp)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert "403" in result.error.summary + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raised", + [ + httpx.ConnectError("connection refused", request=MagicMock()), + httpx.ReadTimeout("timed out", request=MagicMock()), + litellm.Timeout( + message="Connection timed out", + model="default-model-name", + llm_provider="litellm-httpx-handler", + ), + ], +) +async def test_fetch_network_error_maps_to_upstream_unavailable(raised): + client = AsyncMock() + client.post.side_effect = raised + with patch(_PATCH_TARGET, return_value=client): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth", + [ + PrivateKeyJwtAuth(private_key=SecretStr("not-a-pem-key"), signing_alg="RS256"), + PrivateKeyJwtAuth(private_key=SecretStr(_PRIVATE_PEM), signing_alg="XX999"), + ], + ids=["garbage-key", "unknown-alg"], +) +async def test_fetch_unsignable_client_assertion_is_misconfigured_not_a_crash(auth): + client = AsyncMock() + with patch(_PATCH_TARGET, return_value=client): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + auth, + ) + + assert isinstance(result, Error) + assert result.error.tag == "misconfigured" + client.post.assert_not_called() + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +@pytest.mark.asyncio +async def test_fetch_invalid_json_maps_to_upstream_unavailable(): + bad = MagicMock() + bad.raise_for_status = MagicMock() + bad.json.side_effect = json.JSONDecodeError("Expecting value", "", 0) + with patch(_PATCH_TARGET, return_value=_client(bad)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +@pytest.mark.asyncio +async def test_fetch_none_response_is_upstream_unavailable(): + with patch(_PATCH_TARGET, return_value=_client(None)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_fetch_missing_access_token_is_upstream_unavailable(): + bad = MagicMock() + bad.json.return_value = {"token_type": "Bearer"} + bad.raise_for_status = MagicMock() + with patch(_PATCH_TARGET, return_value=_client(bad)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert result.error.tag == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_fetch_http_error_does_not_leak_endpoint_url(): + error_resp = MagicMock() + error_resp.status_code = 403 + error_resp.raise_for_status.side_effect = httpx.HTTPStatusError( + "Forbidden", request=MagicMock(), response=error_resp + ) + with patch(_PATCH_TARGET, return_value=_client(error_resp)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +@pytest.mark.asyncio +async def test_fetch_none_response_does_not_leak_endpoint_url(): + with patch(_PATCH_TARGET, return_value=_client(None)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +@pytest.mark.asyncio +async def test_fetch_missing_access_token_does_not_leak_endpoint_url(): + bad = MagicMock() + bad.json.return_value = {"token_type": "Bearer"} + bad.raise_for_status = MagicMock() + with patch(_PATCH_TARGET, return_value=_client(bad)): + result = await TokenEndpointClient().fetch( + _ENDPOINT, + _CLIENT_ID, + {"grant_type": "g"}, + ClientSecretAuth(client_secret=SecretStr("s")), + ) + + assert isinstance(result, Error) + assert _ENDPOINT not in result.error.summary + assert "idp.example.com" not in result.error.summary + + +def _ok_token(value="cached") -> Result[ExchangedToken, CredError]: + return Ok(ExchangedToken(access_token=value, expires_in=3600)) + + +@pytest.mark.asyncio +async def test_cache_hit_skips_the_second_compute(): + cache = ExchangedTokenCache() + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + return _ok_token("tok") + + first = await cache.get_or_compute("k", compute) + second = await cache.get_or_compute("k", compute) + + assert isinstance(first, Ok) and first.ok == "tok" + assert isinstance(second, Ok) and second.ok == "tok" + assert calls == 1 + + +@pytest.mark.asyncio +async def test_cache_single_flights_concurrent_misses(): + cache = ExchangedTokenCache() + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + await asyncio.sleep(0.01) + return _ok_token("shared") + + results = await asyncio.gather( + cache.get_or_compute("k", compute), + cache.get_or_compute("k", compute), + ) + + assert [r.ok for r in results] == ["shared", "shared"] + assert calls == 1 + + +@pytest.mark.asyncio +async def test_cache_invalidate_forces_the_next_compute(): + cache = ExchangedTokenCache() + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + return _ok_token(f"tok-{calls}") + + first = await cache.get_or_compute("k", compute) + cache.invalidate("k") + second = await cache.get_or_compute("k", compute) + + assert isinstance(first, Ok) and first.ok == "tok-1" + assert isinstance(second, Ok) and second.ok == "tok-2" + assert calls == 2 + + +@pytest.mark.asyncio +async def test_cache_invalidate_only_evicts_the_named_key(): + cache = ExchangedTokenCache() + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + return _ok_token(f"tok-{calls}") + + await cache.get_or_compute("keep", compute) + await cache.get_or_compute("evict", compute) + cache.invalidate("evict") + kept = await cache.get_or_compute("keep", compute) + + assert isinstance(kept, Ok) and kept.ok == "tok-1" + assert calls == 2 + + +@pytest.mark.asyncio +async def test_cache_does_not_store_a_failed_compute(): + cache = ExchangedTokenCache() + calls = 0 + + async def compute(): + nonlocal calls + calls += 1 + if calls == 1: + return Error(CredError.of_upstream_unavailable("down")) + return _ok_token("recovered") + + first = await cache.get_or_compute("k", compute) + second = await cache.get_or_compute("k", compute) + + assert isinstance(first, Error) + assert isinstance(second, Ok) and second.ok == "recovered" + assert calls == 2 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py index 43b3612a5f2..bb25ab6bd3c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py @@ -17,10 +17,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( AuthSpecKind, AwsSigV4Config, Byok, + ClientSecretAuth, CredError, Error, + IdJagConfig, NoneConfig, Ok, + PrivateKeyJwtAuth, ServerSpec, SharedKey, StaticKeys, @@ -29,6 +32,14 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( _AUTH_CONFIG = TypeAdapter(AuthConfig) +_ID_JAG_MINIMAL = { + "kind": "id_jag", + "org_token_endpoint": "https://idp.example.com/token", + "resource_token_endpoint": "https://mcp-as.example.com/token", + "client_id": "litellm", + "client_auth": {"source": "client_secret", "client_secret": "s"}, +} + def test_parse_auth_spec_kind_accepts_known_mode(): result = parse_auth_spec_kind("token_exchange") @@ -148,3 +159,73 @@ def test_secrets_do_not_leak_in_repr(): key = SharedKey(value=SecretStr("super-secret")) assert "super-secret" not in repr(key) assert key.value.get_secret_value() == "super-secret" + + +@pytest.mark.parametrize( + "missing", + ["org_token_endpoint", "resource_token_endpoint", "client_id", "client_auth"], +) +def test_id_jag_config_requires_each_endpoint_client_and_auth(missing): + payload = {k: v for k, v in _ID_JAG_MINIMAL.items() if k != missing} + with pytest.raises(ValidationError): + _AUTH_CONFIG.validate_python(payload) + + +def test_id_jag_client_auth_discriminates_on_source(): + by_secret = _AUTH_CONFIG.validate_python(_ID_JAG_MINIMAL) + assert isinstance(by_secret, IdJagConfig) + assert isinstance(by_secret.client_auth, ClientSecretAuth) + assert by_secret.client_auth.client_secret.get_secret_value() == "s" + + by_key = _AUTH_CONFIG.validate_python( + { + **_ID_JAG_MINIMAL, + "client_auth": { + "source": "private_key_jwt", + "private_key": "PEM", + "key_id": "kid-1", + "signing_alg": "RS384", + }, + } + ) + assert isinstance(by_key, IdJagConfig) + assert isinstance(by_key.client_auth, PrivateKeyJwtAuth) + assert by_key.client_auth.private_key.get_secret_value() == "PEM" + assert by_key.client_auth.key_id == "kid-1" + assert by_key.client_auth.signing_alg == "RS384" + + +def test_id_jag_client_auth_rejects_unknown_source(): + with pytest.raises(ValidationError): + _AUTH_CONFIG.validate_python( + {**_ID_JAG_MINIMAL, "client_auth": {"source": "mystery"}} + ) + + +def test_id_jag_config_defaults_id_token_subject_and_empty_optionals(): + config = _AUTH_CONFIG.validate_python(_ID_JAG_MINIMAL) + assert isinstance(config, IdJagConfig) + assert config.subject_token_type == "urn:ietf:params:oauth:token-type:id_token" + assert config.audience is None + assert config.resource is None + assert config.scopes == () + + +def test_id_jag_secrets_do_not_leak_in_repr(): + config = IdJagConfig( + org_token_endpoint="https://idp.example.com/token", + resource_token_endpoint="https://mcp-as.example.com/token", + client_id="litellm", + client_auth=PrivateKeyJwtAuth(private_key=SecretStr("super-secret-pem")), + ) + assert "super-secret-pem" not in repr(config) + + +def test_id_jag_server_spec_derives_auth_spec_kind(): + config = _AUTH_CONFIG.validate_python(_ID_JAG_MINIMAL) + spec = ServerSpec( + server_id="s", + resource="https://mcp.example.com/mcp", + config=config, + ) + assert spec.auth_spec_kind is AuthSpecKind.id_jag diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 7269774442b..56ca855c814 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -19,6 +19,8 @@ import pytest from litellm.proxy._experimental.mcp_server.db import ( _decode_user_credential, _prepare_mcp_server_data, + decrypt_credentials, + encrypt_credentials, get_user_credential, get_user_oauth_credential, is_oauth_credential_expired, @@ -332,6 +334,29 @@ def _stored_value(prisma) -> str: return create_value +# ── MCP server credentials at rest ────────────────────────────────────────────── + + +def test_client_private_key_encrypted_at_rest(): + """An ID-JAG client_private_key is a secret and must be encrypted in the stored + credentials blob, never persisted in plaintext, and must round-trip back. The + pre-fix code left client_private_key out of encrypt_credentials, so it was stored + verbatim.""" + private_key = ( + "-----BEGIN PRIVATE KEY-----\nsensitive-rsa-material\n-----END PRIVATE KEY-----" + ) + credentials = {"client_secret": "shh", "client_private_key": private_key} + + encrypted = encrypt_credentials(dict(credentials), encryption_key=None) + assert encrypted["client_private_key"] != private_key + assert private_key not in encrypted["client_private_key"] + assert encrypted["client_secret"] != "shh" + + decrypted = decrypt_credentials(dict(encrypted)) + assert decrypted["client_private_key"] == private_key + assert decrypted["client_secret"] == "shh" + + # ── BYOK round-trip ─────────────────────────────────────────────────────────── @@ -978,3 +1003,67 @@ def test_prepare_mcp_server_data_update_carries_token_exchange_columns(): assert data["audience"] == "https://upstream.example.com" assert data["subject_token_type"] == "urn:ietf:params:oauth:token-type:jwt" assert data["token_exchange_profile"] == "entra_obo" + + +@pytest.mark.asyncio +async def test_master_key_rotation_reencrypts_oauth_client_store(monkeypatch): + """The server-scoped DCR client store (LiteLLM_MCPServerOAuthClient) is encrypted at rest, so a + master-key rotation must re-encrypt it alongside the server rows. Skipping it leaves + config-declared DCR clients under the retired key, where they decrypt back to ciphertext and + force a full re-authorization.""" + import litellm.proxy.common_utils.encrypt_decrypt_utils as enc + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.proxy._experimental.mcp_server.db import ( + decrypt_credentials, + encrypt_credentials, + rotate_mcp_server_credentials_master_key, + ) + + key_old, key_new = "salt-old-key", "salt-new-key" + + blob_old = safe_dumps( + encrypt_credentials( + credentials={"client_id": "cid-123", "client_secret": "sec-456"}, + encryption_key=key_old, + ) + ) + + monkeypatch.setattr(enc, "_get_salt_key", lambda: key_old) + + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_mcpserveroauthclient.find_many = AsyncMock( + return_value=[SimpleNamespace(server_id="config_faros", credentials=blob_old)] + ) + store_update = AsyncMock() + prisma.db.litellm_mcpserveroauthclient.update = store_update + + await rotate_mcp_server_credentials_master_key(prisma, touched_by="test", new_master_key=key_new) + + store_update.assert_awaited_once() + assert store_update.await_args.kwargs["where"] == {"server_id": "config_faros"} + rotated_blob = store_update.await_args.kwargs["data"]["credentials"] + + monkeypatch.setattr(enc, "_get_salt_key", lambda: key_new) + recovered = decrypt_credentials(credentials=json.loads(rotated_blob)) + assert recovered["client_id"] == "cid-123" + assert recovered["client_secret"] == "sec-456" + + +@pytest.mark.asyncio +async def test_delete_mcp_server_cleans_oauth_client_store(): + """Deleting a server must remove its server-scoped DCR client store entry alongside the per-user + credential and env-var rows, or a re-created server reusing the same server_id would inherit the + deleted server's OAuth client.""" + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + prisma = MagicMock() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=SimpleNamespace(server_id="s1")) + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock() + prisma.db.litellm_mcpuserenvvars.delete_many = AsyncMock() + prisma.db.litellm_mcpserveroauthclient.delete_many = AsyncMock() + + await delete_mcp_server(prisma, "s1", invalidate_token_cache=AsyncMock()) + + prisma.db.litellm_mcpserveroauthclient.delete_many.assert_awaited_once_with(where={"server_id": "s1"}) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index f5ac229d119..6f2f24df8fa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7130,3 +7130,373 @@ async def test_token_exchange_unreadable_body_still_renders_oauth_fault(): assert response.status_code == 502 body = json.loads(response.body) assert body == {"error": "server_error", "error_description": "upstream token endpoint returned HTTP 400"} + + +@pytest.mark.asyncio +async def test_persist_dcr_client_for_config_server_uses_side_store(): + """A config.yaml-declared OAuth2 DCR server has no LiteLLM_MCPServerTable row, so + update_mcp_server returns None. The minted client must then persist to the server-scoped + OAuth-client store keyed by server_id (never a shadow server row), overlay onto the in-memory + server so refresh can authenticate this process, and never call update_server(None) (which + previously raised AttributeError on .approval_status, was swallowed, and reported a 200 that + persisted nothing).""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _persist_dcr_client_registration, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + config_server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + client_secret=None, + authorization_url="https://provider.example/oauth/authorize", + token_url="https://provider.example/oauth/token", + registration_url="https://provider.example/oauth/register", + ) + + mock_upsert = AsyncMock() + mock_update_server = AsyncMock() + + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.update_mcp_server", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", + new=mock_upsert, + ), + patch.object(global_mcp_server_manager, "update_server", new=mock_update_server), + ): + result = await _persist_dcr_client_registration( + mcp_server=config_server, + registration_response={ + "client_id": "minted-client", + "client_secret": "minted-secret", + "token_endpoint_auth_method": "client_secret_basic", + }, + current_redirect_uri="https://proxy.litellm.example/callback", + ) + + assert result == "persisted" + + mock_upsert.assert_called_once() + assert mock_upsert.call_args.kwargs["server_id"] == "config_faros" + stored = mock_upsert.call_args.kwargs["credentials"] + assert stored["client_id"] == "minted-client" + assert stored["client_secret"] == "minted-secret" + assert stored["token_endpoint_auth_method"] == "client_secret_basic" + assert stored["redirect_uris"] == ["https://proxy.litellm.example/callback"] + + assert config_server.client_id == "minted-client" + assert config_server.client_secret == "minted-secret" + assert config_server.token_endpoint_auth_method == "client_secret_basic" + + mock_update_server.assert_not_called() + + +@pytest.mark.asyncio +async def test_hydrate_config_server_applies_stored_dcr_client(monkeypatch): + """On restart a config server's in-memory object has no client_id; hydration overlays the + persisted DCR client from the server-scoped store, decrypting the encrypted-at-rest blob, so the + refresh_token grant can authenticate as the registered client instead of re-authenticating.""" + import litellm.proxy.common_utils.encrypt_decrypt_utils as enc + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.proxy._experimental.mcp_server.db import encrypt_credentials + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + hydrate_config_server_dcr_client, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-hydrate-key") + stored_blob = safe_dumps( + encrypt_credentials( + credentials={ + "client_id": "stored-client", + "client_secret": "stored-secret", + "token_endpoint_auth_method": "client_secret_basic", + "redirect_uris": ["https://proxy.litellm.example/callback"], + }, + encryption_key="salt-hydrate-key", + ) + ) + assert "stored-client" not in stored_blob and "stored-secret" not in stored_blob + + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=AsyncMock(return_value=stored_blob), + ), + ): + applied = await hydrate_config_server_dcr_client(server) + + assert applied is True + assert server.client_id == "stored-client" + assert server.client_secret == "stored-secret" + assert server.token_endpoint_auth_method == "client_secret_basic" + + +@pytest.mark.asyncio +async def test_reuse_config_server_reads_store_with_real_crypto(monkeypatch): + """A config-declared server (rowless) keeps its DCR client in the store, so the reuse read + resolves it from the store and decrypts the encrypted-at-rest client, mirroring the write path so + a re-authorize reuses the client instead of re-minting one.""" + import litellm.proxy.common_utils.encrypt_decrypt_utils as enc + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.proxy._experimental.mcp_server.db import encrypt_credentials + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_faros", + name="config_faros", + server_name="config_faros", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + monkeypatch.setattr(enc, "_get_salt_key", lambda: "salt-reuse-key") + blob = safe_dumps( + encrypt_credentials( + credentials={"client_id": "stored-client", "client_secret": "sec", "redirect_uris": ["https://x/callback"]}, + encryption_key="salt-reuse-key", + ) + ) + assert "stored-client" not in blob + store_lookup = AsyncMock(return_value=blob) + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=True), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_lookup, + ), + ): + result = await _reuse_persisted_dcr_client_if_available(server, current_redirect_uri="https://x/callback") + + assert result is True + assert server.client_id == "stored-client" + store_lookup.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_temp_server_is_not_persisted_to_store(): + """A rowless server that is NOT config-declared (a throwaway /server/oauth/session server) must + not leave a permanent store row on persist, and the read must never consult the store for it. Its + minted client is overlaid in memory for the session only.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _persist_dcr_client_registration, + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + temp = MCPServer( + server_id="temp-uuid", + name="temp", + server_name="temp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + authorization_url="https://p.example/authorize", + token_url="https://p.example/token", + registration_url="https://p.example/register", + ) + + upsert = AsyncMock() + store_read = AsyncMock(return_value=None) + with ( + patch.object(global_mcp_server_manager, "is_config_declared_server", return_value=False), + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch("litellm.proxy._experimental.mcp_server.db.update_mcp_server", new=AsyncMock(return_value=None)), + patch("litellm.proxy._experimental.mcp_server.db.get_mcp_server", new=AsyncMock(return_value=None)), + patch("litellm.proxy._experimental.mcp_server.db.upsert_mcp_server_oauth_client_credentials", new=upsert), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_read, + ), + patch.object(global_mcp_server_manager, "update_server", new=AsyncMock()), + ): + result = await _persist_dcr_client_registration( + temp, {"client_id": "temp-client", "client_secret": "s"}, "https://x/callback" + ) + reused = await _reuse_persisted_dcr_client_if_available( + MCPServer( + server_id="temp-uuid", + name="temp", + server_name="temp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ), + current_redirect_uri="https://x/callback", + ) + + assert result == "persisted" + assert temp.client_id == "temp-client" + upsert.assert_not_called() + store_read.assert_not_called() + assert reused is False + + +@pytest.mark.asyncio +async def test_hydrate_does_not_overwrite_explicit_config_client_id(): + """An explicit client_id set in config.yaml wins: hydration must not overwrite it with a stale + persisted store client, and must not even read the store when config already supplied a client.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + hydrate_config_server_dcr_client, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="config_static", + name="config_static", + server_name="config_static", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="explicit-from-config", + ) + store_read = AsyncMock( + return_value={"client_id": "stale-store-client", "client_secret": "x", "redirect_uris": []} + ) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_read, + ), + ): + applied = await hydrate_config_server_dcr_client(server) + + assert applied is False + assert server.client_id == "explicit-from-config" + store_read.assert_not_called() + + +@pytest.mark.asyncio +async def test_reuse_does_not_inherit_store_client_when_a_row_exists(): + """Security: a server that HAS a LiteLLM_MCPServerTable row reads its DCR client only from that + row, never from the server-scoped store. server_id is caller-settable on create, so a submitted + server whose id collides with a config-declared server must not be able to load that config + server's client from the store and send it to its own token endpoint. A row that exists but has + no client_id yields no reusable client and must not fall back to the store.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _reuse_persisted_dcr_client_if_available, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + submitted = MCPServer( + server_id="collides_with_config", + name="submitted", + server_name="submitted", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id=None, + ) + + row_without_client = MagicMock() + row_without_client.credentials = None + row_without_client.server_id = "collides_with_config" + store_lookup = AsyncMock( + return_value={"client_id": "config-secret-client", "client_secret": "leak", "redirect_uris": []} + ) + with ( + patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server", + new=AsyncMock(return_value=row_without_client), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.get_mcp_server_oauth_client_credentials", + new=store_lookup, + ), + ): + result = await _reuse_persisted_dcr_client_if_available(submitted, current_redirect_uri="https://x/callback") + + assert result is False + assert submitted.client_id is None + store_lookup.assert_not_called() + + +@pytest.mark.asyncio +async def test_load_servers_from_config_hydrates_dcr_clients(): + """load_servers_from_config must invoke DCR-client hydration so config servers pick up their + persisted client on startup; deleting the call site leaves a restarted server with no client_id + and forces re-authentication on every token expiry.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + hydrate_spy = AsyncMock() + with patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy): + await global_mcp_server_manager.load_servers_from_config({}) + + hydrate_spy.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_reload_servers_from_database_hydrates_dcr_clients(): + """load_servers_from_config runs before the DB connects at startup, so its hydration no-ops; + reload_servers_from_database runs after the DB connects and must hydrate config servers' persisted + DCR clients too, or a fresh pod has no client_id for a config server and forces re-authentication + on the first token refresh.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + + hydrate_spy = AsyncMock() + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma, + ), + patch.object(global_mcp_server_manager, "_hydrate_config_servers_dcr_clients", new=hydrate_spy), + ): + await global_mcp_server_manager.reload_servers_from_database() + + hydrate_spy.assert_awaited_once() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 55b6bbbdbc2..491fa023031 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -176,6 +176,89 @@ class TestMCPServerManager: assert calls == [("", "authz-srv")] assert client is not None + @pytest.mark.asyncio + async def test_caller_auth_header_cannot_bypass_id_jag_exchange(self): + """A caller-supplied per-request override must not disable the ID-JAG exchange and forward an + arbitrary bearer upstream: _create_mcp_client keeps the v2 spec and resolves through the + injected provider rather than deferring to the v1 caller-override path.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Ok, + ) + from litellm.types.mcp import MCPAuth + + calls = [] + + class _FakeProvider: + async def resolve_credentials(self, subject, server): + calls.append((subject.subject_id, server.server_id)) + return Ok(StaticHeaderAuth("Bearer minted-id-jag-token")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = MCPServer( + server_id="id-jag-srv", + name="id-jag", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.example/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token", + ) + + client = await manager._create_mcp_client( + server, + mcp_auth_header="Bearer caller-supplied-token", + subject_token="caller-id-token", + ) + + assert calls == [("", "id-jag-srv")] + assert client is not None + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "missing_field", + ["token_exchange_endpoint", "id_jag_resource_token_endpoint", "client_id", "client_secret"], + ) + async def test_half_configured_id_jag_fails_closed_instead_of_deferring_to_v1(self, missing_field): + """ID-JAG has no v1 arm, so a half-configured oauth2_id_jag server must not silently fall + through to resolve_mcp_auth, where a caller x-mcp-* override or the static + authentication_token would bypass the per-user identity assertion. It must be refused as an + operator misconfiguration (HTTP 500) before any client is built.""" + from fastapi import HTTPException + + from litellm.types.mcp import MCPAuth + + fields = { + "client_id": "gateway-client", + "client_secret": "gateway-secret", + "token_exchange_endpoint": "https://org-idp.example/oauth2/token", + "id_jag_resource_token_endpoint": "https://resource-as.example/oauth2/token", + } + fields.pop(missing_field) + server = MCPServer( + server_id="id-jag-srv", + name="id-jag", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.oauth2_id_jag, + authentication_token="static-server-secret", + **fields, + ) + + with pytest.raises(HTTPException) as exc_info: + await MCPServerManager()._create_mcp_client( + server, + mcp_auth_header="Bearer caller-supplied-token", + subject_token="caller-id-token", + ) + + assert exc_info.value.status_code == 500 + assert "oauth2_id_jag" in str(exc_info.value.detail) + async def test_create_mcp_client_stdio_injects_npm_config_cache(self): """Test that _create_mcp_client injects NPM_CONFIG_CACHE when not already set, and preserves user-provided NPM_CONFIG_CACHE when present.""" @@ -257,6 +340,35 @@ class TestMCPServerManager: assert env == {} @pytest.mark.asyncio + async def test_load_servers_from_config_debug_dump_redacts_secrets(self, caplog): + """The registry debug dump must not leak long-lived credentials: the ID-JAG signing key, + client secret, and static token are masked while non-secret fields stay readable.""" + + manager = MCPServerManager() + config = { + "idjag": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2_id_jag, + "client_id": "gateway-client", + "client_secret": "SECRET-CLIENT-SECRET", + "client_private_key": "-----BEGIN PRIVATE KEY-----SECRET-PEM-----END PRIVATE KEY-----", + "token_exchange_endpoint": "https://org-idp.example/oauth2/token", + "id_jag_resource_token_endpoint": "https://resource-as.example/oauth2/token", + "authentication_token": "SECRET-STATIC-TOKEN", + } + } + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + await manager.load_servers_from_config(config) + + dump = next(m for m in caplog.messages if "Loaded MCP Servers" in m) + assert "SECRET-PEM" not in dump + assert "SECRET-CLIENT-SECRET" not in dump + assert "SECRET-STATIC-TOKEN" not in dump + assert "gateway-client" in dump + assert "https://org-idp.example/oauth2/token" in dump + async def test_load_servers_from_config_warns_on_invalid_alias(self, caplog): """Invalid aliases from config should emit warnings during load.""" @@ -8122,6 +8234,86 @@ class TestOBOCallToolRetry: manager._create_mcp_client.assert_awaited_once() assert first.attempts == 1 and retry.attempts == 1 + @pytest.mark.asyncio + async def test_upstream_401_on_id_jag_evicts_the_cached_bearer_and_retries(self): + """The retry path must invalidate the ID-JAG leg-2 bearer too: without eviction the rebuilt + client resolves the same rejected token from the cache and the retry 401s identically.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( + IdJagConfig, + ) + + manager = self._manager() + success = CallToolResult(content=[], isError=False) + first = _RetryFakeClient(raises=_UpstreamAuthError(401)) + retry = _RetryFakeClient(result=success) + manager._create_mcp_client = AsyncMock(return_value=retry) + server = MCPServer( + server_id="id-jag-srv", + name="id-jag", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.example/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token", + ) + + result = await manager._obo_call_tool_with_retry( + client=first, + call_tool_params=MagicMock(), + host_progress_callback=None, + mcp_server=server, + server_auth_header=None, + extra_headers=None, + stdio_env=None, + subject_token="caller-id-token", + user_api_key_auth=None, + ) + + assert result is success + manager._cred_provider.invalidate_credentials.assert_awaited_once() + invalidated_spec = manager._cred_provider.invalidate_credentials.await_args.args[1] + assert isinstance(invalidated_spec.config, IdJagConfig) + assert first.attempts == 1 and retry.attempts == 1 + + @pytest.mark.asyncio + async def test_call_regular_routes_id_jag_through_the_retry_path(self): + """An oauth2_id_jag tool call with a subject token must take the invalidate-and-retry branch + of _call_regular_mcp_tool, not the plain single call, so an upstream 401 re-exchanges.""" + manager = self._manager() + success = CallToolResult(content=[], isError=False) + first = _RetryFakeClient(raises=_UpstreamAuthError(401)) + retry = _RetryFakeClient(result=success) + manager._create_mcp_client = AsyncMock(side_effect=[first, retry]) + server = MCPServer( + server_id="id-jag-srv", + name="id-jag", + url="https://upstream.example/mcp", + transport=MCPTransport.sse, + auth_type=MCPAuth.oauth2_id_jag, + client_id="gateway-client", + client_secret="gateway-secret", + token_exchange_endpoint="https://org-idp.example/oauth2/token", + id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token", + ) + + result = await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="tool", + arguments={}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer caller-id-token"}, + raw_headers=None, + proxy_logging_obj=None, + ) + + assert result is success + manager._cred_provider.invalidate_credentials.assert_awaited_once() + assert first.attempts == 1 and retry.attempts == 1 + @pytest.mark.asyncio async def test_non_auth_error_does_not_retry(self): manager = self._manager() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py index 5992fd1814f..f6b61c1d9f7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -989,6 +989,7 @@ class TestRotateCredentials: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[server]) mock_prisma.db.litellm_mcpservertable.update = AsyncMock() + mock_prisma.db.litellm_mcpserveroauthclient.find_many = AsyncMock(return_value=[]) with ( patch( @@ -1036,6 +1037,7 @@ class TestRotateCredentials: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[server]) mock_prisma.db.litellm_mcpservertable.update = AsyncMock() + mock_prisma.db.litellm_mcpserveroauthclient.find_many = AsyncMock(return_value=[]) with ( patch( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py new file mode 100644 index 00000000000..ca57118ee9d --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py @@ -0,0 +1,733 @@ +import json +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from litellm.exceptions import GuardrailRaisedException, ModifyResponseException +from litellm.proxy.guardrails.guardrail_hooks.straiker import initialize_guardrail +from litellm.proxy.guardrails.guardrail_hooks.straiker.straiker import ( + StraikerGuardrail, +) +from litellm.proxy.guardrails.guardrail_registry import ( + guardrail_class_registry, + guardrail_initializer_registry, +) +from litellm.types.proxy.guardrails.guardrail_hooks.straiker import ( + StraikerGuardrailConfigModel, + StraikerGuardrailConfigModelOptionalParams, +) +from litellm.types.utils import Choices, Message, ModelResponse, Usage + + +def _mock_response(action: str, turn_id: str = "turn-1", schema_version: str = "1", **extra) -> MagicMock: + resp = MagicMock(spec=httpx.Response) + resp.status_code = 200 + resp.json.return_value = { + "schema_version": schema_version, + "action": action, + "turn_id": turn_id, + **extra, + } + resp.text = "" + return resp + + +def _make_guardrail(**overrides) -> StraikerGuardrail: + defaults = dict( + api_key="test-key", + api_base="https://test.straiker.ai", + max_retries=0, + guardrail_name="straiker", + event_hook="pre_call", + async_handler=MagicMock(spec=httpx.AsyncClient), + ) + defaults.update(overrides) + g = StraikerGuardrail(**defaults) + g.async_handler.post = AsyncMock() + return g + + +def _logging_obj() -> MagicMock: + obj = MagicMock() + obj.litellm_call_id = "call-123" + obj.litellm_trace_id = "trace-456" + obj.call_type = "acompletion" + return obj + + +def _posted_payload(g: StraikerGuardrail) -> dict: + return json.loads(g.async_handler.post.call_args.kwargs["content"]) + + +def test_registry_membership(): + assert "straiker" in guardrail_initializer_registry + assert guardrail_class_registry["straiker"] is StraikerGuardrail + + +def test_config_model_wiring(): + assert StraikerGuardrailConfigModel.ui_friendly_name() == "Straiker" + assert StraikerGuardrail.get_config_model() is StraikerGuardrailConfigModel + fields = StraikerGuardrailConfigModel.model_fields + assert "api_key" in fields + assert "api_base" in fields + assert "default_app" in fields + assert "source" not in fields + assert "optional_params" in fields + assert "timeout" not in fields + assert "verbose" not in fields + + +def test_init_rejects_empty_api_key(): + with pytest.raises(ValueError): + StraikerGuardrail(api_key="") + + +def test_init_rejects_invalid_fallback(): + with pytest.raises(ValueError): + StraikerGuardrail(api_key="k", unreachable_fallback="nope") + + +def test_supported_hooks_limited_to_pre_and_post(): + from litellm.types.guardrails import GuardrailEventHooks + + assert StraikerGuardrail.get_supported_event_hooks() == [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ] + + +def test_during_call_mode_rejected_at_init(): + with pytest.raises(ValueError): + StraikerGuardrail(api_key="k", event_hook="during_call") + + +def test_streaming_attrs_hardcoded_to_buffered(): + g = _make_guardrail() + assert g.streaming_buffer_until_moderated is True + assert g.streaming_end_of_stream_only is True + + +def test_streaming_flags_not_configurable(): + fields = StraikerGuardrailConfigModelOptionalParams.model_fields + assert "streaming_buffer_until_moderated" not in fields + assert "streaming_end_of_stream_only" not in fields + assert "streaming_sampling_rate" not in fields + + +def test_initializer_builds_working_callback(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams(guardrail="straiker", mode="pre_call", api_key="abc", api_base="https://x.straiker.ai") + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.api_base == "https://x.straiker.ai" + + +def test_initializer_maps_default_app_to_source(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams( + guardrail="straiker", + mode="pre_call", + api_key="abc", + default_app="My App", + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert callback.source == "My App" + + +def test_initializer_reads_optional_params_flattened_like_ui(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams( + guardrail="straiker", + mode="pre_call", + api_key="abc", + api_base="https://x.straiker.ai", + timeout=9.5, + verbose=True, + unreachable_fallback="fail_open", + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.timeout == 9.5 + assert callback.verbose is True + assert callback.unreachable_fallback == "fail_open" + assert callback.api_base == "https://x.straiker.ai" + + +def test_initializer_reads_nested_optional_params(): + from types import SimpleNamespace + + from litellm.types.guardrails import LitellmParams + + params = LitellmParams.model_construct( + guardrail="straiker", + mode="pre_call", + api_key="abc", + api_base="https://x.straiker.ai", + optional_params=SimpleNamespace( + timeout=7.25, + verbose=True, + unreachable_fallback="fail_open", + ), + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.timeout == 7.25 + assert callback.verbose is True + assert callback.unreachable_fallback == "fail_open" + + +def test_initializer_reads_dict_optional_params(): + from litellm.types.guardrails import LitellmParams + + params = LitellmParams.model_construct( + guardrail="straiker", + mode="pre_call", + api_key="abc", + api_base="https://x.straiker.ai", + optional_params={"timeout": 7.25, "verbose": True, "unreachable_fallback": "fail_open"}, + ) + callback = initialize_guardrail(params, {"guardrail_name": "straiker"}) + assert isinstance(callback, StraikerGuardrail) + assert callback.timeout == 7.25 + assert callback.verbose is True + assert callback.unreachable_fallback == "fail_open" + + +@pytest.mark.asyncio +async def test_request_envelope_transport_and_shape(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + inputs = {"texts": ["hello world"], "model": "gpt-4o-mini"} + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello world"}], + "metadata": {"user_api_key_alias": "team-key", "agent_id": "chatbot-app", "app_name": "Chatbot"}, + } + + out = await g.apply_guardrail(inputs=inputs, request_data=request_data, input_type="request", logging_obj=_logging_obj()) + + assert out is inputs + url = g.async_handler.post.call_args.args[0] + assert url == "https://test.straiker.ai/api/v1/detect/webhook" + headers = g.async_handler.post.call_args.kwargs["headers"] + assert headers["X-Straiker-Webhook-Format"] == "litellm" + assert headers["Authorization"] == "Bearer test-key" + + payload = _posted_payload(g) + assert payload["schema_version"] == "1" + assert payload["event"]["type"] == "pre_call" + assert payload["event"]["id"] == "call-123:request" + assert payload["request"]["texts"] == ["hello world"] + assert payload["context"]["litellm_call_id"] == "call-123" + assert payload["identity"]["litellm_key"] == "team-key" + assert payload["application"] == {"source": "chatbot-app", "name": "Chatbot"} + assert "session_id" not in payload["application"] + assert "user_name" not in payload["application"] + assert "user_role" not in payload["application"] + assert "response" not in payload + assert "metadata" not in payload + + +@pytest.mark.asyncio +async def test_webhook_metadata_session_id_and_opaque_passthrough(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "litellm_session_id": "sess-from-litellm", + "metadata": { + "agent_id": "chatbot-app", + "app_name": "Chatbot", + "user_api_key_alias": "team-key", + "custom_tag": "experiment-7", + "client_ip": "10.0.0.1", + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + payload = _posted_payload(g) + assert payload["application"] == {"source": "chatbot-app", "name": "Chatbot"} + assert payload["identity"]["litellm_key"] == "team-key" + assert payload["context"]["session_id"] == "sess-from-litellm" + assert "session_id" not in payload["metadata"] + assert payload["metadata"] == { + "custom_tag": "experiment-7", + "client_ip": "10.0.0.1", + } + + +@pytest.mark.asyncio +async def test_webhook_metadata_never_forwards_proxy_internal_keys(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "custom_tag": "experiment-7", + "user_api_key": "sk-hashed-secret", + "user_api_end_user_max_budget": 12.5, + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["metadata"] == {"custom_tag": "experiment-7"} + + +@pytest.mark.asyncio +async def test_default_metadata_injected_and_config_wins_on_clash(): + g = _make_guardrail(metadata={"tenant": "acme", "custom_tag": "config-value"}) + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": {"custom_tag": "request-value", "client_ip": "10.0.0.1"}, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["metadata"] == { + "client_ip": "10.0.0.1", + "custom_tag": "config-value", + "tenant": "acme", + } + + +@pytest.mark.asyncio +async def test_default_metadata_present_without_request_metadata(): + g = _make_guardrail(metadata={"tenant": "acme"}) + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m"}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["metadata"] == {"tenant": "acme"} + + +@pytest.mark.asyncio +async def test_context_session_id_from_request_metadata(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m", "metadata": {"session_id": "sess-meta"}}, + input_type="request", + logging_obj=_logging_obj(), + ) + payload = _posted_payload(g) + assert payload["context"]["session_id"] == "sess-meta" + assert "metadata" not in payload + + +@pytest.mark.asyncio +async def test_identity_key_and_team_coalesce_alias_over_id(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "user_api_key_alias": "prod-key", + "user_api_key_hash": "hash-abc", + "user_api_key_team_alias": "growth", + "user_api_key_team_id": "team-9", + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + identity = _posted_payload(g)["identity"] + assert identity["litellm_key"] == "prod-key" + assert identity["litellm_team"] == "growth" + assert "key" not in identity + assert "team" not in identity + + +@pytest.mark.asyncio +async def test_identity_key_and_team_fall_back_to_hash_and_id(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "user_api_key_hash": "hash-abc", + "user_api_key_team_id": "team-9", + }, + }, + input_type="request", + logging_obj=_logging_obj(), + ) + identity = _posted_payload(g)["identity"] + assert identity["litellm_key"] == "hash-abc" + assert identity["litellm_team"] == "team-9" + + +@pytest.mark.asyncio +async def test_identity_end_user_from_resolved_metadata(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={ + "model": "m", + "metadata": { + "user_api_key_end_user_id": "eu-meta", + "user_api_key_user_id": "default_user_id", + }, + "user": "eu-body", + }, + input_type="request", + logging_obj=_logging_obj(), + ) + identity = _posted_payload(g)["identity"] + assert identity["end_user_id"] == "eu-meta" + assert identity["litellm_user_id"] == "default_user_id" + assert _posted_payload(g)["application"] == {"source": g.source} + + +@pytest.mark.asyncio +async def test_identity_end_user_absent_without_resolved_metadata(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m", "user": "eu-body", "metadata": {"user_api_key_user_id": "default_user_id"}}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert "end_user_id" not in _posted_payload(g)["identity"] + + +@pytest.mark.asyncio +async def test_application_source_from_agent_id(): + g = _make_guardrail(source="litellm") + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["x"]}, + request_data={"model": "m", "metadata": {"agent_id": "analytics-app", "app_name": "Analytics"}}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert _posted_payload(g)["application"] == {"source": "analytics-app", "name": "Analytics"} + +@pytest.mark.asyncio +async def test_request_block_raises_guardrail_exception_with_reason(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("BLOCKED", blocked_reason="prompt injection") + with pytest.raises(GuardrailRaisedException) as exc: + await g.apply_guardrail( + inputs={"texts": ["attack"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert "prompt injection" in str(exc.value) + + +@pytest.mark.asyncio +async def test_guardrail_intervened_writes_back_modified_text_only(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + inputs = {"texts": ["my ssn is 123"], "images": ["img-a"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out["texts"] == ["[redacted]"] + assert out["images"] == ["img-a"] + + +@pytest.mark.asyncio +async def test_streamed_response_intervention_converts_to_block(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "p"}], + "stream": True, + "response": response, + } + with pytest.raises(ModifyResponseException): + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + + +@pytest.mark.asyncio +async def test_non_streamed_response_intervention_redacts(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "p"}], + "response": response, + } + out = await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + assert out["texts"] == ["[redacted]"] + + +@pytest.mark.asyncio +async def test_guardrail_intervened_without_texts_blocks(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED") + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["my ssn is 123"]}, + request_data={"model": "m"}, + input_type="request", + logging_obj=_logging_obj(), + ) + + +@pytest.mark.asyncio +async def test_streamed_via_proxy_server_request_body_converts_to_block(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("GUARDRAIL_INTERVENED", texts=["[redacted]"]) + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "proxy_server_request": {"body": {"stream": True}}, + "response": response, + } + with pytest.raises(ModifyResponseException): + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + + +@pytest.mark.asyncio +async def test_response_envelope_and_block_replaces_response(): + g = _make_guardrail(verbose=True) + g.async_handler.post.return_value = _mock_response("BLOCKED") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "original prompt"}], + "stream": True, + "response": response, + } + with pytest.raises(ModifyResponseException) as exc: + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + + assert exc.value.original_response is response + payload = _posted_payload(g) + assert payload["event"]["type"] == "post_call" + assert payload["event"]["stream"]["phase"] == "assembled" + assert payload["response"]["texts"] == ["secret"] + assert payload["response"]["finish_reason"] == "stop" + assert payload["request"]["structured_messages"] == [{"role": "user", "content": "original prompt"}] + + +@pytest.mark.asyncio +async def test_post_call_fail_closed_raises_modify_response_exception(): + g = _make_guardrail(unreachable_fallback="fail_closed") + g.async_handler.post.side_effect = httpx.ConnectError("boom") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="secret", role="assistant"))], + model="gpt-4o-mini", + ) + request_data = {"model": "gpt-4o-mini", "response": response} + with pytest.raises(ModifyResponseException) as exc: + await g.apply_guardrail( + inputs={"texts": ["secret"], "model": "gpt-4o-mini"}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj(), + ) + assert exc.value.original_response is response + + +@pytest.mark.asyncio +async def test_usage_tokens_on_post_call(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="hi", role="assistant"))], + model="gpt-4o-mini", + usage=Usage(prompt_tokens=11, completion_tokens=7, total_tokens=18), + ) + await g.apply_guardrail( + inputs={"texts": ["hi"], "model": "gpt-4o-mini"}, + request_data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hey"}], "response": response}, + input_type="response", + logging_obj=_logging_obj(), + ) + usage = _posted_payload(g)["usage"] + assert usage == {"input_tokens": 11, "output_tokens": 7} + + +@pytest.mark.asyncio +async def test_usage_absent_on_pre_call(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + await g.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={"model": "gpt-4o-mini"}, + input_type="request", + logging_obj=_logging_obj(), + ) + assert "usage" not in _posted_payload(g) + + +@pytest.mark.asyncio +async def test_allow_returns_inputs_unchanged(): + g = _make_guardrail() + g.async_handler.post.return_value = _mock_response("NONE") + inputs = {"texts": ["fine"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_unreachable_fail_closed_blocks(): + g = _make_guardrail(unreachable_fallback="fail_closed") + g.async_handler.post.side_effect = httpx.ConnectError("boom") + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + + +@pytest.mark.asyncio +async def test_unreachable_fail_open_passes_through(): + g = _make_guardrail(unreachable_fallback="fail_open") + g.async_handler.post.side_effect = httpx.ConnectError("boom") + inputs = {"texts": ["x"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_fail_on_error_false_allows_on_bad_status(): + g = _make_guardrail(unreachable_fallback="fail_closed", fail_on_error=False) + bad = MagicMock(spec=httpx.Response) + bad.status_code = 400 + bad.text = "bad request" + g.async_handler.post.return_value = bad + inputs = {"texts": ["x"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_non_retryable_status_fail_closed_blocks(): + g = _make_guardrail(unreachable_fallback="fail_closed", fail_on_error=True) + bad = MagicMock(spec=httpx.Response) + bad.status_code = 401 + bad.text = "unauthorized" + g.async_handler.post.return_value = bad + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + + +@pytest.mark.asyncio +async def test_payload_size_guard_fails_closed(): + g = _make_guardrail(max_payload_bytes=10) + inputs = {"texts": ["x" * 5000]} + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + g.async_handler.post.assert_not_called() + + +@pytest.mark.asyncio +async def test_payload_size_guard_blocks_even_with_fail_open(): + g = _make_guardrail(max_payload_bytes=10, unreachable_fallback="fail_open") + inputs = {"texts": ["x" * 5000]} + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + g.async_handler.post.assert_not_called() + + +@pytest.mark.asyncio +async def test_invalid_response_schema_blocks_even_with_fail_open(): + g = _make_guardrail(unreachable_fallback="fail_open") + bad = MagicMock(spec=httpx.Response) + bad.status_code = 200 + bad.json.return_value = {"action": "NOT_A_VALID_ACTION"} + bad.text = "" + g.async_handler.post.return_value = bad + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + + +@pytest.mark.asyncio +async def test_unreachable_http_status_fail_open_passes(): + g = _make_guardrail(unreachable_fallback="fail_open") + resp = MagicMock(spec=httpx.Response) + resp.status_code = 503 + resp.text = "service unavailable" + g.async_handler.post.return_value = resp + inputs = {"texts": ["x"]} + out = await g.apply_guardrail( + inputs=inputs, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) + assert out is inputs + + +@pytest.mark.asyncio +async def test_unreachable_http_status_fail_closed_blocks(): + g = _make_guardrail(unreachable_fallback="fail_closed") + resp = MagicMock(spec=httpx.Response) + resp.status_code = 503 + resp.text = "service unavailable" + g.async_handler.post.return_value = resp + with pytest.raises(GuardrailRaisedException): + await g.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request", logging_obj=_logging_obj() + ) diff --git a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py index cf92f9cd12b..e4bf06991b4 100644 --- a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py +++ b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py @@ -124,6 +124,47 @@ def test_close_dangling_otel_server_span_records_status_and_ends(monkeypatch): } +def test_close_dangling_otel_server_span_v2_stamps_error_without_ending(monkeypatch): + """LIT-4179: under OTel v2 the FastAPI instrumentor owns the SERVER span, so + the handler must only stamp error.* on it (via record_error_attributes_on_span) + and must NOT set status, end the span, or clear request state — otherwise the + instrumentor's http.* attributes and span close are lost.""" + import litellm.integrations.otel.model.config as otel_config + import litellm.proxy.proxy_server as ps + + span = MagicMock() + fake_logger = MagicMock() + monkeypatch.setattr(ps, "open_telemetry_logger", fake_logger, raising=False) + monkeypatch.setattr(otel_config, "is_otel_v2_enabled", lambda: True) + request = _make_request(parent_otel_span=span) + exc = ProxyException(message="bad", type="bad_request_error", param=None, code=400) + + _close_dangling_otel_server_span(request=request, status_code=422, exc=exc) + + fake_logger.record_error_attributes_on_span.assert_called_once_with(span, exc, 422) + assert not span.end.called + assert not span.set_status.called + assert not fake_logger.set_response_status_code_attribute.called + assert request.state.parent_otel_span is span + + +def test_close_dangling_otel_server_span_v2_success_does_not_stamp(monkeypatch): + """Under v2 a sub-400 status must not stamp an error onto the SERVER span.""" + import litellm.integrations.otel.model.config as otel_config + import litellm.proxy.proxy_server as ps + + span = MagicMock() + fake_logger = MagicMock() + monkeypatch.setattr(ps, "open_telemetry_logger", fake_logger, raising=False) + monkeypatch.setattr(otel_config, "is_otel_v2_enabled", lambda: True) + request = _make_request(parent_otel_span=span) + + _close_dangling_otel_server_span(request=request, status_code=200) + + assert not fake_logger.record_error_attributes_on_span.called + assert not span.end.called + + def test_close_dangling_otel_server_span_missing_span_is_noop_error(): """When parent_otel_span is missing the call short-circuits — no error.""" request = _make_request(parent_otel_span=None) diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 25e84fb59a7..577af3dcffc 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -725,6 +725,7 @@ async def test_v1_models_translates_team_model_for_access_group_key(monkeypatch) router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() + router.get_configured_token_limits.return_value = (None, None) router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] @@ -766,6 +767,7 @@ async def test_v1_models_keeps_internal_names_when_public_name_flag_disabled( router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() + router.get_configured_token_limits.return_value = (None, None) router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] @@ -800,6 +802,7 @@ async def test_v1_models_translates_team_model_with_metadata(monkeypatch): router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() + router.get_configured_token_limits.return_value = (None, None) router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] router.get_model_group_info.return_value = None @@ -845,6 +848,7 @@ async def test_v1_models_metadata_fallbacks_use_internal_routing_key(monkeypatch router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() + router.get_configured_token_limits.return_value = (None, None) router.model_list = [team_dep] router.get_model_list.return_value = [team_dep] # Fallbacks are keyed on the internal routing name, as the router stores them. @@ -901,6 +905,7 @@ async def test_v1_models_metadata_does_not_leak_other_team_fallbacks(monkeypatch router.get_model_names.return_value = ["model_name_teamX_uuid9"] router.get_model_access_groups.return_value = {"grp-a": ["model_name_teamX_uuid9"]} router.get_fully_blocked_model_names.return_value = set() + router.get_configured_token_limits.return_value = (None, None) router.model_list = [team_x, team_y] router.get_model_list.return_value = [team_x, team_y] router.fallbacks = [ @@ -1155,6 +1160,7 @@ def test_translate_team_model_names_for_listing_respects_legacy_flag(): def _public_named_router(*team_rows: dict) -> MagicMock: router = MagicMock() router.get_model_list.return_value = list(team_rows) + router.get_configured_token_limits.return_value = (None, None) return router diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9a8f8146d6f..51d72aa2ab2 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -46,6 +46,69 @@ from litellm.types.utils import ( ) +def _get_additional_usage_values_for_usage(usage: litellm.Usage) -> dict: + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj=litellm.ModelResponse( + id="chatcmpl-test", + choices=[], + usage=usage, + ), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + metadata = json.loads(payload["metadata"]) + return metadata["additional_usage_values"] + + +def test_get_logging_payload_maps_openai_cached_tokens_to_cache_read_input_tokens(): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + prompt_tokens_details={"cached_tokens": 123}, + ) + ) + + assert additional_usage_values["cache_read_input_tokens"] == 123 + assert additional_usage_values["prompt_tokens_details"]["cached_tokens"] == 123 + + +def test_get_logging_payload_preserves_anthropic_cache_read_input_tokens(): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + prompt_tokens_details={"cached_tokens": 123}, + cache_read_input_tokens=456, + ) + ) + + assert additional_usage_values["cache_read_input_tokens"] == 456 + + +@pytest.mark.parametrize( + "prompt_tokens_details", + [None, {"cached_tokens": 0}], +) +def test_get_logging_payload_does_not_map_missing_or_zero_cached_tokens(prompt_tokens_details): + additional_usage_values = _get_additional_usage_values_for_usage( + litellm.Usage( + prompt_tokens=10, + completion_tokens=2, + total_tokens=12, + prompt_tokens_details=prompt_tokens_details, + ) + ) + + assert "cache_read_input_tokens" not in additional_usage_values + + def test_sanitize_request_body_for_spend_logs_payload_basic(): request_body = { "messages": [{"role": "user", "content": "Hello, how are you?"}], diff --git a/tests/test_litellm/proxy/test_model_id_header_propagation.py b/tests/test_litellm/proxy/test_model_id_header_propagation.py index e48168f89b3..f7f3eabae6d 100644 --- a/tests/test_litellm/proxy/test_model_id_header_propagation.py +++ b/tests/test_litellm/proxy/test_model_id_header_propagation.py @@ -200,6 +200,60 @@ def test_get_custom_headers_without_model_id(): assert headers["x-litellm-model-id"] in [None, ""] +class _FakeLoggingObj: + def __init__(self, litellm_params): + self.litellm_params = litellm_params + self.litellm_call_id = "test-call-id" + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +def test_get_custom_headers_includes_deployment_model_name(metadata_key): + """ + x-litellm-model-name should expose the underlying deployment model string, + which the router records under litellm_params[metadata]["deployment"]. + """ + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.tpm_limit = 1000 + mock_user_api_key_dict.rpm_limit = 100 + + logging_obj = _FakeLoggingObj( + litellm_params={metadata_key: {"deployment": "azure/gpt-4o-2024-08-06"}} + ) + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + model_id="deployment-uuid", + request_data={}, + hidden_params={}, + litellm_logging_obj=logging_obj, + ) + + assert headers["x-litellm-model-name"] == "azure/gpt-4o-2024-08-06" + assert headers["x-litellm-model-id"] == "deployment-uuid" + + +def test_get_custom_headers_omits_model_name_when_deployment_missing(): + """ + Without a deployment model string, x-litellm-model-name must not be emitted + (rather than leaking an empty/None value). + """ + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.tpm_limit = 1000 + mock_user_api_key_dict.rpm_limit = 100 + + logging_obj = _FakeLoggingObj(litellm_params={"metadata": {}}) + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + model_id="deployment-uuid", + request_data={}, + hidden_params={}, + litellm_logging_obj=logging_obj, + ) + + assert "x-litellm-model-name" not in headers + + def test_get_custom_headers_with_empty_string_model_id(): """ Test that get_custom_headers handles empty string model_id correctly. diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 54db0c0fd4f..a100e7837f4 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8984,6 +8984,261 @@ async def test_update_config_field_throttle_persists_to_litellm_settings(monkeyp assert saved["litellm_settings"]["budget_exceeded_throttle_percentage"] == 0.1 +def test_get_config_list_includes_anthropic_prompt_caching_fields(monkeypatch): + """The auto prompt caching flag and its ttl are litellm_settings globals surfaced on the + General Settings table, so an admin can turn caching on without hand-writing config. The + ttl is a Select and must ship its allowed values, or the table renders no editor for it.""" + import types + from unittest.mock import AsyncMock, MagicMock + + from fastapi.testclient import TestClient + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.proxy_server import app + + mock_prisma = MagicMock() + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + monkeypatch.setattr(litellm, "anthropic_prompt_caching_ttl", "1h") + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + client = TestClient(app) + resp = client.get("/config/list", params={"config_type": "general_settings"}) + assert resp.status_code == 200, resp.text + fields = {item["field_name"]: item for item in resp.json()} + + assert fields["enable_anthropic_prompt_caching"]["field_type"] == "Boolean" + assert fields["enable_anthropic_prompt_caching"]["field_value"] is True + + assert fields["anthropic_prompt_caching_ttl"]["field_type"] == "Select" + assert fields["anthropic_prompt_caching_ttl"]["field_value"] == "1h" + assert fields["anthropic_prompt_caching_ttl"]["field_options"] == ["5m", "1h"] + + # Both caching fields carry their sub-tab so the Admin UI can render them on a + # dedicated Prompt Caching tab, while ungrouped fields stay on General. + assert fields["enable_anthropic_prompt_caching"]["field_tab"] == "prompt_caching" + assert fields["anthropic_prompt_caching_ttl"]["field_tab"] == "prompt_caching" + assert fields["budget_exceeded_throttle_percentage"]["field_tab"] is None + finally: + app.dependency_overrides.clear() + + +def test_general_settings_ui_fields_are_db_overridable(): + """Every field the Admin UI can edit is a `litellm.` set via setattr on the handling + worker (`_persist_general_settings_ui_litellm_field`). Unless it is also in + LITELLM_SETTINGS_SAFE_DB_OVERRIDES, a config reload on a peer worker merges the DB value but + never applies it to the live attribute, so peer workers stay on their startup value. + + This invariant is the guard against the two registries drifting: adding a UI-editable field + without enrolling it in the DB-override allowlist silently breaks cross-worker propagation. + """ + from litellm.constants import LITELLM_SETTINGS_SAFE_DB_OVERRIDES + from litellm.proxy.proxy_server import _GENERAL_SETTINGS_UI_LITELLM_FIELDS + + missing = set(_GENERAL_SETTINGS_UI_LITELLM_FIELDS) - set(LITELLM_SETTINGS_SAFE_DB_OVERRIDES) + assert not missing, ( + f"UI-editable litellm_settings fields missing from LITELLM_SETTINGS_SAFE_DB_OVERRIDES: {sorted(missing)}. " + "Add them, or they will not propagate to other workers when changed from the UI." + ) + + +@pytest.mark.parametrize( + "field_name, db_value", + [ + ("enable_anthropic_prompt_caching", True), + ("anthropic_prompt_caching_ttl", "1h"), + ], +) +def test_prompt_caching_settings_propagate_on_config_reload(monkeypatch, field_name, db_value): + """A UI toggle on one worker persists to the DB; a peer worker picks it up only when the + config reload applies the safe-override allowlist. Regression for the fields being absent + from that allowlist, which left peer workers stale.""" + import litellm.proxy.proxy_server as ps + + # peer worker booted with the opposite/absent value + monkeypatch.setattr(litellm, field_name, False if isinstance(db_value, bool) else None) + + pc = ps.ProxyConfig() + pc._update_config_fields( + current_config={"litellm_settings": {}}, + param_name="litellm_settings", + db_param_value={field_name: db_value}, + ) + + assert getattr(litellm, field_name) == db_value + + +def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypatch): + """The flag defaults to False rather than None, so a plain 'is not None' check would + report the default as 'In Config' and imply an admin had set it.""" + import types + from unittest.mock import AsyncMock, MagicMock + + from fastapi.testclient import TestClient + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.proxy_server import app + + mock_prisma = MagicMock() + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + client = TestClient(app) + resp = client.get("/config/list", params={"config_type": "general_settings"}) + fields = {item["field_name"]: item for item in resp.json()} + assert fields["enable_anthropic_prompt_caching"]["stored_in_db"] is None + finally: + app.dependency_overrides.clear() + + +@pytest.mark.parametrize( + "field_name, field_value", + [ + ("enable_anthropic_prompt_caching", True), + ("enable_anthropic_prompt_caching", False), + ("anthropic_prompt_caching_ttl", "5m"), + ("anthropic_prompt_caching_ttl", "1h"), + ], +) +@pytest.mark.asyncio +async def test_update_config_field_prompt_caching_persists_to_litellm_settings(monkeypatch, field_name, field_value): + """Toggling either row must set litellm. live and persist under litellm_settings, + so the running proxy caches immediately and still does after a restart.""" + from unittest.mock import MagicMock + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import ( + ConfigFieldUpdate, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.proxy_server import update_config_general_settings + + saved: dict = {} + + async def fake_get_config(): + return {"litellm_settings": {}} + + async def fake_save_config(new_config=None): + saved.update(new_config or {}) + + monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config) + monkeypatch.setattr(ps.proxy_config, "save_config", fake_save_config) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(litellm, "store_audit_logs", False) + monkeypatch.setattr(litellm, field_name, None) + + admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN) + await update_config_general_settings( + data=ConfigFieldUpdate(field_name=field_name, field_value=field_value, config_type="general_settings"), + user_api_key_dict=admin, + ) + + assert getattr(litellm, field_name) == field_value + assert saved["litellm_settings"][field_name] == field_value + + +@pytest.mark.parametrize( + "field_name, bad_value", + [ + ("enable_anthropic_prompt_caching", "yes"), + ("enable_anthropic_prompt_caching", 1), + ("anthropic_prompt_caching_ttl", "10m"), + ("anthropic_prompt_caching_ttl", "1H"), + ("anthropic_prompt_caching_ttl", 3600), + ], +) +@pytest.mark.asyncio +async def test_update_config_field_prompt_caching_rejects_invalid(monkeypatch, field_name, bad_value): + """An unsupported ttl must be refused here rather than reaching Anthropic verbatim.""" + from unittest.mock import MagicMock + + from fastapi import HTTPException + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import ( + ConfigFieldUpdate, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.proxy_server import update_config_general_settings + + async def fake_get_config(): + return {"litellm_settings": {}} + + monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(litellm, field_name, None) + + admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN) + with pytest.raises(HTTPException) as exc: + await update_config_general_settings( + data=ConfigFieldUpdate(field_name=field_name, field_value=bad_value, config_type="general_settings"), + user_api_key_dict=admin, + ) + assert exc.value.status_code == 400 + assert getattr(litellm, field_name) is None + + +@pytest.mark.parametrize( + "field_name, expected_default", + [ + ("enable_anthropic_prompt_caching", False), + ("anthropic_prompt_caching_ttl", None), + ("budget_exceeded_throttle_percentage", None), + ], +) +@pytest.mark.asyncio +async def test_reset_config_field_restores_type_default(monkeypatch, field_name, expected_default): + """Reset must restore each field's own default. Blanket None would leave the boolean flag + set to None, which is not a bool and would read as neither on nor off.""" + from unittest.mock import MagicMock + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import ( + ConfigFieldDelete, + LitellmUserRoles, + UserAPIKeyAuth, + ) + from litellm.proxy.proxy_server import delete_config_general_settings + + saved: dict = {} + + async def fake_get_config(): + return {"litellm_settings": {field_name: "stale"}} + + async def fake_save_config(new_config=None): + saved.update(new_config or {}) + + monkeypatch.setattr(ps.proxy_config, "get_config", fake_get_config) + monkeypatch.setattr(ps.proxy_config, "save_config", fake_save_config) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(litellm, "store_audit_logs", False) + monkeypatch.setattr(litellm, field_name, "stale") + + admin = UserAPIKeyAuth(api_key="k", user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN) + await delete_config_general_settings( + data=ConfigFieldDelete(field_name=field_name, config_type="general_settings"), + user_api_key_dict=admin, + ) + + assert getattr(litellm, field_name) is expected_default + assert field_name not in saved["litellm_settings"] + + @pytest.mark.parametrize("bad_value", [0, -0.1, 1.5, True]) @pytest.mark.asyncio async def test_update_config_field_throttle_rejects_invalid(monkeypatch, bad_value): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 55c09e6cac4..76a6e3c1bbe 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2864,6 +2864,185 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch assert calls == [1] +def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): + """ + Responses API calls pass `input` (str) instead of `messages`. Context-window + checks must count tokens from `input` and filter deployments over the limit. Uses + the real token_counter so the transform + counting path is a true regression guard. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} + ) + + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="m", + healthy_deployments=deployments, + input="a very long prompt that exceeds the tiny context window", + ) + + +def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): + """ + Responses API `input` can be a list of input items. It must be normalized to + chat messages and counted so oversized requests are filtered out. Uses the real + token_counter (no mock) so the transform + counting path is a true regression guard. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} + ) + + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="m", + healthy_deployments=deployments, + input=[ + {"role": "user", "content": "count these tokens against the one token limit please"}, + ], + ) + + +def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): + """ + Responses API `instructions` become a system message the model receives, so their + tokens must be counted too. A request whose `input` alone fits under the limit but + whose `input` + `instructions` exceeds it must be filtered (regression for the + context-window check under-filtering when instructions were ignored). + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + + short_input = "hi" + long_instructions = "you are a helpful assistant. " * 20 + + input_only_tokens = router._count_pre_call_check_tokens(messages=None, input=short_input) + with_instructions_tokens = router._count_pre_call_check_tokens( + messages=None, input=short_input, instructions=long_instructions + ) + assert with_instructions_tokens > input_only_tokens + + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} + ) + with pytest.raises(litellm.ContextWindowExceededError): + router._pre_call_checks( + model="m", + healthy_deployments=deployments, + input=short_input, + request_kwargs={"instructions": long_instructions}, + ) + + +def test_count_pre_call_check_tokens_across_api_surfaces(): + """ + _count_pre_call_check_tokens must count tokens from chat `messages`, a Responses + API string `input`, and a Responses API list `input`, and raise when given neither. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + ) + + messages_tokens = router._count_pre_call_check_tokens( + messages=[{"role": "user", "content": "hello world"}], input=None + ) + string_input_tokens = router._count_pre_call_check_tokens(messages=None, input="hello world") + list_input_tokens = router._count_pre_call_check_tokens( + messages=None, input=[{"role": "user", "content": "hello world"}] + ) + + assert messages_tokens > 0 + assert string_input_tokens > 0 + assert list_input_tokens > 0 + + with pytest.raises(ValueError): + router._count_pre_call_check_tokens(messages=None, input=None) + + +def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): + """ + When neither messages nor input is provided (e.g. endpoints without prompt text), + token counting is skipped gracefully and all deployments are returned. + """ + router = litellm.Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, + ], + enable_pre_call_checks=True, + ) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) + + counted: list[dict] = [] + original = router._count_pre_call_check_tokens + monkeypatch.setattr( + router, + "_count_pre_call_check_tokens", + lambda **kwargs: counted.append(kwargs) or original(**kwargs), + ) + + deployments = [ + {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, + ] + result = router._pre_call_checks(model="m", healthy_deployments=deployments) + assert len(result) == 1 + assert counted == [] # token counting skipped entirely, so no misleading error is logged + + +@pytest.mark.asyncio +async def test_aresponses_enforces_context_window_pre_call_check(): + """ + End-to-end router regression: a Responses API call whose `input` exceeds the + deployment's max_input_tokens must be filtered by the pre-call check, raising + ContextWindowExceededError instead of being silently routed. This guards the + wiring that forwards `input` from the generic-call path into deployment selection + (the deployment uses mock_response, so the check must trip before any real call). + """ + router = litellm.Router( + model_list=[ + { + "model_name": "small-ctx", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + "model_info": {"max_input_tokens": 5}, + } + ], + enable_pre_call_checks=True, + ) + with pytest.raises(litellm.ContextWindowExceededError): + await router.aresponses( + model="small-ctx", + input="this responses input is definitely much longer than five tokens for sure", + ) + + def test_get_deployment_model_info_base_model_flow(): """Test that get_deployment_model_info correctly handles the base model flow""" from unittest.mock import patch diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index edd93cbebe0..a1a9448cc58 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4317,7 +4317,7 @@ _FIREWORKS_MODELS = [ "accounts/fireworks/models/glm-5p2", 1.4e-06, 4.4e-06, - 2.6e-07, + 1.4e-07, 1048576, 131072, False, diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 32e9a03da95..c775af81ba8 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -881,19 +881,6 @@ "count": 1 } }, - "src/app/(dashboard)/policies/_components/attachment_table.test.tsx": { - "react/display-name": { - "count": 1 - } - }, - "src/app/(dashboard)/policies/_components/attachment_table.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/policies/_components/guardrail_selection_modal.tsx": { "no-nested-ternary": { "count": 1 @@ -1079,7 +1066,7 @@ }, "src/app/(dashboard)/router-settings/_components/general_settings.tsx": { "no-nested-ternary": { - "count": 3 + "count": 1 }, "no-restricted-imports": { "count": 2 @@ -1885,19 +1872,11 @@ "count": 1 } }, - "src/components/model_add/AddCredentialModal.tsx": { + "src/components/model_add/CredentialModal.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/model_add/EditCredentialModal.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/model_add/credentials.tsx": { "no-restricted-imports": { "count": 1 diff --git a/ui/litellm-dashboard/public/assets/logos/straiker.svg b/ui/litellm-dashboard/public/assets/logos/straiker.svg new file mode 100644 index 00000000000..bdfe0405736 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/straiker.svg @@ -0,0 +1,9 @@ + + + + + + + + + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts index 2ad5819b5f0..a40587cb3ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_configs.ts @@ -300,4 +300,10 @@ export const GUARDRAIL_PRESETS: Record = { mode: "pre_call", defaultOn: false, }, + straiker: { + provider: "Straiker", + guardrailNameSuggestion: "Straiker Guardrail", + mode: "pre_call", + defaultOn: false, + }, }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts index f81277f13c3..ba11d3d400d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_data.ts @@ -442,6 +442,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ tags: ["Security", "Policy", "Prompt Injection"], providerKey: "Repelloai", }, + { + id: "straiker", + name: "Straiker", + description: + "Defend AI Agentic Guardrails: Indirect/Direct Prompt Injection, Tool Misuse, Malicious MCP and Skills", + category: "partner", + logo: `${ASSET_PREFIX}straiker.svg`, + tags: ["Agentic", "Prompt Injection", "Tool Misuse", "MCP", "Skills"], + providerKey: "Straiker", + }, ]; export const ALL_CARDS = [...LITELLM_CONTENT_FILTER_CARDS, ...PARTNER_GUARDRAIL_CARDS]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index e8c4810ca69..a2873797096 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -166,6 +166,7 @@ export const guardrailLogoMap: Record = { Akto: `${asset_logos_folder}akto.svg`, "Qostodian Nexus": `${asset_logos_folder}qohash.jpg`, "RepelloAI Argus": `${asset_logos_folder}repelloai.png`, + Straiker: `${asset_logos_folder}straiker.svg`, }; export const getGuardrailLogoAndName = (guardrailValue: string): { logo: string; displayName: string } => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/attachment_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx similarity index 55% rename from ui/litellm-dashboard/src/app/(dashboard)/policies/_components/attachment_table.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx index c53881e5cce..b544c44d190 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/attachment_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.test.tsx @@ -1,63 +1,17 @@ import React from "react"; -import { screen } from "@testing-library/react"; +import { screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "@/../tests/test-utils"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import AttachmentTable from "./attachment_table"; +import AttachmentTable from "./AttachmentTable"; import { PolicyAttachment } from "@/components/policies/types"; vi.mock("./impact_popover", () => ({ - default: () =>