Merge litellm_internal_staging into devin/1784399610-rust-anthropic-messages

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

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")
+);

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

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/<route>/transformation.rs` (e.g. `AnthropicMessagesProviderConfig`, mirroring `OcrProviderConfig`).
+4. Each provider implements that trait as a `const <PROVIDER>_<ROUTE>_CONFIG` in `core/src/providers/<provider>/<route>/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/<provider>/<route>/`; a route is a module, never a new crate.
+9. Route entry point stays thin: `<route>()` -> `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<String>` 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
+    ```

litellm-rust/crates/core/tests/workspace_crate_allowlist.rs

@@ -62,12 +62,17 @@ fn parse_members(manifest: &str) -> BTreeSet<String> {
     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<String> {
     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()
 }

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))

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,

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

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:

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

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,

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,

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,
     )
 
 

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(

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()

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",

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:

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"))

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,
+    )

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

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):

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(
         *,

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,
+}

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

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,
             )
         )

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

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:

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"
 

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,
             )
 

litellm/types/guardrails.py

@@ -131,6 +131,7 @@ class SupportedGuardrailIntegrations(Enum):
     SINGULR = "singulr"
     HEADROOM = "headroom"
     COMPRESR = "compresr"
+    STRAIKER = "straiker"
 
 
 class Role(Enum):

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"

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"

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"

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,

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",

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

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
 

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
 

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)

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)

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 |

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)

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

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
     ]

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

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 `<property>` 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()

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

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

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)

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)

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)

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)

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)

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)

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,

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"

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,

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,
     )

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)
 
 

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,
     )

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,

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,
     )

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

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())

tests/e2e/load/load_constants.py

@@ -0,0 +1,3 @@
+from __future__ import annotations
+
+LOAD_MODEL = "load-mock"

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)

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",
+        )

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"
+        )

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")

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)

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]}"
     )

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]}"

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, (

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")

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)

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}"
 

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}"

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)

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)

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
 
 

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"]

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,

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

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)

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)

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

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), (

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

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 = []

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]

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]

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)

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)

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")

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)

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)

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)

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)

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

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.

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