chore: merge staging into OCR cutover

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

basedpyright-code-budget.json

@@ -24,7 +24,7 @@
     "limit": 42
   },
   "reportExplicitAny": {
-    "limit": 10397
+    "limit": 10389
   },
   "reportFunctionMemberAccess": {
     "limit": 11

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/ai-gateway/src/constants.rs

@@ -28,7 +28,6 @@ pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500;
 /// Provider attributed to realtime sessions in the logging payload.
 #[cfg(feature = "server")]
 pub(crate) const DEFAULT_PROVIDER: &str = "openai";
-
 pub(crate) const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
 
 pub(crate) const VERTEXAI_CREDENTIALS_ENV: &str = "VERTEXAI_CREDENTIALS";
@@ -46,3 +45,15 @@ pub(crate) const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
 pub(crate) const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0;
 
 pub(crate) const MAX_SAFE_FETCH_REDIRECTS: usize = 10;
+
+/// Full-request timeout ceiling for Anthropic Messages provider calls, in
+/// seconds. Mirrors the Python Anthropic Messages default. The per-request
+/// timeout from `litellm_params` still overrides this on the request builder.
+pub(crate) const MESSAGES_TIMEOUT_SECS: u64 = 600;
+
+/// Connect timeout for Anthropic Messages provider calls, in seconds.
+pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10;
+
+/// Max characters of an upstream error body echoed across the host boundary
+/// before truncation, so provider bodies are bounded and data-minimized.
+pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256;

litellm-rust/crates/ai-gateway/src/io/messages.rs

@@ -0,0 +1 @@
+pub use crate::messages::{messages, MessagesRequest};

litellm-rust/crates/ai-gateway/src/io/mod.rs

@@ -1,3 +1,4 @@
+pub mod messages;
 pub mod ocr;
 pub mod realtime;
 pub mod realtime_pool;

litellm-rust/crates/ai-gateway/src/lib.rs

@@ -13,6 +13,7 @@
 
 pub(crate) mod config;
 pub mod io;
+pub mod messages;
 pub mod ocr;
 
 /// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and

litellm-rust/crates/ai-gateway/src/messages/client.rs

@@ -0,0 +1,15 @@
+use std::sync::OnceLock;
+use std::time::Duration;
+
+use crate::constants::{MESSAGES_CONNECT_TIMEOUT_SECS, MESSAGES_TIMEOUT_SECS};
+
+pub(super) fn http_client() -> &'static reqwest::Client {
+    static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
+    CLIENT.get_or_init(|| {
+        reqwest::Client::builder()
+            .timeout(Duration::from_secs(MESSAGES_TIMEOUT_SECS))
+            .connect_timeout(Duration::from_secs(MESSAGES_CONNECT_TIMEOUT_SECS))
+            .build()
+            .unwrap_or_else(|_| reqwest::Client::new())
+    })
+}

litellm-rust/crates/ai-gateway/src/messages/common_utils.rs

@@ -0,0 +1,50 @@
+use litellm_core::error::{json_type_name, CoreError};
+use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
+use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
+use litellm_core::CoreResult;
+use serde_json::{Map, Value};
+
+use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS;
+
+pub(super) fn truncate_error_body(body: &str) -> String {
+    if body.chars().count() <= MESSAGES_ERROR_BODY_MAX_CHARS {
+        return body.to_string();
+    }
+    let truncated: String = body.chars().take(MESSAGES_ERROR_BODY_MAX_CHARS).collect();
+    format!("{truncated}... (truncated)")
+}
+
+pub(super) fn messages_provider_config(
+    provider: &str,
+) -> Option<&'static dyn AnthropicMessagesProviderConfig> {
+    match provider {
+        "azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG),
+        _ => None,
+    }
+}
+
+pub(super) fn string_headers(
+    extra_headers: Option<Map<String, Value>>,
+) -> CoreResult<Vec<(String, String)>> {
+    extra_headers
+        .unwrap_or_default()
+        .into_iter()
+        .map(|(key, value)| {
+            value
+                .as_str()
+                .map(|value| (key.clone(), value.to_string()))
+                .ok_or_else(|| {
+                    CoreError::InvalidRequest(format!(
+                        "messages extra_headers.{key} must be a string, got {}",
+                        json_type_name(&value)
+                    ))
+                })
+        })
+        .collect()
+}
+
+pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
+    headers
+        .iter()
+        .any(|(key, _)| key.eq_ignore_ascii_case(name))
+}

litellm-rust/crates/ai-gateway/src/messages/handler.rs

@@ -0,0 +1,47 @@
+use litellm_core::error::CoreError;
+use litellm_core::CoreResult;
+use serde_json::Value;
+
+use super::client::http_client;
+use super::common_utils::truncate_error_body;
+use super::types::ProviderMessagesRequest;
+
+pub(super) async fn execute_messages_provider_call(
+    request: ProviderMessagesRequest,
+) -> CoreResult<Value> {
+    let mut request_builder = http_client().post(&request.url).json(&request.body);
+    for (key, value) in &request.upstream_headers {
+        request_builder = request_builder.header(key, value);
+    }
+    if let Some(duration) = request.timeout {
+        request_builder = request_builder.timeout(duration);
+    }
+
+    let response = request_builder
+        .send()
+        .await
+        .map_err(|err| CoreError::Network(err.to_string()))?;
+
+    let status = response.status();
+    let text = response
+        .text()
+        .await
+        .map_err(|err| CoreError::Network(err.to_string()))?;
+
+    if !status.is_success() {
+        return Err(CoreError::Http {
+            status: status.as_u16(),
+            body: truncate_error_body(&text),
+        });
+    }
+
+    let response = serde_json::from_str(&text).map_err(|err| {
+        CoreError::InvalidResponse(format!("invalid messages response JSON: {err}"))
+    })?;
+    let transformed = request
+        .config
+        .transform_response(&request.model, response)?;
+    serde_json::to_value(transformed).map_err(|err| {
+        CoreError::InvalidResponse(format!("failed to serialize messages response: {err}"))
+    })
+}

litellm-rust/crates/ai-gateway/src/messages/mod.rs

@@ -0,0 +1,21 @@
+use litellm_core::CoreResult;
+use serde_json::Value;
+
+mod client;
+mod common_utils;
+mod handler;
+mod prepare;
+mod types;
+
+pub use types::MessagesRequest;
+
+use handler::execute_messages_provider_call;
+use prepare::prepare_messages_call;
+
+pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<Value> {
+    let prepared = prepare_messages_call(request)?;
+    execute_messages_provider_call(prepared).await
+}
+
+#[cfg(test)]
+mod tests;

litellm-rust/crates/ai-gateway/src/messages/prepare.rs

@@ -0,0 +1,72 @@
+use litellm_core::messages::transformation::MessagesAuthStrategy;
+use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider};
+use litellm_core::CoreError;
+use litellm_core::CoreResult;
+
+use super::common_utils::{has_header, messages_provider_config, string_headers};
+use super::types::{MessagesRequest, ProviderMessagesRequest};
+
+pub(super) fn prepare_messages_call(
+    request: MessagesRequest<'_>,
+) -> CoreResult<ProviderMessagesRequest> {
+    let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
+        .or_else(|| {
+            request
+                .custom_llm_provider
+                .map(|provider| CustomLlmProvider {
+                    model: request.model,
+                    custom_llm_provider: provider,
+                })
+        })
+        .ok_or_else(|| {
+            CoreError::InvalidProvider(
+                "unable to resolve custom_llm_provider for messages request".to_string(),
+            )
+        })?;
+    let model = provider_info.model.to_string();
+    let provider = provider_info.custom_llm_provider;
+
+    let config = messages_provider_config(provider)
+        .ok_or_else(|| CoreError::InvalidProvider(provider.to_string()))?;
+    let env_lookup = |key: &str| std::env::var(key).ok();
+
+    let mut headers = string_headers(request.extra_headers)?;
+
+    let auth_strategy = config.auth_strategy();
+    if !has_header(&headers, auth_strategy.header_name()) {
+        let api_key = config.resolve_api_key(request.api_key, &env_lookup)?;
+        let auth_header = match auth_strategy {
+            MessagesAuthStrategy::Bearer => {
+                ("authorization".to_string(), format!("Bearer {api_key}"))
+            }
+            MessagesAuthStrategy::Header(name) => (name.to_string(), api_key),
+        };
+        headers.push(auth_header);
+    }
+
+    for (name, value) in config.default_headers() {
+        if !has_header(&headers, name) {
+            headers.push((name.to_string(), value.to_string()));
+        }
+    }
+
+    let url = config.complete_url(request.api_base, &model, &env_lookup)?;
+    let typed_request = serde_json::from_value(request.body).map_err(|err| {
+        CoreError::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
+    })?;
+    let transformed = config.transform_request(typed_request)?;
+    let body = serde_json::to_value(transformed).map_err(|err| {
+        CoreError::InvalidRequest(format!(
+            "failed to serialize Anthropic messages request: {err}"
+        ))
+    })?;
+
+    Ok(ProviderMessagesRequest {
+        model,
+        config,
+        url,
+        body,
+        upstream_headers: headers,
+        timeout: request.timeout,
+    })
+}

litellm-rust/crates/ai-gateway/src/messages/tests.rs

@@ -0,0 +1,259 @@
+use std::time::Duration;
+
+use litellm_core::error::CoreError;
+use serde_json::{json, Map, Value};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::{TcpListener, TcpStream};
+
+use super::common_utils::{
+    has_header, messages_provider_config, string_headers, truncate_error_body,
+};
+use super::{messages, MessagesRequest};
+
+async fn read_http_request(socket: &mut TcpStream) -> String {
+    let mut request = Vec::new();
+    let mut buffer = [0_u8; 1024];
+    let header_end = loop {
+        let n = socket.read(&mut buffer).await.expect("reads request");
+        if n == 0 {
+            break request.len();
+        }
+        request.extend_from_slice(&buffer[..n]);
+        if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
+            break position + 4;
+        }
+    };
+    let headers = String::from_utf8_lossy(&request[..header_end]);
+    let content_length = headers
+        .lines()
+        .find_map(|line| {
+            let (name, value) = line.split_once(':')?;
+            name.eq_ignore_ascii_case("content-length")
+                .then(|| value.trim().parse::<usize>().ok())
+                .flatten()
+        })
+        .unwrap_or(0);
+    while request.len().saturating_sub(header_end) < content_length {
+        let n = socket.read(&mut buffer).await.expect("reads body");
+        if n == 0 {
+            break;
+        }
+        request.extend_from_slice(&buffer[..n]);
+    }
+    String::from_utf8(request).expect("request is utf8")
+}
+
+fn write_response(body: &str) -> String {
+    format!(
+        "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
+        body.len(),
+        body
+    )
+}
+
+#[test]
+fn provider_config_only_resolves_azure_ai() {
+    assert!(messages_provider_config("azure_ai").is_some());
+    assert!(messages_provider_config("anthropic").is_none());
+    assert!(messages_provider_config("openai").is_none());
+}
+
+#[test]
+fn truncate_error_body_caps_long_payloads() {
+    let body = "x".repeat(400);
+    let truncated = truncate_error_body(&body);
+    assert!(truncated.ends_with("... (truncated)"));
+    let prefix_chars = truncated
+        .strip_suffix("... (truncated)")
+        .expect("truncated marker present")
+        .chars()
+        .count();
+    assert_eq!(prefix_chars, 256);
+}
+
+#[test]
+fn string_headers_rejects_non_string_values() {
+    let headers = json!({"x-count": 3}).as_object().unwrap().clone();
+    let err = string_headers(Some(headers)).expect_err("non-string header rejected");
+    assert!(matches!(err, CoreError::InvalidRequest(_)));
+}
+
+#[test]
+fn has_header_is_case_insensitive() {
+    let headers = vec![("X-Api-Key".to_string(), "secret".to_string())];
+    assert!(has_header(&headers, "x-api-key"));
+    assert!(!has_header(&headers, "authorization"));
+}
+
+#[tokio::test]
+async fn messages_round_trip_builds_azure_request_and_passes_response_through() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
+    let addr = listener.local_addr().expect("addr");
+
+    let server = tokio::spawn(async move {
+        let (mut socket, _) = listener.accept().await.expect("accepts request");
+        let request = read_http_request(&mut socket).await;
+        let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#;
+        socket
+            .write_all(write_response(response_body).as_bytes())
+            .await
+            .expect("writes response");
+        request
+    });
+
+    let response = messages(MessagesRequest {
+        model: "claude-sonnet-4-5",
+        body: json!({
+            "model": "claude-sonnet-4-5",
+            "max_tokens": 1024,
+            "messages": [{
+                "role": "user",
+                "content": [{
+                    "type": "text",
+                    "text": "hi",
+                    "cache_control": {"type": "ephemeral", "scope": "global"}
+                }]
+            }]
+        }),
+        api_key: Some("sk-azure"),
+        api_base: Some(&format!("http://{addr}")),
+        custom_llm_provider: Some("azure_ai"),
+        extra_headers: None,
+        timeout: Some(Duration::from_secs(5)),
+    })
+    .await
+    .expect("messages request succeeds");
+
+    assert_eq!(response["content"][0]["text"], "hi");
+    assert_eq!(response["stop_reason"], "end_turn");
+
+    let request = server.await.expect("server task completes");
+    let (head, body) = request.split_once("\r\n\r\n").expect("has body");
+    assert!(head.starts_with("POST /anthropic/v1/messages "), "{head}");
+    let head_lower = head.to_ascii_lowercase();
+    assert!(head_lower.contains("x-api-key: sk-azure"), "{head}");
+    assert!(
+        head_lower.contains("anthropic-version: 2023-06-01"),
+        "{head}"
+    );
+    assert!(
+        head_lower.contains("content-type: application/json"),
+        "{head}"
+    );
+
+    let sent_body: Value = serde_json::from_str(body).expect("body is json");
+    assert_eq!(
+        sent_body["messages"][0]["content"][0]["cache_control"],
+        json!({"type": "ephemeral"})
+    );
+}
+
+#[tokio::test]
+async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
+    let addr = listener.local_addr().expect("addr");
+
+    let server = tokio::spawn(async move {
+        let (mut socket, _) = listener.accept().await.expect("accepts request");
+        let request = read_http_request(&mut socket).await;
+        let response_body =
+            r#"{"id":"msg_2","type":"message","role":"assistant","content":[],"model":"m"}"#;
+        socket
+            .write_all(write_response(response_body).as_bytes())
+            .await
+            .expect("writes response");
+        request
+    });
+
+    let mut headers = Map::new();
+    headers.insert(
+        "x-api-key".to_string(),
+        Value::String("from-python".to_string()),
+    );
+    headers.insert(
+        "anthropic-beta".to_string(),
+        Value::String("token-efficient-tools-2025-02-19".to_string()),
+    );
+
+    messages(MessagesRequest {
+        model: "claude-sonnet-4-5",
+        body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
+        api_key: Some("rust-fallback-key"),
+        api_base: Some(&format!("http://{addr}")),
+        custom_llm_provider: Some("azure_ai"),
+        extra_headers: Some(headers),
+        timeout: Some(Duration::from_secs(5)),
+    })
+    .await
+    .expect("messages request succeeds");
+
+    let request = server.await.expect("server task completes");
+    let head = request
+        .split_once("\r\n\r\n")
+        .expect("has body")
+        .0
+        .to_ascii_lowercase();
+    let api_key_count = head
+        .lines()
+        .filter(|line| line.starts_with("x-api-key:"))
+        .count();
+    assert_eq!(api_key_count, 1, "{head}");
+    assert!(head.contains("x-api-key: from-python"), "{head}");
+    assert!(
+        head.contains("anthropic-beta: token-efficient-tools-2025-02-19"),
+        "{head}"
+    );
+    assert!(!head.contains("rust-fallback-key"), "{head}");
+}
+
+#[tokio::test]
+async fn messages_maps_provider_error_status_to_http_error() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
+    let addr = listener.local_addr().expect("addr");
+
+    tokio::spawn(async move {
+        let (mut socket, _) = listener.accept().await.expect("accepts request");
+        let _ = read_http_request(&mut socket).await;
+        let body = "unauthorized";
+        let response = format!(
+            "HTTP/1.1 401 Unauthorized\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
+            body.len(),
+            body
+        );
+        socket
+            .write_all(response.as_bytes())
+            .await
+            .expect("writes response");
+    });
+
+    let err = messages(MessagesRequest {
+        model: "claude-sonnet-4-5",
+        body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
+        api_key: Some("sk-azure"),
+        api_base: Some(&format!("http://{addr}")),
+        custom_llm_provider: Some("azure_ai"),
+        extra_headers: None,
+        timeout: Some(Duration::from_secs(5)),
+    })
+    .await
+    .expect_err("provider error propagates");
+
+    assert!(matches!(err, CoreError::Http { status: 401, .. }));
+}
+
+#[tokio::test]
+async fn messages_rejects_unsupported_provider() {
+    let err = messages(MessagesRequest {
+        model: "claude-3-5-sonnet",
+        body: json!({"model": "claude-3-5-sonnet", "max_tokens": 8, "messages": []}),
+        api_key: Some("sk"),
+        api_base: Some("http://127.0.0.1:1"),
+        custom_llm_provider: Some("anthropic"),
+        extra_headers: None,
+        timeout: Some(Duration::from_millis(50)),
+    })
+    .await
+    .expect_err("unsupported provider errors");
+
+    assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "anthropic"));
+}

litellm-rust/crates/ai-gateway/src/messages/types.rs

@@ -0,0 +1,23 @@
+use std::time::Duration;
+
+use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
+use serde_json::{Map, Value};
+
+pub struct MessagesRequest<'a> {
+    pub model: &'a str,
+    pub body: Value,
+    pub api_key: Option<&'a str>,
+    pub api_base: Option<&'a str>,
+    pub custom_llm_provider: Option<&'a str>,
+    pub extra_headers: Option<Map<String, Value>>,
+    pub timeout: Option<Duration>,
+}
+
+pub(crate) struct ProviderMessagesRequest {
+    pub(crate) model: String,
+    pub(crate) config: &'static dyn AnthropicMessagesProviderConfig,
+    pub(crate) url: String,
+    pub(crate) body: Value,
+    pub(crate) upstream_headers: Vec<(String, String)>,
+    pub(crate) timeout: Option<Duration>,
+}

litellm-rust/crates/core/src/error.rs

@@ -19,11 +19,11 @@ pub enum CoreError {
     InvalidRequest(String),
     #[error("{0}")]
     Auth(String),
-    #[error("OCR request failed with status {status}: {body}")]
+    #[error("upstream request failed with status {status}: {body}")]
     Http { status: u16, body: String },
     #[error("OCR request timed out")]
     Timeout,
-    #[error("OCR network error: {0}")]
+    #[error("upstream network error: {0}")]
     Network(String),
     #[error("routing error: {0}")]
     Routing(String),

litellm-rust/crates/core/src/lib.rs

@@ -2,6 +2,7 @@ pub mod cache;
 pub mod call_lifecycle;
 pub(crate) mod constants;
 pub mod error;
+pub mod messages;
 pub mod ocr;
 pub mod providers;
 pub mod realtime;

litellm-rust/crates/core/src/messages/mod.rs

@@ -0,0 +1,2 @@
+pub mod transformation;
+pub mod types;

litellm-rust/crates/core/src/messages/transformation.rs

@@ -0,0 +1,59 @@
+use crate::error::CoreResult;
+
+use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse};
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum MessagesAuthStrategy {
+    Bearer,
+    Header(&'static str),
+}
+
+impl MessagesAuthStrategy {
+    pub fn header_name(self) -> &'static str {
+        match self {
+            Self::Bearer => "authorization",
+            Self::Header(header_name) => header_name,
+        }
+    }
+}
+
+pub trait AnthropicMessagesProviderConfig: Sync {
+    fn complete_url(
+        &self,
+        api_base: Option<&str>,
+        model: &str,
+        env_lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> CoreResult<String>;
+
+    fn resolve_api_key(
+        &self,
+        api_key: Option<&str>,
+        env_lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> CoreResult<String>;
+
+    fn auth_strategy(&self) -> MessagesAuthStrategy {
+        MessagesAuthStrategy::Header("x-api-key")
+    }
+
+    fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
+        &[
+            ("anthropic-version", "2023-06-01"),
+            ("content-type", "application/json"),
+        ]
+    }
+
+    fn transform_request(
+        &self,
+        request: AnthropicMessagesRequest,
+    ) -> CoreResult<AnthropicMessagesRequest> {
+        Ok(request)
+    }
+
+    fn transform_response(
+        &self,
+        _model: &str,
+        response: AnthropicMessagesResponse,
+    ) -> CoreResult<AnthropicMessagesResponse> {
+        Ok(response)
+    }
+}

litellm-rust/crates/core/src/messages/types.rs

@@ -0,0 +1,110 @@
+use serde::{Deserialize, Serialize};
+use serde_json::{Map, Value};
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum SystemPrompt {
+    Text(String),
+    Blocks(Vec<ContentBlock>),
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum MessageContent {
+    Text(String),
+    Blocks(Vec<ContentBlock>),
+}
+
+#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
+pub struct ContentBlock {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub cache_control: Option<CacheControl>,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}
+
+#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
+pub struct CacheControl {
+    #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
+    pub cache_type: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub ttl: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub scope: Option<String>,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct AnthropicMessage {
+    pub role: String,
+    pub content: MessageContent,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct AnthropicMessagesRequest {
+    pub model: String,
+    pub messages: Vec<AnthropicMessage>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub max_tokens: Option<u64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub system: Option<SystemPrompt>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub metadata: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub stop_sequences: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub stream: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub temperature: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub top_p: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub top_k: Option<i64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tools: Option<Vec<Value>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_choice: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub thinking: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub service_tier: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub container: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub mcp_servers: Option<Vec<Value>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub context_management: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub output_format: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub output_config: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub speed: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub inference_geo: Option<String>,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct AnthropicMessagesResponse {
+    pub id: String,
+    #[serde(rename = "type")]
+    pub message_type: String,
+    pub role: String,
+    pub model: String,
+    pub content: Vec<Value>,
+    // Anthropic always includes stop_reason / stop_sequence, null until the turn
+    // ends; serialize them even when None so callers see the same shape as Python.
+    pub stop_reason: Option<String>,
+    pub stop_sequence: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub usage: Option<Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub container: Option<Value>,
+    #[serde(flatten)]
+    pub extra: Map<String, Value>,
+}

litellm-rust/crates/core/src/providers/anthropic/messages/mod.rs

@@ -0,0 +1 @@
+pub mod transformation;

litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs

@@ -0,0 +1,142 @@
+use crate::error::{CoreError, CoreResult};
+use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
+
+const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY";
+const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE";
+const DEFAULT_ANTHROPIC_API_BASE: &str = "https://api.anthropic.com";
+const MESSAGES_PATH_SUFFIX: &str = "/v1/messages";
+
+pub struct AnthropicMessagesConfig;
+
+pub const ANTHROPIC_MESSAGES_CONFIG: AnthropicMessagesConfig = AnthropicMessagesConfig;
+
+pub fn non_empty(value: Option<&str>) -> Option<&str> {
+    value.map(str::trim).filter(|value| !value.is_empty())
+}
+
+pub fn resolve_anthropic_api_key(
+    api_key: Option<&str>,
+    env_lookup: &dyn Fn(&str) -> Option<String>,
+) -> CoreResult<String> {
+    non_empty(api_key)
+        .map(str::to_string)
+        .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
+        .ok_or_else(|| {
+            CoreError::Auth(
+                "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY \
+                 environment variable"
+                    .to_string(),
+            )
+        })
+}
+
+pub fn complete_anthropic_url(
+    api_base: Option<&str>,
+    env_lookup: &dyn Fn(&str) -> Option<String>,
+) -> String {
+    let api_base = non_empty(api_base)
+        .map(str::to_string)
+        .or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
+        .unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string());
+
+    let api_base = api_base.trim_end_matches('/');
+    if api_base.ends_with(MESSAGES_PATH_SUFFIX) {
+        return api_base.to_string();
+    }
+    format!("{api_base}{MESSAGES_PATH_SUFFIX}")
+}
+
+impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
+    fn complete_url(
+        &self,
+        api_base: Option<&str>,
+        _model: &str,
+        env_lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> CoreResult<String> {
+        Ok(complete_anthropic_url(api_base, env_lookup))
+    }
+
+    fn resolve_api_key(
+        &self,
+        api_key: Option<&str>,
+        env_lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> CoreResult<String> {
+        resolve_anthropic_api_key(api_key, env_lookup)
+    }
+
+    fn auth_strategy(&self) -> MessagesAuthStrategy {
+        MessagesAuthStrategy::Header("x-api-key")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn url_defaults_to_public_anthropic_endpoint() {
+        assert_eq!(
+            complete_anthropic_url(None, &|_| None),
+            "https://api.anthropic.com/v1/messages"
+        );
+    }
+
+    #[test]
+    fn url_appends_messages_suffix_to_custom_base() {
+        assert_eq!(
+            complete_anthropic_url(Some("https://proxy.internal"), &|_| None),
+            "https://proxy.internal/v1/messages"
+        );
+    }
+
+    #[test]
+    fn url_leaves_complete_messages_endpoint_untouched() {
+        assert_eq!(
+            complete_anthropic_url(Some("https://proxy.internal/v1/messages"), &|_| None),
+            "https://proxy.internal/v1/messages"
+        );
+    }
+
+    #[test]
+    fn url_falls_back_to_env_base() {
+        let with_env = |key: &str| {
+            (key == ANTHROPIC_API_BASE_ENV).then(|| "https://env.anthropic".to_string())
+        };
+        assert_eq!(
+            complete_anthropic_url(Some("  "), &with_env),
+            "https://env.anthropic/v1/messages"
+        );
+    }
+
+    #[test]
+    fn api_key_prefers_param_then_env_then_errors() {
+        assert_eq!(
+            resolve_anthropic_api_key(Some("sk-param"), &|_| None).unwrap(),
+            "sk-param"
+        );
+        let with_env = |key: &str| (key == ANTHROPIC_API_KEY_ENV).then(|| "sk-env".to_string());
+        assert_eq!(
+            resolve_anthropic_api_key(Some("  "), &with_env).unwrap(),
+            "sk-env"
+        );
+        assert!(matches!(
+            resolve_anthropic_api_key(None, &|_| None).expect_err("missing key"),
+            CoreError::Auth(_)
+        ));
+    }
+
+    #[test]
+    fn auth_strategy_and_default_headers_match_anthropic() {
+        assert_eq!(
+            ANTHROPIC_MESSAGES_CONFIG.auth_strategy().header_name(),
+            "x-api-key"
+        );
+        assert_eq!(
+            ANTHROPIC_MESSAGES_CONFIG.default_headers(),
+            &[
+                ("anthropic-version", "2023-06-01"),
+                ("content-type", "application/json"),
+            ]
+        );
+    }
+}

litellm-rust/crates/core/src/providers/anthropic/mod.rs

@@ -0,0 +1 @@
+pub mod messages;

litellm-rust/crates/core/src/providers/azure_ai/messages/mod.rs

@@ -0,0 +1 @@
+pub mod transformation;

litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs

@@ -0,0 +1,512 @@
+use crate::error::{CoreError, CoreResult};
+use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
+use crate::messages::types::{
+    AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock,
+    MessageContent, SystemPrompt,
+};
+use crate::providers::anthropic::messages::transformation::{
+    non_empty, AnthropicMessagesConfig, ANTHROPIC_MESSAGES_CONFIG,
+};
+use serde_json::{Map, Value};
+
+const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY";
+const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE";
+const ANTHROPIC_PATH_SEGMENT: &str = "/anthropic";
+const MESSAGES_PATH_SUFFIX: &str = "/v1/messages";
+const SYSTEM_ROLE: &str = "system";
+const TEXT_BLOCK_TYPE: &str = "text";
+
+pub struct AzureAnthropicMessagesConfig {
+    anthropic: AnthropicMessagesConfig,
+}
+
+pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig =
+    AzureAnthropicMessagesConfig {
+        anthropic: ANTHROPIC_MESSAGES_CONFIG,
+    };
+
+pub fn resolve_azure_api_key(
+    api_key: Option<&str>,
+    env_lookup: &dyn Fn(&str) -> Option<String>,
+) -> CoreResult<String> {
+    non_empty(api_key)
+        .map(str::to_string)
+        .or_else(|| env_lookup(AZURE_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
+        .ok_or_else(|| {
+            CoreError::Auth(
+                "Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable"
+                    .to_string(),
+            )
+        })
+}
+
+pub fn complete_azure_anthropic_url(
+    api_base: Option<&str>,
+    env_lookup: &dyn Fn(&str) -> Option<String>,
+) -> CoreResult<String> {
+    let api_base = non_empty(api_base)
+        .map(str::to_string)
+        .or_else(|| env_lookup(AZURE_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
+        .ok_or_else(|| {
+            CoreError::Auth(
+                "Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. \
+                 Expected format: https://<resource-name>.services.ai.azure.com/anthropic"
+                    .to_string(),
+            )
+        })?;
+
+    let api_base = api_base.trim_end_matches('/');
+
+    if api_base.ends_with(MESSAGES_PATH_SUFFIX) {
+        return Ok(api_base.to_string());
+    }
+
+    let with_anthropic = match api_base.split_once(ANTHROPIC_PATH_SEGMENT) {
+        Some((prefix, _)) => format!("{prefix}{ANTHROPIC_PATH_SEGMENT}"),
+        None => format!("{api_base}{ANTHROPIC_PATH_SEGMENT}"),
+    };
+    Ok(format!("{with_anthropic}{MESSAGES_PATH_SUFFIX}"))
+}
+
+fn strip_scope_from_block(block: &mut ContentBlock) {
+    if let Some(cache_control) = block.cache_control.as_mut() {
+        cache_control.scope = None;
+    }
+}
+
+fn strip_scope_from_system(system: &mut SystemPrompt) {
+    if let SystemPrompt::Blocks(blocks) = system {
+        blocks.iter_mut().for_each(strip_scope_from_block);
+    }
+}
+
+fn strip_scope_from_message(message: &mut AnthropicMessage) {
+    if let MessageContent::Blocks(blocks) = &mut message.content {
+        blocks.iter_mut().for_each(strip_scope_from_block);
+    }
+}
+
+fn text_content_block(text: String) -> ContentBlock {
+    let extra = Map::from_iter([
+        (
+            "type".to_string(),
+            Value::String(TEXT_BLOCK_TYPE.to_string()),
+        ),
+        ("text".to_string(), Value::String(text)),
+    ]);
+    ContentBlock {
+        cache_control: None,
+        extra,
+    }
+}
+
+fn content_into_blocks(content: MessageContent) -> Vec<ContentBlock> {
+    match content {
+        MessageContent::Text(text) => vec![text_content_block(text)],
+        MessageContent::Blocks(blocks) => blocks,
+    }
+}
+
+fn system_into_blocks(system: Option<SystemPrompt>) -> Vec<ContentBlock> {
+    match system {
+        None => Vec::new(),
+        Some(SystemPrompt::Text(text)) => vec![text_content_block(text)],
+        Some(SystemPrompt::Blocks(blocks)) => blocks,
+    }
+}
+
+fn fold_system_role_messages(request: AnthropicMessagesRequest) -> AnthropicMessagesRequest {
+    if !request.messages.iter().any(|msg| msg.role == SYSTEM_ROLE) {
+        return request;
+    }
+
+    let (system_messages, chat_messages): (Vec<AnthropicMessage>, Vec<AnthropicMessage>) = request
+        .messages
+        .into_iter()
+        .partition(|msg| msg.role == SYSTEM_ROLE);
+
+    let folded_system: Vec<ContentBlock> = system_into_blocks(request.system)
+        .into_iter()
+        .chain(
+            system_messages
+                .into_iter()
+                .flat_map(|msg| content_into_blocks(msg.content)),
+        )
+        .collect();
+
+    AnthropicMessagesRequest {
+        messages: chat_messages,
+        system: (!folded_system.is_empty()).then_some(SystemPrompt::Blocks(folded_system)),
+        ..request
+    }
+}
+
+impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
+    fn complete_url(
+        &self,
+        api_base: Option<&str>,
+        _model: &str,
+        env_lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> CoreResult<String> {
+        complete_azure_anthropic_url(api_base, env_lookup)
+    }
+
+    fn resolve_api_key(
+        &self,
+        api_key: Option<&str>,
+        env_lookup: &dyn Fn(&str) -> Option<String>,
+    ) -> CoreResult<String> {
+        resolve_azure_api_key(api_key, env_lookup)
+    }
+
+    fn auth_strategy(&self) -> MessagesAuthStrategy {
+        self.anthropic.auth_strategy()
+    }
+
+    fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
+        self.anthropic.default_headers()
+    }
+
+    fn transform_request(
+        &self,
+        request: AnthropicMessagesRequest,
+    ) -> CoreResult<AnthropicMessagesRequest> {
+        let mut request = fold_system_role_messages(request);
+        if let Some(system) = request.system.as_mut() {
+            strip_scope_from_system(system);
+        }
+        request
+            .messages
+            .iter_mut()
+            .for_each(strip_scope_from_message);
+        self.anthropic.transform_request(request)
+    }
+
+    fn transform_response(
+        &self,
+        model: &str,
+        response: AnthropicMessagesResponse,
+    ) -> CoreResult<AnthropicMessagesResponse> {
+        self.anthropic.transform_response(model, response)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    fn request_from(value: serde_json::Value) -> AnthropicMessagesRequest {
+        serde_json::from_value(value).expect("valid request")
+    }
+
+    fn to_value(request: AnthropicMessagesRequest) -> serde_json::Value {
+        serde_json::to_value(request).expect("serializable request")
+    }
+
+    #[test]
+    fn url_appends_anthropic_and_messages_suffix() {
+        let url =
+            complete_azure_anthropic_url(Some("https://resource.services.ai.azure.com"), &|_| None)
+                .expect("url builds");
+        assert_eq!(
+            url,
+            "https://resource.services.ai.azure.com/anthropic/v1/messages"
+        );
+    }
+
+    #[test]
+    fn url_keeps_existing_anthropic_segment() {
+        let url = complete_azure_anthropic_url(
+            Some("https://resource.services.ai.azure.com/anthropic"),
+            &|_| None,
+        )
+        .expect("url builds");
+        assert_eq!(
+            url,
+            "https://resource.services.ai.azure.com/anthropic/v1/messages"
+        );
+    }
+
+    #[test]
+    fn url_leaves_complete_messages_endpoint_untouched() {
+        for base in [
+            "https://resource.services.ai.azure.com/anthropic/v1/messages",
+            "https://resource.services.ai.azure.com/v1/messages",
+        ] {
+            assert_eq!(
+                complete_azure_anthropic_url(Some(base), &|_| None).expect("url builds"),
+                base
+            );
+        }
+    }
+
+    #[test]
+    fn url_trims_trailing_slash_and_truncates_after_anthropic() {
+        let url = complete_azure_anthropic_url(
+            Some("https://resource.services.ai.azure.com/anthropic/extra/"),
+            &|_| None,
+        )
+        .expect("url builds");
+        assert_eq!(
+            url,
+            "https://resource.services.ai.azure.com/anthropic/v1/messages"
+        );
+    }
+
+    #[test]
+    fn url_falls_back_to_env_then_errors_when_absent() {
+        let with_env = |key: &str| {
+            (key == AZURE_API_BASE_ENV).then(|| "https://env.services.ai.azure.com".to_string())
+        };
+        assert_eq!(
+            complete_azure_anthropic_url(None, &with_env).expect("url builds"),
+            "https://env.services.ai.azure.com/anthropic/v1/messages"
+        );
+        let err = complete_azure_anthropic_url(Some("  "), &|_| None).expect_err("missing base");
+        assert!(matches!(err, CoreError::Auth(_)));
+    }
+
+    #[test]
+    fn resolve_api_key_prefers_param_then_env() {
+        assert_eq!(
+            resolve_azure_api_key(Some("sk-param"), &|_| None).unwrap(),
+            "sk-param"
+        );
+        let with_env = |key: &str| (key == AZURE_API_KEY_ENV).then(|| "sk-env".to_string());
+        assert_eq!(
+            resolve_azure_api_key(Some("  "), &with_env).unwrap(),
+            "sk-env"
+        );
+        assert!(matches!(
+            resolve_azure_api_key(None, &|_| None).expect_err("missing key"),
+            CoreError::Auth(_)
+        ));
+    }
+
+    #[test]
+    fn auth_strategy_is_x_api_key() {
+        assert_eq!(
+            AZURE_ANTHROPIC_MESSAGES_CONFIG
+                .auth_strategy()
+                .header_name(),
+            "x-api-key"
+        );
+    }
+
+    #[test]
+    fn default_headers_match_python() {
+        assert_eq!(
+            AZURE_ANTHROPIC_MESSAGES_CONFIG.default_headers(),
+            &[
+                ("anthropic-version", "2023-06-01"),
+                ("content-type", "application/json"),
+            ]
+        );
+    }
+
+    #[test]
+    fn transform_request_strips_scope_from_system_and_messages() {
+        let request = request_from(json!({
+            "model": "claude-sonnet-4-5",
+            "max_tokens": 1024,
+            "system": [
+                {
+                    "type": "text",
+                    "text": "sys",
+                    "cache_control": {"type": "ephemeral", "ttl": "1h", "scope": "global"}
+                }
+            ],
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": "hi",
+                            "cache_control": {"type": "ephemeral", "scope": "global"}
+                        },
+                        {"type": "text", "text": "no cache control"}
+                    ]
+                }
+            ]
+        }));
+
+        let transformed = to_value(
+            AZURE_ANTHROPIC_MESSAGES_CONFIG
+                .transform_request(request)
+                .expect("request transforms"),
+        );
+
+        assert_eq!(
+            transformed["system"][0]["cache_control"],
+            json!({"type": "ephemeral", "ttl": "1h"})
+        );
+        assert_eq!(
+            transformed["messages"][0]["content"][0]["cache_control"],
+            json!({"type": "ephemeral"})
+        );
+        assert_eq!(
+            transformed["messages"][0]["content"][1],
+            json!({"type": "text", "text": "no cache control"})
+        );
+    }
+
+    #[test]
+    fn transform_request_is_idempotent_and_preserves_string_system() {
+        let request = request_from(json!({
+            "model": "claude-sonnet-4-5",
+            "max_tokens": 16,
+            "system": "plain string system",
+            "messages": [{"role": "user", "content": "hi"}]
+        }));
+        let once = AZURE_ANTHROPIC_MESSAGES_CONFIG
+            .transform_request(request)
+            .expect("request transforms");
+        let twice = AZURE_ANTHROPIC_MESSAGES_CONFIG
+            .transform_request(once.clone())
+            .expect("request transforms");
+        assert_eq!(once, twice);
+        assert_eq!(to_value(once)["system"], json!("plain string system"));
+    }
+
+    #[test]
+    fn transform_request_preserves_all_supported_params() {
+        let body = json!({
+            "model": "claude-sonnet-4-5",
+            "max_tokens": 256,
+            "messages": [{"role": "user", "content": "hi"}],
+            "system": "be terse",
+            "metadata": {"user_id": "u1"},
+            "stop_sequences": ["STOP"],
+            "stream": false,
+            "temperature": 0.4,
+            "top_p": 0.9,
+            "top_k": 40,
+            "tools": [{"name": "get_weather", "input_schema": {"type": "object"}}],
+            "tool_choice": {"type": "auto"},
+            "thinking": {"type": "enabled", "budget_tokens": 1024},
+            "service_tier": "auto",
+            "container": {"id": "c1"},
+            "mcp_servers": [{"type": "url", "url": "https://mcp.example", "name": "x"}],
+            "context_management": {"edits": []},
+            "output_format": {"type": "json_schema"},
+            "output_config": {"effort": "high"},
+            "speed": "fast",
+            "inference_geo": "us",
+            "litellm_metadata": {"trace": "abc"}
+        });
+        let transformed = to_value(
+            AZURE_ANTHROPIC_MESSAGES_CONFIG
+                .transform_request(request_from(body.clone()))
+                .expect("request transforms"),
+        );
+        assert_eq!(transformed, body);
+    }
+
+    #[test]
+    fn transform_request_folds_system_role_message_into_top_level_system() {
+        let request = request_from(json!({
+            "model": "claude-sonnet-4-5",
+            "max_tokens": 256,
+            "system": [{"type": "text", "text": "base system"}],
+            "messages": [
+                {"role": "user", "content": "fix the bug"},
+                {"role": "system", "content": "Available agent types: claude"}
+            ]
+        }));
+
+        let transformed = to_value(
+            AZURE_ANTHROPIC_MESSAGES_CONFIG
+                .transform_request(request)
+                .expect("request transforms"),
+        );
+
+        assert_eq!(
+            transformed["messages"],
+            json!([{"role": "user", "content": "fix the bug"}])
+        );
+        assert_eq!(
+            transformed["system"],
+            json!([
+                {"type": "text", "text": "base system"},
+                {"type": "text", "text": "Available agent types: claude"}
+            ])
+        );
+    }
+
+    #[test]
+    fn transform_request_folds_system_role_when_no_top_level_system() {
+        let request = request_from(json!({
+            "model": "claude-sonnet-4-5",
+            "max_tokens": 256,
+            "messages": [
+                {"role": "user", "content": [{"type": "text", "text": "hi"}]},
+                {"role": "system", "content": [{"type": "text", "text": "sys block"}]}
+            ]
+        }));
+
+        let transformed = to_value(
+            AZURE_ANTHROPIC_MESSAGES_CONFIG
+                .transform_request(request)
+                .expect("request transforms"),
+        );
+
+        assert_eq!(
+            transformed["messages"],
+            json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}])
+        );
+        assert_eq!(
+            transformed["system"],
+            json!([{"type": "text", "text": "sys block"}])
+        );
+    }
+
+    #[test]
+    fn transform_request_leaves_requests_without_system_role_untouched() {
+        let body = json!({
+            "model": "claude-sonnet-4-5",
+            "max_tokens": 256,
+            "system": "be terse",
+            "messages": [
+                {"role": "user", "content": "hi"},
+                {"role": "assistant", "content": "hello"}
+            ]
+        });
+        let transformed = to_value(
+            AZURE_ANTHROPIC_MESSAGES_CONFIG
+                .transform_request(request_from(body.clone()))
+                .expect("request transforms"),
+        );
+        assert_eq!(transformed, body);
+    }
+
+    #[test]
+    fn transform_request_rejects_non_object_body() {
+        let err = serde_json::from_value::<AnthropicMessagesRequest>(json!("bad"))
+            .expect_err("non-object body should error");
+        assert!(err.is_data());
+    }
+
+    #[test]
+    fn transform_response_passes_through() {
+        let response: AnthropicMessagesResponse = serde_json::from_value(json!({
+            "id": "msg_1",
+            "type": "message",
+            "role": "assistant",
+            "content": [{"type": "text", "text": "hello"}],
+            "model": "claude-sonnet-4-5",
+            "stop_reason": "end_turn",
+            "stop_sequence": null,
+            "usage": {"input_tokens": 1, "output_tokens": 2}
+        }))
+        .expect("valid response");
+        let transformed = AZURE_ANTHROPIC_MESSAGES_CONFIG
+            .transform_response("claude-sonnet-4-5", response)
+            .expect("response transforms");
+        let value = serde_json::to_value(transformed).expect("serializable");
+        assert_eq!(value["stop_reason"], json!("end_turn"));
+        assert_eq!(value["stop_sequence"], json!(null));
+        assert_eq!(value["content"][0]["text"], json!("hello"));
+    }
+}

litellm-rust/crates/core/src/providers/azure_ai/mod.rs

@@ -1 +1,2 @@
+pub mod messages;
 pub mod ocr;

litellm-rust/crates/core/src/providers/mod.rs

@@ -1,3 +1,4 @@
+pub mod anthropic;
 pub mod azure_ai;
 pub mod mistral;
 pub mod openai;

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-rust/crates/python-bridge/src/lib.rs

@@ -1,8 +1,9 @@
 use std::time::Duration;
 
+use litellm_ai_gateway::io::messages::{messages as run_messages, MessagesRequest};
 use litellm_ai_gateway::io::ocr::{ocr as run_ocr, OcrRequest};
 use litellm_core::error::CoreError;
-use pyo3::exceptions::PyValueError;
+use pyo3::exceptions::{PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::{PyAny, PyDict};
 use serde_json::{Map, Value};
@@ -29,12 +30,23 @@ fn json_to_py(py: Python<'_>, value: Value) -> PyResult<Py<PyAny>> {
     Ok(json.call_method1("loads", (encoded,))?.unbind())
 }
 
-fn core_error_to_pyerr(py: Python<'_>, err: CoreError) -> PyErr {
+fn ocr_error_to_pyerr(py: Python<'_>, err: CoreError) -> PyErr {
     let status_code = err.public_status_code();
     let message = err.public_message();
     build_rust_ocr_error(py, &message, status_code).unwrap_or_else(|import_err| import_err)
 }
 
+fn core_error_to_pyerr(err: CoreError) -> PyErr {
+    match err {
+        CoreError::Auth(message) => PyValueError::new_err(message),
+        CoreError::InvalidProvider(_)
+        | CoreError::InvalidRequest(_)
+        | CoreError::InvalidType { .. }
+        | CoreError::MissingField(_) => PyValueError::new_err(err.to_string()),
+        other => PyRuntimeError::new_err(other.to_string()),
+    }
+}
+
 fn build_rust_ocr_error(
     py: Python<'_>,
     message: &str,
@@ -130,7 +142,7 @@ fn ocr(
 
     match result {
         Ok(value) => json_to_py(py, value),
-        Err(err) => Err(core_error_to_pyerr(py, err)),
+        Err(err) => Err(ocr_error_to_pyerr(py, err)),
     }
 }
 
@@ -172,7 +184,93 @@ fn aocr(
             litellm_call_id: None,
         })
         .await
-        .map_err(|err| Python::attach(|py| core_error_to_pyerr(py, err)))?;
+        .map_err(|err| Python::attach(|py| ocr_error_to_pyerr(py, err)))?;
+
+        Python::attach(|py| json_to_py(py, value))
+    })
+}
+
+type MarshaledMessagesInputs = (Value, Option<Map<String, Value>>, Option<Duration>);
+
+fn marshal_messages_inputs(
+    py: Python<'_>,
+    body: Py<PyAny>,
+    extra_headers: Option<Py<PyAny>>,
+    timeout_seconds: Option<f64>,
+) -> PyResult<MarshaledMessagesInputs> {
+    let body = py_to_json(py, body.bind(py))?;
+    if !body.is_object() {
+        return Err(PyValueError::new_err("body must be a dict"));
+    }
+    let extra_headers = match extra_headers {
+        Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
+        None => None,
+    };
+    Ok((body, extra_headers, optional_timeout(timeout_seconds)))
+}
+
+#[pyfunction]
+#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
+#[allow(clippy::too_many_arguments)]
+fn messages(
+    py: Python<'_>,
+    model: String,
+    body: Py<PyAny>,
+    api_key: Option<String>,
+    api_base: Option<String>,
+    custom_llm_provider: Option<String>,
+    extra_headers: Option<Py<PyAny>>,
+    timeout_seconds: Option<f64>,
+) -> PyResult<Py<PyAny>> {
+    let (body, extra_headers, timeout) =
+        marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?;
+
+    let result = gil::release_gil(py, || {
+        pyo3_async_runtimes::tokio::get_runtime().block_on(run_messages(MessagesRequest {
+            model: &model,
+            body,
+            api_key: api_key.as_deref(),
+            api_base: api_base.as_deref(),
+            custom_llm_provider: custom_llm_provider.as_deref(),
+            extra_headers,
+            timeout,
+        }))
+    });
+
+    match result {
+        Ok(value) => json_to_py(py, value),
+        Err(err) => Err(core_error_to_pyerr(err)),
+    }
+}
+
+#[pyfunction]
+#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
+#[allow(clippy::too_many_arguments)]
+fn amessages(
+    py: Python<'_>,
+    model: String,
+    body: Py<PyAny>,
+    api_key: Option<String>,
+    api_base: Option<String>,
+    custom_llm_provider: Option<String>,
+    extra_headers: Option<Py<PyAny>>,
+    timeout_seconds: Option<f64>,
+) -> PyResult<Bound<'_, PyAny>> {
+    let (body, extra_headers, timeout) =
+        marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?;
+
+    pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        let value = run_messages(MessagesRequest {
+            model: &model,
+            body,
+            api_key: api_key.as_deref(),
+            api_base: api_base.as_deref(),
+            custom_llm_provider: custom_llm_provider.as_deref(),
+            extra_headers,
+            timeout,
+        })
+        .await
+        .map_err(core_error_to_pyerr)?;
 
         Python::attach(|py| json_to_py(py, value))
     })
@@ -189,6 +287,8 @@ fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
 fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add_function(wrap_pyfunction!(ocr, module)?)?;
     module.add_function(wrap_pyfunction!(aocr, module)?)?;
+    module.add_function(wrap_pyfunction!(messages, module)?)?;
+    module.add_function(wrap_pyfunction!(amessages, module)?)?;
     module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
     Ok(())
 }

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/custom_httpx/llm_http_handler.py

@@ -152,6 +152,9 @@ if TYPE_CHECKING:
     from aiohttp import ClientSession
 
     from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
+    from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
+        AnthropicMessagesStreamingResponse,
+    )
     from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
     from litellm.types.llms.openai_evals import (
         CancelEvalResponse,
@@ -1807,6 +1810,37 @@ class BaseLLMHTTPHandler:
             },
         )
 
+        rust_messages_response = await self._maybe_rust_anthropic_messages(
+            custom_llm_provider=custom_llm_provider,
+            litellm_params=litellm_params,
+            stream=stream or False,
+            rust_stream_eligible=bool(stream) and not self._has_agentic_completion_hook(logging_obj),
+            model=model,
+            api_key=api_key,
+            api_base=api_base,
+            headers=headers,
+            request_body=request_body,
+            timeout=self._resolve_anthropic_messages_timeout(
+                litellm_params=litellm_params,
+                stream=stream or False,
+                custom_llm_provider=custom_llm_provider,
+            ),
+        )
+        if rust_messages_response is not None:
+            if stream:
+                return self._rust_anthropic_messages_fake_stream(rust_messages_response)
+            return await self._finalize_anthropic_messages_response(
+                initial_response=rust_messages_response,
+                model=model,
+                messages=messages,
+                anthropic_messages_provider_config=anthropic_messages_provider_config,
+                anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
+                logging_obj=logging_obj,
+                custom_llm_provider=custom_llm_provider,
+                api_key=api_key,
+                kwargs=kwargs,
+            )
+
         response = await self._async_post_anthropic_messages_with_http_error_retry(
             async_httpx_client=async_httpx_client,
             request_url=request_url,
@@ -1881,6 +1915,31 @@ class BaseLLMHTTPHandler:
                 logging_obj=logging_obj,
             )
 
+        return await self._finalize_anthropic_messages_response(
+            initial_response=initial_response,
+            model=model,
+            messages=messages,
+            anthropic_messages_provider_config=anthropic_messages_provider_config,
+            anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
+            logging_obj=logging_obj,
+            custom_llm_provider=custom_llm_provider,
+            api_key=api_key,
+            kwargs=kwargs,
+        )
+
+    async def _finalize_anthropic_messages_response(
+        self,
+        *,
+        initial_response: AnthropicMessagesResponse,
+        model: str,
+        messages: list[dict],
+        anthropic_messages_provider_config: BaseAnthropicMessagesConfig,
+        anthropic_messages_optional_request_params: dict,
+        logging_obj: LiteLLMLoggingObj,
+        custom_llm_provider: str,
+        api_key: str | None,
+        kwargs: dict,
+    ) -> AnthropicMessagesResponse | AsyncIterator:
         # Inject api_key into kwargs so follow-up calls in agentic hooks can
         # authenticate. api_key is a named param here (not in kwargs), so
         # _prepare_followup_kwargs would miss it otherwise.
@@ -1904,6 +1963,70 @@ class BaseLLMHTTPHandler:
             "anthropic_messages",
         )
 
+    @staticmethod
+    async def _maybe_rust_anthropic_messages(
+        *,
+        custom_llm_provider: str,
+        litellm_params: GenericLiteLLMParams,
+        stream: bool,
+        rust_stream_eligible: bool,
+        model: str,
+        api_key: str | None,
+        api_base: str | None,
+        headers: dict,
+        request_body: dict,
+        timeout: float | httpx.Timeout | None,
+    ) -> AnthropicMessagesResponse | None:
+        if custom_llm_provider != "azure_ai" or litellm_params.get("rust") is not True:
+            return None
+        if stream and not rust_stream_eligible:
+            return None
+
+        from litellm.rust_bridge import messages as rust_messages_bridge
+
+        upstream_body = {key: value for key, value in request_body.items() if key != "stream"}
+        try:
+            rust_response = await rust_messages_bridge.amessages(
+                model=model,
+                body=upstream_body,
+                api_key=api_key,
+                api_base=api_base,
+                custom_llm_provider=custom_llm_provider,
+                extra_headers=headers,
+                timeout=timeout,
+            )
+        except Exception as rust_error:  # noqa: BLE001  # rollout-safety fallback: any Rust bridge failure must fall back to the Python path
+            verbose_logger.debug(
+                "Rust Anthropic messages bridge raised %s; falling back to Python path",
+                type(rust_error).__name__,
+            )
+            return None
+        if rust_response is None:
+            return None
+
+        response_obj = cast(AnthropicMessagesResponse, dict(rust_response))
+        response_obj["_hidden_params"] = {"additional_headers": {"x-litellm-rust": "true"}}
+        return response_obj
+
+    @staticmethod
+    def _rust_anthropic_messages_fake_stream(
+        rust_response: AnthropicMessagesResponse,
+    ) -> "AnthropicMessagesStreamingResponse":
+        from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
+            FakeAnthropicMessagesStreamIterator,
+        )
+        from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
+            AnthropicMessagesStreamHiddenParams,
+            AnthropicMessagesStreamingResponse,
+        )
+
+        completion_stream = cast(AsyncIterator[bytes], FakeAnthropicMessagesStreamIterator(response=rust_response))
+        hidden_params = AnthropicMessagesStreamHiddenParams(additional_headers={"x-litellm-rust": "true"})
+        return AnthropicMessagesStreamingResponse(
+            completion_stream=completion_stream,
+            hidden_params=hidden_params,
+        )
+
     def anthropic_messages_handler(
         self,
         model: str,

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/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/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/rust_bridge/messages.py

@@ -0,0 +1,135 @@
+"""Thin Python wrapper for the native Rust Anthropic Messages bridge."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Awaitable, Final, Protocol, Union, cast
+
+import httpx
+
+from litellm.rust_bridge.timeouts import timeout_to_seconds
+
+
+class RustMessages(Protocol):
+    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]:
+        raise NotImplementedError
+
+
+class RustAmessages(Protocol):
+    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,
+    ) -> Awaitable[dict[str, object]]:
+        raise NotImplementedError
+
+
+class _Unset:
+    pass
+
+
+_UNSET: Final[_Unset] = _Unset()
+
+
+@dataclass(slots=True)
+class _RustMessagesState:
+    messages: RustMessages | None = None
+    amessages: RustAmessages | None = None
+
+
+_STATE: Final[_RustMessagesState] = _RustMessagesState()
+
+
+def set_rust_messages(
+    *,
+    messages: RustMessages | None | _Unset = _UNSET,
+    amessages: RustAmessages | None | _Unset = _UNSET,
+) -> None:
+    if not isinstance(messages, _Unset):
+        _STATE.messages = messages
+    if not isinstance(amessages, _Unset):
+        _STATE.amessages = amessages
+
+
+def load_rust_messages() -> RustMessages | None:
+    if _STATE.messages is not None:
+        return _STATE.messages
+    from litellm.rust_bridge import get_native_bridge
+
+    native_bridge = get_native_bridge()
+    if native_bridge is None:
+        return None
+    return cast(RustMessages, getattr(native_bridge, "messages", None))
+
+
+def load_rust_amessages() -> RustAmessages | None:
+    if _STATE.amessages is not None:
+        return _STATE.amessages
+    from litellm.rust_bridge import get_native_bridge
+
+    native_bridge = get_native_bridge()
+    if native_bridge is None:
+        return None
+    return cast(RustAmessages, getattr(native_bridge, "amessages", None))
+
+
+def messages(
+    *,
+    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: Union[float, httpx.Timeout] | None,
+) -> dict[str, object] | None:
+    rust_messages = load_rust_messages()
+    if rust_messages is None:
+        return None
+    return rust_messages(
+        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_to_seconds(timeout),
+    )
+
+
+async def amessages(
+    *,
+    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: Union[float, httpx.Timeout] | None,
+) -> dict[str, object] | None:
+    rust_amessages = load_rust_amessages()
+    if rust_amessages is None:
+        return None
+    return await rust_amessages(
+        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_to_seconds(timeout),
+    )

litellm/rust_bridge/ocr.py

@@ -2,10 +2,15 @@
 
 from __future__ import annotations
 
-from typing import Any, Awaitable, Final, Protocol, Union, cast
+from typing import TYPE_CHECKING, Awaitable, Final, Protocol, Union, cast
 
 import httpx
 
+from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds
+
+if TYPE_CHECKING:
+    from litellm.rust_bridge.messages import RustAmessages, RustMessages
+
 
 class RustOcrError(Exception):
     def __init__(self, message: str, status_code: int | None = None) -> None:
@@ -69,11 +74,26 @@ def use_litellm_rust(
     *,
     ocr: RustOcr | None | _Unset = _UNSET,
     aocr: RustAocr | None | _Unset = _UNSET,
+    messages: RustMessages | None | _Unset = _UNSET,
+    amessages: RustAmessages | None | _Unset = _UNSET,
 ) -> None:
-    if not enabled:
-        _set_rust_ocr_bridge(ocr=None, aocr=None)
+    configuring_ocr = not isinstance(ocr, _Unset) or not isinstance(aocr, _Unset)
+    configuring_messages = not isinstance(messages, _Unset) or not isinstance(amessages, _Unset)
+    if configuring_ocr or not configuring_messages:
+        if enabled:
+            _set_rust_ocr_bridge(ocr=ocr, aocr=aocr)
+        else:
+            _set_rust_ocr_bridge(ocr=None, aocr=None)
+    if not configuring_messages:
         return
-    _set_rust_ocr_bridge(ocr=ocr, aocr=aocr)
+    from litellm.rust_bridge.messages import set_rust_messages
+
+    if not isinstance(messages, _Unset) and not isinstance(amessages, _Unset):
+        set_rust_messages(messages=messages, amessages=amessages)
+    elif not isinstance(messages, _Unset):
+        set_rust_messages(messages=messages)
+    elif not isinstance(amessages, _Unset):
+        set_rust_messages(amessages=amessages)
 
 
 def rust_ocr_enabled() -> bool:
@@ -102,22 +122,14 @@ def load_rust_aocr() -> RustAocr | None:
     return cast(RustAocr, getattr(native_bridge, "aocr", None))
 
 
-def _timeout_to_seconds(timeout: Union[float, httpx.Timeout] | None) -> float | None:
-    if timeout is None:
-        return None
-    if isinstance(timeout, httpx.Timeout):
-        return timeout.read
-    return float(timeout)
-
-
 def ocr(
     *,
     model: str,
-    document: dict[str, Any],
+    document: dict[str, object],
     api_key: str | None,
     api_base: str | None,
     custom_llm_provider: str | None,
-    extra_headers: dict[str, Any] | None,
+    extra_headers: dict[str, object] | None,
     optional_params: dict[str, object],
     timeout: Union[float, httpx.Timeout] | None,
 ) -> dict[str, object] | None:
@@ -126,11 +138,11 @@ def ocr(
         return None
     return rust_ocr(
         model=model,
-        document=cast(dict[str, object], document),
+        document=document,
         api_key=api_key,
         api_base=api_base,
         custom_llm_provider=custom_llm_provider,
-        extra_headers=cast(dict[str, object] | None, extra_headers),
+        extra_headers=extra_headers,
         optional_params=optional_params,
         timeout_seconds=_timeout_to_seconds(timeout),
     )
@@ -139,11 +151,11 @@ def ocr(
 async def aocr(
     *,
     model: str,
-    document: dict[str, Any],
+    document: dict[str, object],
     api_key: str | None,
     api_base: str | None,
     custom_llm_provider: str | None,
-    extra_headers: dict[str, Any] | None,
+    extra_headers: dict[str, object] | None,
     optional_params: dict[str, object],
     timeout: Union[float, httpx.Timeout] | None,
 ) -> dict[str, object] | None:
@@ -152,11 +164,11 @@ async def aocr(
         return None
     return await rust_aocr(
         model=model,
-        document=cast(dict[str, object], document),
+        document=document,
         api_key=api_key,
         api_base=api_base,
         custom_llm_provider=custom_llm_provider,
-        extra_headers=cast(dict[str, object] | None, extra_headers),
+        extra_headers=extra_headers,
         optional_params=optional_params,
         timeout_seconds=_timeout_to_seconds(timeout),
     )

litellm/rust_bridge/timeouts.py

@@ -0,0 +1,15 @@
+"""Shared timeout conversion for the native Rust bridges."""
+
+from __future__ import annotations
+
+from typing import Union
+
+import httpx
+
+
+def timeout_to_seconds(timeout: Union[float, httpx.Timeout] | None) -> float | None:
+    if timeout is None:
+        return None
+    if isinstance(timeout, httpx.Timeout):
+        return timeout.read
+    return float(timeout)

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"

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)