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