diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index da788bf1ce3..d094c98f5ec 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 14076 + "limit": 14074 }, "reportArgumentType": { "limit": 2216 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 4128 + "limit": 4125 }, "reportFunctionMemberAccess": { "limit": 7 @@ -108,10 +108,10 @@ "limit": 38350 }, "reportUnknownParameterType": { - "limit": 19626 + "limit": 19625 }, "reportUnknownVariableType": { - "limit": 29890 + "limit": 29877 }, "reportUnnecessaryCast": { "limit": 111 @@ -138,7 +138,7 @@ "limit": 138 }, "reportUnusedImport": { - "limit": 543 + "limit": 542 }, "reportUnusedVariable": { "limit": 137 diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py index 1d3268da9a0..190141470df 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py @@ -39,9 +39,9 @@ async def available_enterprise_users( if not premium_user: # check if SSO is enabled - show 5 user limit - from litellm.proxy.auth.auth_utils import _has_user_setup_sso + from litellm.proxy.auth.auth_utils import has_user_setup_sso - if _has_user_setup_sso(): + if has_user_setup_sso(): premium_user_data = EnterpriseLicenseData( max_users=5, ) diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs deleted file mode 100644 index 140bc8aeea8..00000000000 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs +++ /dev/null @@ -1,46 +0,0 @@ -use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig; -use litellm_core::error::Error; -use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG; -use serde_json::{Map, Value}; -use std::collections::BTreeMap; - -pub(super) fn audio_transcription_provider_config( - provider: &str, -) -> Option<&'static dyn AudioTranscriptionProviderConfig> { - match provider { - "bedrock" => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG), - _ => None, - } -} - -pub(super) fn string_headers( - headers: Option>, -) -> Result, Error> { - headers - .unwrap_or_default() - .into_iter() - .map(|(key, value)| { - value - .as_str() - .map(|value| (key.clone(), value.to_string())) - .ok_or_else(|| { - Error::InvalidRequest(format!( - "audio transcription extra_headers.{key} must be a string" - )) - }) - }) - .collect() -} - -pub(super) fn has_header(headers: &BTreeMap, name: &str) -> bool { - headers.keys().any(|key| key.eq_ignore_ascii_case(name)) -} - -pub(super) fn truncate_error_body(body: &str) -> String { - let truncated: String = body.chars().take(256).collect(); - if truncated.chars().count() == body.chars().count() { - truncated - } else { - format!("{truncated}... (truncated)") - } -} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs deleted file mode 100644 index 1bdd4ae72a2..00000000000 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs +++ /dev/null @@ -1,84 +0,0 @@ -use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth; -use litellm_core::error::Error; -use litellm_core::providers::bedrock::audio_transcription::aws_auth_config; -use litellm_core::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post}; -use serde_json::Value; -use std::time::SystemTime; - -use super::common_utils::truncate_error_body; -use super::types::ProviderAudioTranscriptionRequest; -use crate::client::http_client; - -pub(crate) async fn execute_audio_transcription_provider_call( - request: ProviderAudioTranscriptionRequest, -) -> Result { - let body = serde_json::to_vec(&request.body) - .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?; - let mut request_builder = http_client().post(&request.url).body(body.clone()); - 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(|error| Error::Network(error.to_string()))?; - let status = response.status(); - let text = response - .text() - .await - .map_err(|error| Error::Network(error.to_string()))?; - if !status.is_success() { - return Err(Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }); - } - let response_json: Value = serde_json::from_str(&text) - .map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?; - Ok(request - .config - .transform_transcription_response(&request.model, response_json)? - .into_json()) -} - -pub(crate) async fn sign_request( - request: &ProviderAudioTranscriptionRequest, - optional_params: &serde_json::Map, -) -> Result { - let env_lookup = environment_lookup; - let auth = request - .config - .auth_strategy(&request.model, optional_params, &env_lookup)?; - let body = serde_json::to_vec(&request.body) - .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?; - let mut headers = super::common_utils::string_headers(None)?; - headers.insert("Content-Type".to_string(), "application/json".to_string()); - headers.extend(request.upstream_headers.iter().cloned()); - match auth { - AudioTranscriptionAuth::Bearer => {} - AudioTranscriptionAuth::AwsSigV4 { region, .. } => { - let credentials = - resolve_credentials(aws_auth_config(optional_params, &env_lookup), &env_lookup) - .await?; - headers.extend(sign_bedrock_post( - &request.url, - &body, - &headers, - ®ion, - &credentials, - SystemTime::now(), - )?); - } - } - Ok(ProviderAudioTranscriptionRequest { - upstream_headers: headers.into_iter().collect(), - ..request.clone() - }) -} - -pub(super) fn environment_lookup(key: &str) -> Option { - std::env::var(key).ok() -} diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index 5e1240de759..dbe2d3a325b 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -1,13 +1,14 @@ -use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth; +use litellm_core::audio_transcription::{ + AudioTranscriptionRequest as CoreAudioTranscriptionRequest, ProviderAudioTranscriptionRequest, + prepare_audio_transcription_provider_call, +}; use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use litellm_core::error::Error; use serde_json::{Map, Value, json}; use std::future::Future; use std::pin::Pin; -use super::common_utils::{audio_transcription_provider_config, has_header, string_headers}; -use super::handler::sign_request; -use super::types::{PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; +use super::types::PreparedAudioTranscriptionRequest; use crate::integrations::custom_guardrail::{ CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest, }; @@ -88,46 +89,29 @@ impl AudioTranscriptionLifecycleHooks { &self, request: PreparedAudioTranscriptionRequest, ) -> Result { - let config = audio_transcription_provider_config(&request.custom_llm_provider) - .ok_or_else(|| Error::InvalidProvider(request.custom_llm_provider.clone()))?; - let env_lookup = super::handler::environment_lookup; - let headers = string_headers(request.extra_headers)?; - let url = config.complete_url( - request.api_base.as_deref(), - &request.model, - &request.optional_params, - &env_lookup, - )?; - let filtered_params = config.map_transcription_params(&request.optional_params); - let body = config.transform_transcription_request( - &request.model, - request.audio, - filtered_params, - )?; - let auth = config.auth_strategy(&request.model, &request.optional_params, &env_lookup)?; - let mut upstream_headers = headers.into_iter().collect::>(); - if matches!(auth, AudioTranscriptionAuth::Bearer) - && !has_header( - &upstream_headers - .iter() - .cloned() - .collect::>(), - "authorization", - ) - && let Some(api_key) = request.api_key.as_deref() - { - upstream_headers.push(("Authorization".to_string(), format!("Bearer {api_key}"))); - } - let provider_request = ProviderAudioTranscriptionRequest { - model: request.model, - config, - url, - body: body.body, - upstream_headers, - timeout: request.timeout, - }; - let provider_request = self.run_during_call_guardrails(provider_request).await?; - sign_request(&provider_request, &request.optional_params).await + let PreparedAudioTranscriptionRequest { + model, + custom_llm_provider, + audio, + api_key, + api_base, + extra_headers, + optional_params, + timeout, + .. + } = request; + let provider_request = + prepare_audio_transcription_provider_call(CoreAudioTranscriptionRequest { + model: &model, + audio, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: Some(&custom_llm_provider), + extra_headers, + optional_params, + timeout, + })?; + self.run_during_call_guardrails(provider_request).await } async fn run_during_call_guardrails( @@ -142,10 +126,10 @@ impl AudioTranscriptionLifecycleHooks { .run_during_call( &guardrail_context(&self.request_metadata), GuardrailRequest::new(json!({ - "model": request.model, - "custom_llm_provider": "bedrock", - "url": request.url, - "body": request.body, + "model": request.model(), + "custom_llm_provider": request.custom_llm_provider(), + "url": request.url(), + "body": request.body(), })), ) .await @@ -158,7 +142,7 @@ impl AudioTranscriptionLifecycleHooks { let body = data.remove("body").ok_or_else(|| { Error::InvalidRequest("audio transcription guardrail removed body".to_string()) })?; - Ok(ProviderAudioTranscriptionRequest { body, ..request }) + Ok(request.with_body(body)) } fn logging_payload( diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs index 3983846d7b6..03d621b8414 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs @@ -1,16 +1,14 @@ use litellm_core::Error; +use litellm_core::audio_transcription::execute_audio_transcription_provider_call; use litellm_core::call_lifecycle::CallLifecycle; use serde_json::Value; -mod common_utils; -mod handler; mod hooks; mod prepare; mod types; pub use types::AudioTranscriptionRequest; -use handler::execute_audio_transcription_provider_call; use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call}; pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs index 9697aa98b0a..b470638264e 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs @@ -1,7 +1,6 @@ use std::sync::Arc; use std::time::Duration; -use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig; use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest}; use serde_json::{Map, Value}; @@ -46,13 +45,3 @@ impl CallLifecycleRequest for PreparedAudioTranscriptionRequest { ) } } - -#[derive(Clone)] -pub(crate) struct ProviderAudioTranscriptionRequest { - pub(crate) model: String, - pub(crate) config: &'static dyn AudioTranscriptionProviderConfig, - 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/audio_transcription/client.rs b/litellm-rust/crates/core/src/audio_transcription/client.rs new file mode 100644 index 00000000000..0e612628dc6 --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/client.rs @@ -0,0 +1,14 @@ +use std::sync::OnceLock; +use std::time::Duration; + +use crate::constants::AUDIO_TRANSCRIPTION_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(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)) + .build() + .unwrap_or_else(|_| reqwest::Client::new()) + }) +} diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs new file mode 100644 index 00000000000..30ba0da5e68 --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -0,0 +1,91 @@ +use serde_json::Value; + +use crate::error::Error; +use crate::http_utils::truncate_error_body; + +use super::client::http_client; +use super::types::ProviderAudioTranscriptionRequest; + +pub async fn execute_audio_transcription_provider_call( + request: ProviderAudioTranscriptionRequest, +) -> Result { + let body = serde_json::to_vec(&request.body) + .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?; + let headers = signed_headers(&request, &body).await?; + let mut request_builder = http_client().post(&request.url).body(body); + for (key, value) in 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(|error| Error::Network(error.to_string()))?; + let status = response.status(); + let text = response + .text() + .await + .map_err(|error| Error::Network(error.to_string()))?; + if !status.is_success() { + return Err(Error::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + }); + } + let response_json = serde_json::from_str(&text) + .map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?; + Ok(request + .config + .transform_transcription_response(&request.model, response_json)? + .into_json()) +} + +#[cfg(feature = "bedrock-auth")] +async fn signed_headers( + request: &ProviderAudioTranscriptionRequest, + body: &[u8], +) -> Result, Error> { + use std::collections::BTreeMap; + use std::time::SystemTime; + + use crate::audio_transcription::transformation::AudioTranscriptionAuth; + use crate::providers::bedrock::audio_transcription::aws_auth_config; + use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post}; + + let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else { + return Ok(request.upstream_headers.clone()); + }; + let env_lookup = |key: &str| std::env::var(key).ok(); + let credentials = resolve_credentials( + aws_auth_config(&request.optional_params, &env_lookup), + &env_lookup, + ) + .await?; + let unsigned: BTreeMap = request.upstream_headers.iter().cloned().collect(); + let signature = sign_bedrock_post( + &request.url, + body, + &unsigned, + region, + &credentials, + SystemTime::now(), + )?; + Ok(unsigned.into_iter().chain(signature).collect()) +} + +#[cfg(not(feature = "bedrock-auth"))] +async fn signed_headers( + request: &ProviderAudioTranscriptionRequest, + _body: &[u8], +) -> Result, Error> { + use crate::audio_transcription::transformation::AudioTranscriptionAuth; + + match request.auth { + AudioTranscriptionAuth::AwsSigV4 { .. } => Err(Error::Unsupported( + "AWS SigV4 requires the bedrock-auth feature", + )), + AudioTranscriptionAuth::Bearer => Ok(request.upstream_headers.clone()), + } +} diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index ec2fbb969a6..b71748082bf 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -1,2 +1,20 @@ +use crate::Error; +mod client; +mod handler; +mod prepare; pub mod transformation; pub mod types; + +use serde_json::Value; + +pub use handler::execute_audio_transcription_provider_call; +pub use prepare::prepare_audio_transcription_provider_call; +pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; + +pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { + execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?) + .await +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs new file mode 100644 index 00000000000..6288e96b380 --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -0,0 +1,72 @@ +use crate::error::Error; +use crate::http_utils::{has_header, string_headers}; +#[cfg(feature = "bedrock-auth")] +use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG; +use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; + +use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig}; +use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; + +fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> { + #[cfg(feature = "bedrock-auth")] + if provider == "bedrock" { + return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG); + } + let _ = provider; + None +} + +pub fn prepare_audio_transcription_provider_call( + request: AudioTranscriptionRequest<'_>, +) -> Result { + 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(|| { + Error::InvalidProvider( + "unable to resolve custom_llm_provider for audio transcription request".to_string(), + ) + })?; + let model = provider_info.model.to_string(); + let config = provider_config(provider_info.custom_llm_provider) + .ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?; + let env_lookup = |key: &str| std::env::var(key).ok(); + let mut headers = string_headers("audio transcription", request.extra_headers)?; + let auth = config.auth_strategy(&model, &request.optional_params, &env_lookup)?; + if matches!(auth, AudioTranscriptionAuth::Bearer) + && !has_header(&headers, "authorization") + && let Some(api_key) = request.api_key + { + headers.push(("Authorization".to_string(), format!("Bearer {api_key}"))); + } + if !has_header(&headers, "content-type") { + headers.push(("Content-Type".to_string(), "application/json".to_string())); + } + let url = config.complete_url( + request.api_base, + &model, + &request.optional_params, + &env_lookup, + )?; + let filtered_params = config.map_transcription_params(&request.optional_params); + let transformed = + config.transform_transcription_request(&model, request.audio, filtered_params)?; + Ok(ProviderAudioTranscriptionRequest { + model, + custom_llm_provider: provider_info.custom_llm_provider.to_string(), + config, + url, + body: transformed.body, + upstream_headers: headers, + auth, + #[cfg(feature = "bedrock-auth")] + optional_params: request.optional_params, + timeout: request.timeout, + }) +} diff --git a/litellm-rust/crates/core/src/audio_transcription/tests.rs b/litellm-rust/crates/core/src/audio_transcription/tests.rs new file mode 100644 index 00000000000..263d63337b0 --- /dev/null +++ b/litellm-rust/crates/core/src/audio_transcription/tests.rs @@ -0,0 +1,50 @@ +use std::io::{Read, Write}; +use std::net::TcpListener; +use std::thread; + +use serde_json::{Map, json}; + +use super::audio_transcription; +use super::types::AudioTranscriptionRequest; + +#[tokio::test] +async fn bedrock_request_is_signed_and_contains_audio() { + let listener = TcpListener::bind("127.0.0.1:0").expect("listener"); + let address = listener.local_addr().expect("address"); + let server = thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("connection"); + let mut request = Vec::new(); + let mut buffer = [0_u8; 16_384]; + let count = stream.read(&mut buffer).expect("request"); + request.extend_from_slice(&buffer[..count]); + let request = String::from_utf8_lossy(&request); + assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse")); + assert!(request.contains("authorization: AWS4-HMAC-SHA256")); + assert!(request.contains("x-amz-date:")); + assert!(request.contains("\"bytes\":\"AQI=\"")); + assert!(request.contains("Transcribe the audio. Respond with only the transcript.")); + let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}"; + stream.write_all(response).expect("response"); + }); + + let optional_params = Map::from_iter([ + ("aws_access_key_id".to_string(), json!("access-key")), + ("aws_secret_access_key".to_string(), json!("secret-key")), + ("aws_region_name".to_string(), json!("us-east-1")), + ]); + let api_base = format!("http://{address}"); + let response = audio_transcription(AudioTranscriptionRequest { + model: "mistral.voxtral-mini-3b-2507", + audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}), + api_key: None, + api_base: Some(&api_base), + custom_llm_provider: Some("bedrock"), + extra_headers: None, + optional_params, + timeout: None, + }) + .await + .expect("transcription"); + assert_eq!(response, json!({"text": "hello"})); + server.join().expect("server"); +} diff --git a/litellm-rust/crates/core/src/audio_transcription/types.rs b/litellm-rust/crates/core/src/audio_transcription/types.rs index 3a9e1ecd88c..559d7837027 100644 --- a/litellm-rust/crates/core/src/audio_transcription/types.rs +++ b/litellm-rust/crates/core/src/audio_transcription/types.rs @@ -1,5 +1,56 @@ +use std::time::Duration; + use serde::{Deserialize, Serialize}; -use serde_json::Value; +use serde_json::{Map, Value}; + +use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig}; + +pub struct AudioTranscriptionRequest<'a> { + pub model: &'a str, + pub audio: 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 optional_params: Map, + pub timeout: Option, +} + +#[derive(Clone)] +pub struct ProviderAudioTranscriptionRequest { + pub(super) model: String, + pub(super) custom_llm_provider: String, + pub(super) config: &'static dyn AudioTranscriptionProviderConfig, + pub(super) url: String, + pub(super) body: Value, + pub(super) upstream_headers: Vec<(String, String)>, + pub(super) auth: AudioTranscriptionAuth, + #[cfg(feature = "bedrock-auth")] + pub(super) optional_params: Map, + pub(super) timeout: Option, +} + +impl ProviderAudioTranscriptionRequest { + pub fn model(&self) -> &str { + &self.model + } + + pub fn custom_llm_provider(&self) -> &str { + &self.custom_llm_provider + } + + pub fn url(&self) -> &str { + &self.url + } + + pub fn body(&self) -> &Value { + &self.body + } + + pub fn with_body(self, body: Value) -> Self { + Self { body, ..self } + } +} #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct AudioTranscriptionRequestData { diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index e1ac0a4fc8f..a73961060eb 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -30,6 +30,8 @@ pub(crate) const CHAT_COMPLETIONS_TIMEOUT_SECS: u64 = 600; /// Connect timeout for chat completions provider calls, in seconds. pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10; +pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600; + /// `object` field every non-streaming chat completion response carries. pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion"; diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 68aa9436b15..2e2624acbe1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,11 +1,11 @@ use std::collections::HashMap; use std::time::Duration; -use litellm_ai_gateway::io::audio_transcription::{ - AudioTranscriptionRequest, audio_transcription as run_audio_transcription, -}; use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; +use litellm_core::audio_transcription::{ + AudioTranscriptionRequest, audio_transcription as run_audio_transcription, +}; use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse}; use litellm_core::chat_completions::{ chat_completions as run_chat_completions, chat_completions_decline_reason, @@ -325,10 +325,6 @@ fn transcription( extra_headers, optional_params, timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, }, )) }); @@ -369,10 +365,6 @@ fn atranscription( extra_headers, optional_params, timeout, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: Default::default(), - litellm_call_id: None, }) .await .map_err(core_error_to_pyerr)?; diff --git a/litellm/__init__.py b/litellm/__init__.py index 4eeececdb7e..61794dabddc 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -29,7 +29,7 @@ def _dev_env_hot_reload_enabled() -> bool: if os.getenv("LITELLM_MODE", "DEV") == "DEV": _dotenv.load_dotenv(override=_dev_env_hot_reload_enabled()) -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import ( Any, Callable, @@ -490,6 +490,7 @@ public_mcp_hub_strict_whitelist: bool = True public_model_groups: Optional[List[str]] = None public_agent_groups: Optional[List[str]] = None agent_search_embedding_model: Optional[str] = None +mcp_tool_search: Optional[Mapping[str, object]] = None # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]]) # New format: { "displayName": { "url": "...", "index": 0 } } # Old format: { "displayName": "url" } (for backward compatibility) diff --git a/litellm/_logging.py b/litellm/_logging.py index 9435562f890..14cda772234 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -6,6 +6,7 @@ import sys from datetime import datetime from logging import Formatter from typing import Any, Final, TextIO +from urllib.parse import unquote import litellm from litellm.constants import ( @@ -146,6 +147,72 @@ class SecretRedactionFilter(logging.Filter): _secret_filter: Final = SecretRedactionFilter() +_MAX_SCRUBBED_ACCESS_ARG: Final = 512 + +_REDACTION_PLACEHOLDER: Final = "REDACTED" + + +def _hides_a_credential(value: str) -> bool: + """Whether *value* only looks clean until it is percent-decoded.""" + decoded: Final = unquote(value) + return _redact_string(decoded) != decoded + + +def _drop_encoded_credential(scrubbed: str) -> str: + """Drop the part of a request target that only decoding shows to be a secret. + + The request parser decodes query names and values, so `?k%65y=sk%2D...` is a + working credential that the patterns, which match literal text, do not see. + The decoded text is never logged back: it can carry a newline, and forging + log lines is not a trade worth making for a readable request target. + """ + path, separator, _query = scrubbed.partition("?") + if _hides_a_credential(path): + return _REDACTION_PLACEHOLDER + if separator and _hides_a_credential(scrubbed): + return f"{path}?{_REDACTION_PLACEHOLDER}" + return scrubbed + + +def _scrub_access_arg(value: str) -> str: + """Redact one access-log positional arg, bounding the scanned length. + + The request target is the only input to the secret regex an unauthenticated + caller controls end to end, so it is cut back to a whole query parameter + before it is scanned; a half-parameter would be too short to match its + pattern and would then be logged raw. + """ + if len(value) <= _MAX_SCRUBBED_ACCESS_ARG: + return _drop_encoded_credential(_redact_string(value)) + head: Final = value[:_MAX_SCRUBBED_ACCESS_ARG] + kept: Final = head[: max(head.rfind("?"), head.rfind("&"))] if "?" in head else head + scrubbed: Final = _drop_encoded_credential(_redact_string(kept)) + return f"{scrubbed}... ({len(value) - len(kept)} more chars truncated) ..." + + +class AccessLogRedactionFilter(logging.Filter): + """Scrubs known secret/credential patterns from HTTP access-log records. + + uvicorn's AccessFormatter unpacks ``record.args`` as a five-element tuple at + emit time, so SecretRedactionFilter cannot be reused here: it collapses the + record into ``record.msg`` and clears the args, and the formatter then raises. + """ + + def filter(self, record: logging.LogRecord) -> bool: + if not _ENABLE_SECRET_REDACTION: + return True + if isinstance(record.args, tuple) and record.args: + record.args = tuple( # rebind-ok: a Filter scrubs records in place + _scrub_access_arg(arg) if isinstance(arg, str) else arg for arg in record.args + ) + return True + # No positional args means everything is in msg, where collapsing is correct. + return _secret_filter.filter(record) + + +_access_log_filter: Final = AccessLogRedactionFilter() + + def _get_max_string_length_stdout_log() -> int: """Read the limit per record so a value loaded later via proxy config environment_variables is honored.""" @@ -553,6 +620,14 @@ _REDACTED_THIRD_PARTY_LOGGERS: Final[tuple[str, ...]] = ( "uvicorn.error", ) +# Access loggers, which emit the full request target, so a credential passed as a +# query parameter (e.g. `/key/info?key=`) lands on stdout verbatim. uvicorn.access +# covers uvicorn.run, --run_gunicorn (its worker_class is UvicornWorker, so the +# access line is still uvicorn's) and an embedding host app. --run_hypercorn and +# --run_granian log through their own loggers in their own record shapes, and +# both ship with access logging off. +_REDACTED_ACCESS_LOGGERS: Final[tuple[str, ...]] = ("uvicorn.access",) + def _redact_third_party_loggers() -> None: """Extend secret redaction to records litellm does not emit directly. @@ -575,6 +650,8 @@ def _redact_third_party_loggers() -> None: """ for name in _REDACTED_THIRD_PARTY_LOGGERS: logging.getLogger(name).addFilter(_secret_filter) + for name in _REDACTED_ACCESS_LOGGERS: + logging.getLogger(name).addFilter(_access_log_filter) # Call the suppression function diff --git a/litellm/constants.py b/litellm/constants.py index 1c1939bd350..c7b74e176db 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1742,6 +1742,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ "anthropic_prompt_caching_ttl", "max_ui_session_budget", "budget_rollover", + "mcp_tool_search", ] SPECIAL_LITELLM_AUTH_TOKEN: Final = ["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/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 545b0f40018..3519240dda9 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -755,6 +755,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): def record_gateway_injection( request_kwargs: Mapping[str, object], added: int, + injected_for_every_deployment: bool = False, ) -> None: """Name the deployment whose payload the gateway, not the client, put breakpoints on. @@ -771,7 +772,16 @@ class AnthropicCacheControlHook(CustomPromptManagement): A pass that runs before a deployment is chosen, which is what the proxy does for prompt templates, injects into the payload every leg goes on to send, so it marks - the request for all of them rather than for one. + the request for all of them rather than for one. Such a pass says so with + ``injected_for_every_deployment`` instead of relying on the shape of + ``request_kwargs``: the router's prompt-management factory stamps a provisional + deployment's ``model_info`` into kwargs before the prompt pass runs, and billing + the request through any other deployment would silently drop the credit. An + every-deployment mark, once written, also never narrows: a later per-leg stamp + (the Bedrock converse tool_config one included) describes one leg of a payload + every leg sends, so narrowing to it would uncredit whichever leg gets billed + after a failover. Both losses are fail-closed under-crediting, which is why the + guard only protects the sentinel and per-leg marks still overwrite each other. Only what this pass actually placed counts. A ``tool_config`` point is placed by the Bedrock converse transform, and only when the request carries tools, so the @@ -801,13 +811,19 @@ class AnthropicCacheControlHook(CustomPromptManagement): ), None, ) - if bucket is not None: - model_info: Final = request_kwargs.get("model_info") - bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = ( - model_info.get("id", GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT) - if isinstance(model_info, dict) - else GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT - ) + if bucket is None: + return + if bucket.get(GATEWAY_INJECTED_CACHE_METADATA_KEY) == GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT: + return + if injected_for_every_deployment: + bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT + return + model_info: Final = request_kwargs.get("model_info") + bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = ( + model_info.get("id", GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT) + if isinstance(model_info, dict) + else GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT + ) @staticmethod def maybe_inject_cache_control( diff --git a/litellm/integrations/otel/langfuse_logger.py b/litellm/integrations/otel/langfuse_logger.py new file mode 100644 index 00000000000..ed47533e700 --- /dev/null +++ b/litellm/integrations/otel/langfuse_logger.py @@ -0,0 +1,60 @@ +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_logger +from litellm.integrations.otel.logger import OpenTelemetryV2 +from litellm.integrations.otel.mappers.langfuse import LANGFUSE_OBSERVATION_INPUT, LANGFUSE_OBSERVATION_OUTPUT +from litellm.integrations.otel.model.request_io import request_input, response_output, stream_output +from litellm.integrations.otel.plumbing.context import request_root_span + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import ModelResponseStream + + +class LangfuseOpenTelemetryV2(OpenTelemetryV2): + """Stamps the request's input and output on the root observation while it is still recording. + + Langfuse shows a trace's input and output from its root observation. The proxy's root span ends + when the response is sent, before the success callback runs, so both stamps come from the + post-call hooks in the request task: the request as it stands after the pre-call chain and the + response as it is returned, for the call types whose response renders as a message. + """ + + async def async_post_call_success_hook( + self, + data: Mapping[str, object], + user_api_key_dict: "UserAPIKeyAuth", + response: object, + ) -> None: + self._stamp_root_io(data, lambda: response_output(response)) + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: "UserAPIKeyAuth", + response: "AsyncIterator[ModelResponseStream]", + request_data: Mapping[str, object], + ) -> "AsyncGenerator[ModelResponseStream, None]": + relayed: Final[list[ModelResponseStream]] = [] # mutable-ok: relayed as they arrive, assembled at end of stream + async for chunk in response: + relayed.append(chunk) + yield chunk + self._stamp_root_io(request_data, lambda: stream_output(tuple(relayed), request_data)) + + def _stamp_root_io(self, data: Mapping[str, object], render_output: Callable[[], str | None]) -> None: + root: Final = request_root_span() + if root is None or not root.is_recording(): + return + try: + output: Final = render_output() + if output is None: + return + root.set_attribute(LANGFUSE_OBSERVATION_OUTPUT, output) + rendered_input: Final = request_input(data) + except Exception: # noqa: BLE001 # telemetry must never fail the request it describes + verbose_logger.debug( + "otel v2 langfuse: could not render the root observation input or output", exc_info=True + ) + return + if rendered_input is not None: + root.set_attribute(LANGFUSE_OBSERVATION_INPUT, rendered_input) diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index d2a32ef73b6..a550dca6cc8 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -4,6 +4,7 @@ from collections import OrderedDict from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast from opentelemetry.context import Context, attach, get_current @@ -909,3 +910,29 @@ def phase_span(name: str) -> "Iterator[Span | None]": return with logger.start_phase_span(name) as span: yield span + + +def build_otel_v2_logger( + config: OpenTelemetryV2Config, + callback_name: str | None = None, + tracer_provider: TracerProvider | None = None, + logger_provider: LoggerProvider | None = None, + meter_provider: "MeterProvider | None" = None, + settings: Mapping[str, object] = MappingProxyType({}), +) -> OpenTelemetryV2: + return _logger_class(config)( + config=config, + callback_name=callback_name, + tracer_provider=tracer_provider, + logger_provider=logger_provider, + meter_provider=meter_provider, + **settings, + ) + + +def _logger_class(config: OpenTelemetryV2Config) -> type[OpenTelemetryV2]: + if "langfuse" not in config.mapper_names or not config.capture_span_content: + return OpenTelemetryV2 + from litellm.integrations.otel.langfuse_logger import LangfuseOpenTelemetryV2 + + return LangfuseOpenTelemetryV2 diff --git a/litellm/integrations/otel/mappers/langfuse.py b/litellm/integrations/otel/mappers/langfuse.py index 6d4f1b4fd0a..01063d85355 100644 --- a/litellm/integrations/otel/mappers/langfuse.py +++ b/litellm/integrations/otel/mappers/langfuse.py @@ -11,6 +11,7 @@ the JSON-serialized payloads. ``_llm_call`` just applies both tables. import json from collections.abc import Callable +from typing import Final from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData from litellm.integrations.otel.mappers.utils import ( @@ -25,6 +26,9 @@ from litellm.integrations.otel.model.payloads import ( LLMUsage, ) +LANGFUSE_OBSERVATION_INPUT: Final = "langfuse.observation.input" +LANGFUSE_OBSERVATION_OUTPUT: Final = "langfuse.observation.output" + class LangfuseMapper: _LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = { @@ -56,8 +60,8 @@ class LangfuseMapper: "langfuse.observation.model.parameters": lambda d: json_if( collect(LangfuseMapper._MODEL_PARAMS, d.request_params) ), - "langfuse.observation.input": lambda d: serialize_messages(d.messages_in), - "langfuse.observation.output": lambda d: serialize_messages(output_messages(d)), + LANGFUSE_OBSERVATION_INPUT: lambda d: serialize_messages(d.messages_in), + LANGFUSE_OBSERVATION_OUTPUT: lambda d: serialize_messages(output_messages(d)), "langfuse.observation.usage_details": lambda d: json_if(collect(LangfuseMapper._USAGE_FIELDS, d.usage)), "langfuse.observation.cost_details": lambda d: ( json.dumps({"total": d.response_cost}) if d.response_cost is not None else None diff --git a/litellm/integrations/otel/model/request_io.py b/litellm/integrations/otel/model/request_io.py new file mode 100644 index 00000000000..4e80fb91993 --- /dev/null +++ b/litellm/integrations/otel/model/request_io.py @@ -0,0 +1,90 @@ +from collections.abc import Mapping, Sequence +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +import litellm +from litellm.integrations.otel.mappers.utils import json_or_none +from litellm.proxy.guardrails.anthropic_sse import assemble_anthropic_sse_stream, is_raw_sse_stream +from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse +from litellm.types.utils import ModelResponse, ModelResponseStream + +_SYSTEM_KEYS: Final = ("system", "instructions") +_TURNS: Final = TypeAdapter(tuple[object, ...]) +_MESSAGES: Final = TypeAdapter(list[object] | None) + + +class _Turn(TypedDict): + role: ReadOnly[str] + content: ReadOnly[object] + + +class _AnthropicMessage(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["message"] = Field(exclude=True) + role: str = "assistant" + content: object = None + + +def request_input(data: Mapping[str, object]) -> str | None: + turns: Final = data.get("messages", data.get("input")) + if turns is None: + return None + return json_or_none((*_system_turns(data), *_user_turns(turns))) + + +def _system_turns(data: Mapping[str, object]) -> tuple[_Turn, ...]: + return tuple(_Turn(role="system", content=data[key]) for key in _SYSTEM_KEYS if data.get(key) is not None) + + +def _user_turns(turns: object) -> tuple[object, ...]: + if isinstance(turns, str): + return (_Turn(role="user", content=turns),) + try: + return _TURNS.validate_python(turns) + except ValidationError: + return (_Turn(role="user", content=turns),) + + +def response_output(response: object) -> str | None: + match response: + case ModelResponse(): + return json_or_none(tuple(choice.message.model_dump(exclude_none=True) for choice in response.choices)) + case ResponsesAPIResponse(): + return json_or_none(response.model_dump(exclude_none=True).get("output")) + case _: + return _anthropic_message_output(response) + + +def _anthropic_message_output(message: object) -> str | None: + try: + parsed: Final = _AnthropicMessage.model_validate(message) + except ValidationError: + return None + return json_or_none((parsed.model_dump(),)) + + +def stream_output(chunks: Sequence[object], data: Mapping[str, object]) -> str | None: + if not chunks: + return None + if is_raw_sse_stream(chunks): + return response_output(assemble_anthropic_sse_stream(chunks)) + if all(isinstance(chunk, ModelResponseStream) for chunk in chunks): + return response_output(_assembled_chat_stream(chunks, data)) + return response_output(_completed_response(chunks)) + + +def _assembled_chat_stream(chunks: Sequence[object], data: Mapping[str, object]) -> object: + try: + return litellm.stream_chunk_builder( # pyright: ignore[reportUnknownMemberType] # upstream types chunks as a bare list + chunks=list(chunks), # mutable-ok: stream_chunk_builder takes a list + messages=_MESSAGES.validate_python(data.get("messages")), + ) + except (litellm.APIError, ValidationError): + return None + + +def _completed_response(chunks: Sequence[object]) -> ResponsesAPIResponse | None: + return next((chunk.response for chunk in reversed(chunks) if isinstance(chunk, ResponseCompletedEvent)), None) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9a6fb11f978..f54eeca5178 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -901,6 +901,7 @@ class Logging(LiteLLMLoggingBaseClass): prompt_label: str | None = None, prompt_version: int | None = None, request_kwargs: dict[str, object] | None = None, # mutable-ok: marker stamped into live request kwargs + injected_for_every_deployment: bool = False, ) -> tuple[str, list[AllMessageValues], dict]: from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook @@ -933,6 +934,7 @@ class Logging(LiteLLMLoggingBaseClass): AnthropicCacheControlHook.record_gateway_injection( request_kwargs, AnthropicCacheControlHook.count_request_cache_breakpoints(messages) - breakpoints_before, + injected_for_every_deployment=injected_for_every_deployment, ) self.messages = messages return model, messages, non_default_params @@ -950,6 +952,7 @@ class Logging(LiteLLMLoggingBaseClass): prompt_label: str | None = None, prompt_version: int | None = None, request_kwargs: dict[str, object] | None = None, # mutable-ok: marker stamped into live request kwargs + injected_for_every_deployment: bool = False, ) -> tuple[str, list[AllMessageValues], dict]: from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook @@ -985,6 +988,7 @@ class Logging(LiteLLMLoggingBaseClass): AnthropicCacheControlHook.record_gateway_injection( request_kwargs, AnthropicCacheControlHook.count_request_cache_breakpoints(messages) - breakpoints_before, + injected_for_every_deployment=injected_for_every_deployment, ) self.messages = messages return model, messages, non_default_params @@ -4390,13 +4394,15 @@ def _init_custom_logger_compatible_class( from litellm.integrations.otel.model.config import is_otel_v2_enabled if is_otel_v2_enabled(): - from litellm.integrations.otel.logger import OpenTelemetryV2 + from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger + from litellm.integrations.otel.model.config import OpenTelemetryV2Config for callback in _in_memory_loggers: - if type(callback) is OpenTelemetryV2: + if isinstance(callback, OpenTelemetryV2): return callback - otel_logger_v2: Final = OpenTelemetryV2( - **_get_custom_logger_settings_from_proxy_server(callback_name=logging_integration) + otel_settings: Final = _get_custom_logger_settings_from_proxy_server(callback_name=logging_integration) + otel_logger_v2: Final = build_otel_v2_logger( + config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings ) _in_memory_loggers.append(otel_logger_v2) _maybe_auto_initialize_arize_phoenix(_in_memory_loggers) @@ -4759,7 +4765,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[Custom if not is_otel_v2_enabled(): return None - from litellm.integrations.otel.logger import OpenTelemetryV2 + from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger from litellm.integrations.otel.presets import PRESET_BY_CALLBACK preset_fn: Final = PRESET_BY_CALLBACK.get(callback_name) @@ -4774,7 +4780,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[Custom # If env vars are missing or the preset raises, defer to the legacy path # so customers get the same error story they had before V2 landed. return None - v2_logger: Final = OpenTelemetryV2(config=config, callback_name=callback_name) + v2_logger: Final = build_otel_v2_logger(config=config, callback_name=callback_name) _in_memory_loggers.append(v2_logger) return v2_logger diff --git a/litellm/litellm_core_utils/secret_redaction.py b/litellm/litellm_core_utils/secret_redaction.py index 5d5bd547d22..e93ab155786 100644 --- a/litellm/litellm_core_utils/secret_redaction.py +++ b/litellm/litellm_core_utils/secret_redaction.py @@ -30,6 +30,11 @@ def _build_secret_patterns() -> "re.Pattern[str]": r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}", # OpenAI / Anthropic sk- prefixed keys rf"sk-[A-Za-z0-9\-_]{{{MINIMUM_CUSTOM_KEY_LENGTH - len('sk-')},}}", + # Credentials passed as URL query params. Terminated by "&" like the key= + # and sig= patterns below, so the rest of the request line survives in an + # access log. Must precede the generic patterns to win at the same position. + r"(?<=[?&])(?:api[_-]?key|\w*(?:token|password|passwd|client_secret|secret_key|_secret))" + r"=[^\s&'\"]+", # Generic api_key / api-key / apikey (handles 'key': 'value' dict repr) r"(?:api[_-]?key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]{8,}", # x-api-key / api-key header values (handles 'key': 'value' dict repr) @@ -45,8 +50,10 @@ def _build_secret_patterns() -> "re.Pattern[str]": # Word boundary prevents O(n^2) backtracking on long word-char runs. r"(?:^|(?<=\W))\w*(?:password|passwd|client_secret|secret_key|_secret)" r"['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+", - # Database connection string credentials (scheme://user:pass@host) - r"(?<=://)[^\s'\"]*:[^\s'\"@]+(?=@)", + # Database connection string credentials (scheme://user:pass@host). + # The user half stops at the ":" separator and both halves are length-capped, + # so a long attacker-supplied URL cannot backtrack quadratically. + r"(?<=://)[^\s'\":]{0,4096}:[^\s'\"]{1,4096}(?=@)", # Databricks personal access tokens r"dapi[0-9a-f]{32}", # Module-level provider keys logged as litellm._key= @@ -67,8 +74,10 @@ def _build_secret_patterns() -> "re.Pattern[str]": r"""['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+""", # Raw JWTs (without Bearer prefix) r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]*", - # Azure SAS tokens in URLs - r"[?&]sig=[A-Za-z0-9%+/=]+", + # Azure SAS tokens in URLs. The delimiter is a lookbehind, like the + # `key=` pattern above, so the `?` or `&` survives and the redacted URL + # stays well formed (this string is often a request line in a log). + r"(?<=[?&])sig=[A-Za-z0-9%+/=]+", # Full JSON service-account blobs (single-line and multi-line) r'\{[^{}]*"type"\s*:\s*"service_account"[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', ] diff --git a/litellm/llms/base_llm/files/azure_blob_storage_backend.py b/litellm/llms/base_llm/files/azure_blob_storage_backend.py index e22cf528856..a192226110c 100644 --- a/litellm/llms/base_llm/files/azure_blob_storage_backend.py +++ b/litellm/llms/base_llm/files/azure_blob_storage_backend.py @@ -7,16 +7,36 @@ to reuse all authentication and Azure Storage operations. """ import time +from pathlib import Path from typing import Final from urllib.parse import quote, urlparse from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger +from litellm.proxy.common_utils.path_utils import safe_filename from .storage_backend import BaseFileStorageBackend +def _safe_basename(original_filename: str) -> str: + try: + return safe_filename(original_filename) + except ValueError: + return "file" + + +def _safe_extension(original_filename: str) -> str: + """The extension off a basename, with no path separators or traversal sequences. + + original_filename.split(".")[-1] does not parse path structure, so a filename + like "a.jsonl/../../etc/cron.d/x" would put "../../etc/cron.d/x" straight into + the blob path built below. Path.suffix only ever looks at the last path + component, so routing through safe_filename() first closes that off. + """ + return Path(_safe_basename(original_filename)).suffix.lstrip(".") + + class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): """ Azure Blob Storage backend implementation. @@ -81,16 +101,15 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): def _generate_file_name(self, original_filename: str, file_naming_strategy: str) -> str: """Generate file name based on naming strategy.""" if file_naming_strategy == "original_filename": - # Use original filename, but sanitize it - return quote(original_filename, safe="") + return quote(_safe_basename(original_filename), safe="") elif file_naming_strategy == "timestamp": # Use timestamp - extension = original_filename.split(".")[-1] if "." in original_filename else "" + extension = _safe_extension(original_filename) timestamp: Final = int(time.time() * 1000) # milliseconds return f"{timestamp}.{extension}" if extension else str(timestamp) else: # default to "uuid" # Use UUID - extension = original_filename.split(".")[-1] if "." in original_filename else "" + extension = _safe_extension(original_filename) file_uuid: Final = str(uuid.uuid4()) return f"{file_uuid}.{extension}" if extension else file_uuid diff --git a/litellm/main.py b/litellm/main.py index 01c106adc7c..0128e4defe5 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -141,7 +141,6 @@ from litellm.utils import ( convert_to_model_response_object, create_pretrained_tokenizer, create_tokenizer, - get_api_key, get_llm_provider, get_model_info, get_non_default_completion_params, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2846d12db6e..d8a8f84b032 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -33093,6 +33093,88 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "meta/muse-spark-1.3": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "meta/muse-spark-1.3-contributor": { + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 1e-07, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "meta_llama/Llama-3.3-70B-Instruct": { "litellm_provider": "meta_llama", "max_input_tokens": 128000, diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 425f82794e6..66d4aedba06 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1149,10 +1149,11 @@ class MCPRequestHandler: would miss a real outage wrapped inside it.""" from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler - if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): + outage: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(e) + if outage is not None: raise HTTPException( status_code=503, - detail="Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", + detail=PrismaDBExceptionHandler.database_unavailable_message(outage), ) from None @staticmethod diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 09a3703e904..35a30127e27 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -101,18 +101,29 @@ class _ResolvedKey: key: "UserAPIKeyAuth" -_KeyResolutionFailure = Literal["no_active_key", "unavailable", "unresolvable"] +_KeyResolutionFailure = Literal["no_active_key", "unavailable", "faulted", "unresolvable"] """Why a token request yielded no active litellm key, kept distinct so a caller statuses each truthfully instead of blaming the client for a gateway problem: - ``no_active_key``: none was presented, or the presented key is unknown / blocked / expired (the caller's request is at fault) - ``unavailable``: the auth database was transiently unreachable while resolving (retryable) +- ``faulted``: the auth database's query engine reported a fault that retrying will not clear (still a + 503, but the wording must not tell the operator to wait) - ``unresolvable``: the gateway cannot resolve identity right now (no DB connection, or an unexpected error) -- a gateway fault, not the caller's The classification mirrors admission's ``_reload_admitted_key`` so the mint (ingress) and admission (egress) never disagree on the status of the same outage.""" +def _database_failure(exc: Exception) -> Literal["unavailable", "faulted"]: + from litellm.proxy.db.exception_handler import ( # noqa: PLC0415 # inline import avoids a module-load circular import + PrismaDBExceptionHandler, + ) + + fault: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(exc) or exc + return "faulted" if PrismaDBExceptionHandler.is_permanent_database_fault(fault) else "unavailable" + + async def _resolve_active_litellm_key(request: Request) -> "_ResolvedKey | _KeyResolutionFailure": """Resolve the presented litellm key to an active key record, or say precisely why not. @@ -170,7 +181,7 @@ async def _reload_active_key_by_hash(key_hash: str) -> "_ResolvedKey | _KeyResol return "no_active_key" except Exception as exc: # noqa: BLE001 # classify: a DB outage is retryable, anything else is an opaque gateway fault if PrismaDBExceptionHandler.is_database_service_unavailable_error(exc): - return "unavailable" + return _database_failure(exc) verbose_logger.debug( "_reload_active_key_by_hash: unexpected key-resolution error (%s)", type(exc).__name__, @@ -225,8 +236,9 @@ async def load_active_user_by_id(user_id: str) -> "LiteLLM_UserTable | _KeyResol except (ProxyException, HTTPException): return "no_active_key" except Exception as exc: # noqa: BLE001 # a DB outage is retryable; a missing user (get_user_object's wrapped ValueError) or any other resolution failure fails closed as no_active_key, never a 500 - if PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(exc): - return "unavailable" + outage: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(exc) + if outage is not None: + return _database_failure(outage) verbose_logger.debug("_reload_active_user_by_id: user-resolution error (%s)", type(exc).__name__) return "no_active_key" if user_object is None: @@ -383,6 +395,7 @@ _BridgeMintError = Literal[ "no_identity", "invalid_refresh", "identity_unavailable", + "identity_faulted", "identity_unresolvable", "not_configured", "no_upstream_token", @@ -433,6 +446,13 @@ def _bridge_mint_error_response(error: _BridgeMintError) -> JSONResponse: "temporarily_unavailable", "the authentication database is temporarily unreachable; retry shortly", ) + case "identity_faulted": + status, code, desc = ( + 503, + "temporarily_unavailable", + "the authentication database reported a fault that is not a transient outage; " + "retrying will not help until the gateway deployment is repaired", + ) case "identity_unresolvable": status, code, desc = ( 500, @@ -485,6 +505,8 @@ def _key_resolution_failure_to_mint_error(failure: _KeyResolutionFailure) -> _Br return "no_identity" case "unavailable": return "identity_unavailable" + case "faulted": + return "identity_faulted" case "unresolvable": return "identity_unresolvable" case _: @@ -569,6 +591,8 @@ def _refresh_key_failure_to_mint_error(failure: _KeyResolutionFailure) -> _Bridg return "invalid_refresh" case "unavailable": return "identity_unavailable" + case "faulted": + return "identity_faulted" case "unresolvable": return "identity_unresolvable" case _: diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index a43e762a456..c7b0045dde5 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -150,11 +150,18 @@ _CLIENT_RECORD_DEBUG_KEY: Final = "gateway_dcr_client" _CONNECT_FLOW_DEBUG_KEY: Final = "gateway_connect_flow" _AUTH_CODE_DEBUG_KEY: Final = "gateway_authorization_code" -ReloadUserFailure = Literal["unresolvable", "unavailable", "no_active_key"] +ReloadUserFailure = Literal["unresolvable", "unavailable", "faulted", "no_active_key"] ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]] """Injected live-user revalidation (the token endpoint's mirror of admission): -``None`` means the user is active; ``unavailable`` is a retryable DB outage; anything -else fails the grant closed.""" +``None`` means the user is active; ``unavailable`` is a retryable DB outage; ``faulted`` is +a DB fault retrying will not clear (still 503, worded so nobody just waits); anything else +fails the grant closed.""" + +_DB_UNAVAILABLE_DESCRIPTION: Final = "the gateway database is unavailable; retry" +_DB_FAULTED_DESCRIPTION: Final = ( + "the gateway database reported a fault that is not a transient outage; " + "retrying will not help until the gateway deployment is repaired" +) PROXY_API_AUDIENCE: Final[SessionAudience] = "proxy_api" """The audience a native client (``lite login --pkce``, a Go CLI) asks for by sending the @@ -659,7 +666,9 @@ def _set_flow_cookie(response: Response, request: Request, handle: str, flow: _C def _consent_lookup_failure_response(failure: ReloadUserFailure) -> Response: match failure: case "unavailable": - return _oauth_error(503, "temporarily_unavailable", "the gateway database is unavailable; retry") + return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) + case "faulted": + return _oauth_error(503, "temporarily_unavailable", _DB_FAULTED_DESCRIPTION) case "unresolvable": return _oauth_error(500, "server_error", "the gateway is not configured to resolve users") case "no_active_key": @@ -962,7 +971,9 @@ def _reload_failure_response(failure: ReloadUserFailure) -> Response: ``ReloadUserFailure`` member is a type error here rather than silently 400ing.""" match failure: case "unavailable": - return _oauth_error(503, "temporarily_unavailable", "the gateway database is unavailable; retry") + return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) + case "faulted": + return _oauth_error(503, "temporarily_unavailable", _DB_FAULTED_DESCRIPTION) case "unresolvable": return _oauth_error(500, "server_error", "the gateway is not configured to resolve users") case "no_active_key": @@ -981,7 +992,7 @@ def _mint_failure_response(failure: ProxyCredentialMintFailure) -> Response: return _oauth_error( 400, "invalid_grant", "this user belongs to a team; sign in again and pick the team for this credential" ) - case "unavailable" | "unresolvable" | "no_active_key": + case "unavailable" | "faulted" | "unresolvable" | "no_active_key": return _reload_failure_response(failure) case _: assert_never(failure) @@ -1297,8 +1308,8 @@ async def introspect_gateway_token( if peeked == "claimed": return _inactive_introspection_response() failure: Final = await reload_user(opened.principal.user_id) - if failure == "unavailable": - return _oauth_error(503, "temporarily_unavailable", "the gateway database is unavailable; retry") + if failure == "unavailable" or failure == "faulted": + return _reload_failure_response(failure) if failure is not None: return _inactive_introspection_response() return _active_introspection_response(opened) diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index f79765f6d01..4f6305d88cf 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -2,20 +2,31 @@ from __future__ import annotations import json from collections.abc import Mapping, Sequence +from dataclasses import dataclass from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypedDict, assert_never +from pydantic import ValidationError from typing_extensions import ReadOnly, Required import litellm from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K +from litellm.proxy.common_utils.semantic_text_index import ( + Embedder, + EmbeddingFailed, + SemanticTextIndex, + router_embedder, +) +from litellm.types.mcp import MCPToolSearchSettings if TYPE_CHECKING: - from mcp.types import CallToolResult + from mcp.types import CallToolResult, Tool from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth +MCP_TOOL_SEARCH_SETTINGS_KEY: Final[str] = "mcp_tool_search" MCP_TOOL_SEARCH_TOOL_NAME: Final[str] = "mcp_tool_search" MCP_TOOL_CALL_TOOL_NAME: Final[str] = "mcp_tool_call" AGENT_SEARCH_TOOL_NAME: Final[str] = "agent_search" @@ -29,17 +40,91 @@ def coerce_top_k(value: Any, default: int = 5) -> int: return default -def search_tools(query: str, tools: list[dict[str, Any]], top_k: int = 5) -> list[dict[str, Any]]: +class ToolSearchResult(TypedDict, total=False): + name: Required[ReadOnly[str]] + description: Required[ReadOnly[str]] + inputSchema: Required[ReadOnly[Mapping[str, object]]] + score: ReadOnly[float] + + +@dataclass(frozen=True, slots=True) +class SemanticToolRanker: + embed: Embedder + embedding_model: str + index: SemanticTextIndex + + +global_mcp_tool_search_index: Final = SemanticTextIndex() + + +def mcp_tool_search_settings() -> MCPToolSearchSettings | ValidationError: + try: + return MCPToolSearchSettings.model_validate(litellm.mcp_tool_search or {}) + except ValidationError as exc: + return exc + + +def _tool_result(tool: Tool) -> ToolSearchResult: + return {"name": tool.name, "description": tool.description or "", "inputSchema": tool.inputSchema} + + +def _scored_result(tool: Tool, score: float) -> ToolSearchResult: + return {"name": tool.name, "description": tool.description or "", "inputSchema": tool.inputSchema, "score": score} + + +def _tool_text(tool: Tool) -> str: + return "\n".join(part for part in (tool.name, tool.description or "") if part) + + +def _keyword_score(query: str, tool: Tool) -> float: + haystack: Final = _tool_text(tool).lower() + return float(sum(1 for token in query.lower().split() if token in haystack)) + + +def _split_core_tools(tools: Sequence[Tool], core_tools: Sequence[str]) -> tuple[tuple[Tool, ...], tuple[Tool, ...]]: + by_name: Final = MappingProxyType({tool.name: tool for tool in tools}) + core: Final = tuple(by_name[name] for name in dict.fromkeys(core_tools) if name in by_name) + rest: Final = tuple(tool for tool in tools if tool.name not in frozenset(core_tools)) + return core, rest + + +def _top_hits( + tools: Sequence[Tool], scores: Sequence[float], minimum: float, limit: int +) -> tuple[tuple[float, Tool], ...]: + hits: Final = ((score, tool) for score, tool in zip(scores, tools, strict=True) if score >= minimum) + return tuple(sorted(hits, key=lambda hit: hit[0], reverse=True)[:limit]) + + +def search_tools(query: str, tools: Sequence[Tool], top_k: int = 5) -> tuple[ToolSearchResult, ...]: + """Keyword fallback used when no embedding model is configured: one point per query token found in the tool.""" if not query: - return [] - tokens: Final = query.lower().split() + return () + scores: Final = tuple(_keyword_score(query, tool) for tool in tools) + return tuple(_tool_result(tool) for _, tool in _top_hits(tools, scores, minimum=1.0, limit=top_k)) - def _score(tool: dict[str, Any]) -> int: - haystack: Final = (tool.get("name", "") + " " + tool.get("description", "")).lower() - return sum(1 for t in tokens if t in haystack) - scored: Final = ((s, tool) for tool in tools if (s := _score(tool)) > 0) - return [tool for _, tool in sorted(scored, key=lambda x: x[0], reverse=True)[:top_k]] +async def search_mcp_tools( + query: str, + tools: Sequence[Tool], + top_k: int, + settings: MCPToolSearchSettings, + ranker: SemanticToolRanker | None, +) -> tuple[ToolSearchResult, ...] | EmbeddingFailed: + """Core tools the caller can access come first, then up to `top_k` ranked matches from the remaining tools.""" + core, rest = _split_core_tools(tools, settings.core_tools) + limit: Final = min(top_k, settings.top_k) + core_results: Final = tuple(_tool_result(tool) for tool in core) + if ranker is None: + return (*core_results, *search_tools(query, rest, limit)) + if not query: + return core_results + scores: Final = await ranker.index.scores( + query, tuple(_tool_text(tool) for tool in rest), ranker.embed, ranker.embedding_model + ) + if isinstance(scores, EmbeddingFailed): + return scores + hits: Final = _top_hits(rest, scores, minimum=settings.similarity_threshold, limit=limit) + return (*core_results, *(_scored_result(tool, score) for score, tool in hits)) class _ToolParamSchema(TypedDict, total=False): @@ -66,11 +151,17 @@ def _json_array(*items: str) -> Sequence[str]: _MCP_TOOL_SEARCH_DEFINITION: Final[VirtualToolDefinition] = { "name": MCP_TOOL_SEARCH_TOOL_NAME, - "description": "Search for MCP tools by keyword. Returns top matching tools with names, descriptions, and input schemas.", + "description": ( + "Search for MCP tools by describing what you need. " + "Returns top matching tools with names, descriptions, and input schemas." + ), "inputSchema": { "type": "object", "properties": { - "query": {"type": "string", "description": "Keywords to search for in tool names and descriptions."}, + "query": { + "type": "string", + "description": "What the tool should do, matched against names and descriptions.", + }, "top_k": {"type": "integer", "description": "Maximum number of results to return.", "default": 5}, }, "required": _json_array("query"), @@ -165,10 +256,28 @@ async def handle_mcp_tool_search( oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, ) -> CallToolResult: - from mcp.types import CallToolResult, TextContent - from litellm.proxy._experimental.mcp_server.server import _list_mcp_tools + from litellm.proxy.proxy_server import llm_router + settings: Final = mcp_tool_search_settings() + if isinstance(settings, ValidationError): + return _text_tool_result( + f"litellm_settings.{MCP_TOOL_SEARCH_SETTINGS_KEY} is invalid: {settings}", is_error=True + ) + if settings.embedding_model is not None and llm_router is None: + return _text_tool_result( + f"litellm_settings.{MCP_TOOL_SEARCH_SETTINGS_KEY}.embedding_model needs a model_list so it can be called", + is_error=True, + ) + ranker: Final = ( + SemanticToolRanker( + embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict), + embedding_model=settings.embedding_model, + index=global_mcp_tool_search_index, + ) + if settings.embedding_model is not None and llm_router is not None + else None + ) mcp_listing: Final = await _list_mcp_tools( user_api_key_auth=user_api_key_dict, mcp_servers=mcp_servers, @@ -178,17 +287,10 @@ async def handle_mcp_tool_search( oauth2_headers=oauth2_headers, raw_headers=raw_headers, ) - mcp_tools: Final = mcp_listing.tools - tools: Final = [ - { - "name": t.name, - "description": t.description or "", - "inputSchema": t.inputSchema, - } - for t in mcp_tools - ] - results: Final = search_tools(query, tools, top_k) - return CallToolResult(content=[TextContent(type="text", text=json.dumps(results))], isError=False) + results: Final = await search_mcp_tools(query, mcp_listing.tools, top_k, settings, ranker) + if isinstance(results, EmbeddingFailed): + return _text_tool_result(results.reason, is_error=True) + return _text_tool_result(json.dumps(results), is_error=False) async def handle_mcp_tool_call( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 18714256a8f..2da7ceb2d50 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2508,6 +2508,14 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="max batch input file size in MB for /v1/files uploads with purpose=batch, if a file is larger than this size it will be rejected before being forwarded to the provider", ) + max_file_size_mb: int | None = Field( + None, + description="max file size in MB for /v1/files uploads, for any purpose, if a file is larger than this size it will be rejected before being forwarded to the provider", + ) + blocked_file_extensions: tuple[str, ...] | None = Field( + None, + description="file extensions (e.g. ['.exe', '.sh']) rejected on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename", + ) max_response_size_mb: int | None = Field( None, description="max response size in MB, if a response is larger than this size it will be rejected", @@ -2696,6 +2704,40 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).", ) + password_policy_min_length: int | None = Field( + None, + description=( + "Minimum length required for a locally-managed user's password. Default is 12; " + "a value below 8 is floored to 8 rather than weakening the requirement further." + ), + ) + password_policy_require_uppercase: bool | None = Field( + None, + description="If True (default), a locally-managed user's password must contain an uppercase letter.", + ) + password_policy_require_lowercase: bool | None = Field( + None, + description="If True (default), a locally-managed user's password must contain a lowercase letter.", + ) + password_policy_require_numbers: bool | None = Field( + None, + description="If True (default), a locally-managed user's password must contain a number.", + ) + password_policy_require_special_characters: bool | None = Field( + None, + description="If True (default), a locally-managed user's password must contain a special (non-alphanumeric) character.", + ) + disable_password_login_when_sso_enabled: bool | None = Field( + None, + description=( + "If True and SSO is configured (MICROSOFT_CLIENT_ID, GOOGLE_CLIENT_ID, " + "GENERIC_CLIENT_ID, or SAML_IDP_METADATA_URL/XML), disables username/password " + "login on /login, /v2/login, and /v3/login so SSO is the only way to reach the " + "Admin UI. An admin locked out of the UI can still administer the proxy over the " + "API with the master key; unset this setting and restart the proxy to restore " + "UI username/password login. Default is False." + ), + ) disable_budget_reservation: bool | None = Field( None, description=( diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 144de52d0d2..c7b6bca72cf 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -6,8 +6,12 @@ from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypedDict +from pydantic import TypeAdapter, ValidationError + import litellm +from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) @@ -52,6 +56,9 @@ class AgentRecord(Protocol): @property def agent_name(self) -> str: ... + @property + def litellm_params(self) -> Mapping[str, object] | None: ... + @property def object_permission_id(self) -> str | None: ... @@ -121,6 +128,188 @@ def _dump_agent_params(raw: Mapping[str, object]) -> dict[str, object]: return dict(raw) if raw else {} +_AGENT_PARAMS_MASKER: Final = SensitiveDataMasker() +_REDACT_AGENT_PARAMS_MAX_DEPTH: Final = 10 +_AGENT_PARAMS_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter( + dict[str, object] +) # mutable-ok: safe_dumps() and AgentResponse.litellm_params both require a real dict, not a Mapping +_AGENT_PARAMS_SEQUENCE_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) +_EMPTY_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def redact_sensitive_agent_litellm_params(litellm_params: object, _depth: int = 0) -> object: + """ + Replace credential-bearing values in an agent's litellm_params with + ``REDACTED_BY_LITELM_STRING`` while preserving non-secret keys (``model``, + ``is_public``, rate-limit config). Used so list/get/create/update + responses never echo a stored provider credential back to the caller. + + Handles a plain dict, a JSON-serialized string (some callers hold the + in-memory registry's params that way), and ``None`` at the top level; + anything else is passed through. Recursion depth is bounded to match the + convention documented in ``tests/code_coverage_tests/recursive_detector.py``. + """ + if litellm_params is None: + return None + if isinstance(litellm_params, str): + if _depth >= _REDACT_AGENT_PARAMS_MAX_DEPTH: + return REDACTED_BY_LITELM_STRING + try: + parsed_params: Final = _AGENT_PARAMS_ADAPTER.validate_json(litellm_params) + except ValidationError: + return REDACTED_BY_LITELM_STRING + return json.dumps(_redact_agent_params_tree(parsed_params, _depth + 1)) + return _redact_agent_params_tree(litellm_params, _depth) + + +def _redact_agent_params_tree(value: object, _depth: int) -> object: + """Structural recursion over an already-parsed litellm_params value: a + dict redacts sensitive keys and recurses into the rest, a list redacts + each element (so a secret nested inside a list of provider configs is + still caught), and anything else -- including a plain string leaf, which + must never be re-interpreted as a JSON blob -- passes through unchanged. + """ + if _depth >= _REDACT_AGENT_PARAMS_MAX_DEPTH: + return REDACTED_BY_LITELM_STRING + if isinstance(value, list): + typed_items: Final = _AGENT_PARAMS_SEQUENCE_ADAPTER.validate_python(value) + return tuple(_redact_agent_params_tree(item, _depth + 1) for item in typed_items) + if not isinstance(value, dict): + return value + typed_params: Final = _AGENT_PARAMS_ADAPTER.validate_python(value) + return { + key: ( + REDACTED_BY_LITELM_STRING + if _AGENT_PARAMS_MASKER.is_sensitive_key(key) + else _redact_agent_params_tree(nested_value, _depth + 1) + ) + for key, nested_value in typed_params.items() + } # mutable-ok: consumed by json.dumps()/AgentResponse.litellm_params, both of which require a real dict + + +def parse_agent_litellm_params(value: object) -> Mapping[str, object]: + """Normalize a stored litellm_params column to a read-only mapping. + + The prisma Json column comes back as either an already-parsed dict or a + JSON string depending on the read path, so handle both rather than + assuming one. Only ever read from (merge-source lookups), never mutated + or re-serialized directly, so a read-only view is enough here. + """ + if isinstance(value, str): + try: + return _AGENT_PARAMS_ADAPTER.validate_json(value) + except ValidationError: + return _EMPTY_LITELLM_PARAMS + if isinstance(value, Mapping): + try: + return _AGENT_PARAMS_ADAPTER.validate_python(value) + except ValidationError: + return _EMPTY_LITELLM_PARAMS + return _EMPTY_LITELLM_PARAMS + + +_MISSING_AGENT_PARAM: Final = object() +_RESTORE_AGENT_PARAMS_MAX_DEPTH: Final = 10 + + +def _restore_redacted_nested_value(incoming_value: object, existing_value: object, _depth: int) -> object: + """Recurse into a non-sensitively-named dict/list value so a secret + nested underneath it (e.g. inside a list of per-provider configs) is + still restored, not just top-level keys. Mirrors the shapes + ``redact_sensitive_agent_litellm_params`` recurses into on read, so + restore and redact stay symmetric. + + List elements are paired with the existing list by position: with no + stable per-element identity in an arbitrary ``dict[str, object]`` schema, + index is the same correspondence every other part of this restore (and + the endpoints' existing full-replace-on-PUT semantics) already assumes. + This correctly preserves a masked secret across an ordinary edit of that + same entry's other fields; it does not protect against a caller who both + reorders/resizes the list AND echoes back a masked marker in the same + request, which is a known, narrow limitation (see LIT-6736 PR discussion) + rather than a cross-entry credential leak in the common case. + + A value collapsed to the flat marker by the read side's depth cap is + recovered wholesale from ``existing_value`` (rather than the marker + string itself getting persisted) whenever ``existing_value`` isn't + already that same flat marker. Depth-bounded like its read-side + counterpart; a value at the cap is returned unchanged rather than + corrupted. + """ + if incoming_value == REDACTED_BY_LITELM_STRING and existing_value != REDACTED_BY_LITELM_STRING: + return existing_value + if _depth >= _RESTORE_AGENT_PARAMS_MAX_DEPTH: + return incoming_value + if isinstance(incoming_value, Mapping): + typed_incoming_map: Final = _AGENT_PARAMS_ADAPTER.validate_python(incoming_value) + existing_map: Final = ( + _AGENT_PARAMS_ADAPTER.validate_python(existing_value) + if isinstance(existing_value, Mapping) + else _EMPTY_LITELLM_PARAMS + ) + return _restore_redacted_litellm_params(typed_incoming_map, existing_map, _depth + 1) + if isinstance(incoming_value, (list, tuple)): + typed_incoming_seq: Final = _AGENT_PARAMS_SEQUENCE_ADAPTER.validate_python(incoming_value) + existing_seq: Final = ( + _AGENT_PARAMS_SEQUENCE_ADAPTER.validate_python(existing_value) + if isinstance(existing_value, (list, tuple)) + else () + ) + return tuple( + _restore_redacted_nested_value( + item, + existing_seq[index] if index < len(existing_seq) else None, + _depth + 1, + ) + for index, item in enumerate(typed_incoming_seq) + ) + return incoming_value + + +def _resolved_agent_param_value( + key: str, + incoming: Mapping[str, object], + existing: Mapping[str, object], + _depth: int, +) -> object: + """The value ``key`` should end up with in a restored litellm_params, or + ``_MISSING_AGENT_PARAM`` when it should be dropped entirely.""" + if key in incoming: + value: Final = incoming[key] + if _AGENT_PARAMS_MASKER.is_sensitive_key(key): + return existing.get(key, _MISSING_AGENT_PARAM) if value == REDACTED_BY_LITELM_STRING else value + return _restore_redacted_nested_value(value, existing.get(key), _depth) + if _AGENT_PARAMS_MASKER.is_sensitive_key(key): + return existing.get(key, _MISSING_AGENT_PARAM) + return _MISSING_AGENT_PARAM + + +def _restore_redacted_litellm_params( + incoming: Mapping[str, object], + existing: Mapping[str, object], + _depth: int = 0, +) -> dict[str, object]: + """Restore the real credential behind any litellm_params value the caller + echoed back as ``REDACTED_BY_LITELM_STRING``, and behind any sensitive key + omitted entirely, so an edit to an unrelated field never overwrites (or + silently drops) a stored provider credential -- the UI never has to + read-and-resend a secret to keep it. Recurses into nested dicts and lists + so a secret nested under a non-sensitively-named key is restored too. + + A sensitive key given a real (non-marker) value, including an explicit + empty string, is treated as a deliberate update -- that's how a caller + clears a credential. Non-sensitive keys always take the incoming value + (recursed into), matching the endpoints' existing full-replace-on-PUT / + merge-on-PATCH semantics for everything that isn't a secret. + """ + all_keys: Final = frozenset(incoming) | frozenset(existing) + return { + key: value + for key in all_keys + if (value := _resolved_agent_param_value(key, incoming, existing, _depth)) is not _MISSING_AGENT_PARAM + } # mutable-ok: fed to safe_dumps() for JSON-column storage, which requires a real dict + + class GrantMigrationResult(NamedTuple): rewritten: int missed: int @@ -301,9 +490,14 @@ class AgentRegistry: try: agent_name: Final = agent.get("agent_name") - # Serialize litellm_params + # Serialize litellm_params. A create has no stored row to restore a + # secret behind, so a sensitive key submitted as the redaction + # marker (e.g. a stray client re-post) is dropped rather than + # persisted as the literal placeholder string. litellm_params_obj: Final = agent.get("litellm_params", {}) - litellm_params_dict: Final[dict[str, object]] = _dump_agent_params(litellm_params_obj) + litellm_params_dict: Final = _restore_redacted_litellm_params( + _dump_agent_params(litellm_params_obj), _EMPTY_LITELLM_PARAMS + ) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params @@ -410,8 +604,14 @@ class AgentRegistry: update_data: Final[dict[str, object]] = {} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") - if augment_agent.get("litellm_params"): - update_data["litellm_params"] = safe_dumps(augment_agent.get("litellm_params")) + if "litellm_params" in agent: + existing_litellm_params: Final = parse_agent_litellm_params(existing_agent.get("litellm_params")) + update_data["litellm_params"] = safe_dumps( + _restore_redacted_litellm_params( + _dump_agent_params(agent.get("litellm_params") or _EMPTY_LITELLM_PARAMS), + existing_litellm_params, + ) + ) if augment_agent.get("agent_card_params"): update_data["agent_card_params"] = safe_dumps(augment_agent.get("agent_card_params")) @@ -474,9 +674,22 @@ class AgentRegistry: try: agent_name: Final = agent.get("agent_name") + # A PUT fully replaces litellm_params from the request body, so the + # existing row is read up front to restore any sensitive key the + # caller echoed back redacted (or omitted) rather than persisting + # the marker -- or nothing -- over the real stored credential. + existing_row: Final = await agents_table(prisma_client).find_unique( + where={"agent_id": agent_id} # mutable-ok: prisma's query builder rejects a Mapping/MappingProxyType + ) + existing_litellm_params: Final = parse_agent_litellm_params( + existing_row.litellm_params if existing_row is not None else None + ) + # Serialize litellm_params litellm_params_obj: Final = agent.get("litellm_params", {}) - litellm_params_dict: Final[dict[str, object]] = _dump_agent_params(litellm_params_obj) + litellm_params_dict: Final = _restore_redacted_litellm_params( + _dump_agent_params(litellm_params_obj), existing_litellm_params + ) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params @@ -512,9 +725,8 @@ class AgentRegistry: update_data[rate_field] = _val if agent.get("object_permission") is not None: - existing_agent: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) existing_object_permission_id: Final = ( - existing_agent.object_permission_id if existing_agent is not None else None + existing_row.object_permission_id if existing_row is not None else None ) agent_copy: Final = dict(agent) object_permission_id: Final = await handle_update_object_permission_common( diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py index 46ab36d7b72..76e3fe6c5ad 100644 --- a/litellm/proxy/agent_endpoints/agent_search.py +++ b/litellm/proxy/agent_endpoints/agent_search.py @@ -2,17 +2,18 @@ from __future__ import annotations -import math -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import Sequence from dataclasses import dataclass -from itertools import chain -from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol, TypeAlias +from typing import TYPE_CHECKING, Final, TypeAlias -from openai import OpenAIError from pydantic import BaseModel, ConfigDict, ValidationError -from litellm.exceptions import BudgetExceededError +from litellm.proxy.common_utils.semantic_text_index import ( + Embedder, + EmbeddingFailed, + SemanticTextIndex, + router_embedder, +) from litellm.types.agents import AgentResponse if TYPE_CHECKING: @@ -21,12 +22,6 @@ if TYPE_CHECKING: DEFAULT_AGENT_SEARCH_TOP_K: Final = 5 -Vector: TypeAlias = tuple[float, ...] - - -class Embedder(Protocol): - def __call__(self, texts: Sequence[str]) -> Awaitable[Sequence[Vector]]: ... - @dataclass(frozen=True, slots=True) class AgentSearchHit: @@ -67,18 +62,6 @@ class _SearchableCard(BaseModel): skills: tuple[_SearchableSkill, ...] = () -class _EmbeddingItem(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore") - - embedding: tuple[float, ...] - - -class _EmbeddingData(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore") - - data: tuple[_EmbeddingItem, ...] - - class AgentSearchResult(BaseModel): model_config = ConfigDict(frozen=True) @@ -117,110 +100,21 @@ def agent_search_result(hit: AgentSearchHit) -> AgentSearchResult: ) -def cosine_similarity(left: Vector, right: Vector) -> float: - dot: Final = sum(a * b for a, b in zip(left, right, strict=True)) - norms: Final = math.sqrt(sum(a * a for a in left)) * math.sqrt(sum(b * b for b in right)) - return dot / norms if norms else 0.0 - - -def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]: - from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup - - return { # mutable-ok: the router mutates the metadata dict it is handed - **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict), - "user_api_key": user_api_key_dict.api_key, - } - - -def router_embedder(router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth) -> Embedder: - async def embed(texts: Sequence[str]) -> Sequence[Vector]: - batch: Final = list(texts) # mutable-ok: Router.aembedding accepts only str | list input - response: Final = await router.aembedding( - model=embedding_model, input=batch, metadata=embedding_spend_metadata(user_api_key_dict) - ) - return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data) - - return embed - - -_NO_VECTORS: Final[Mapping[str, Vector]] = MappingProxyType({}) - - -async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | AgentSearchEmbeddingFailed: - try: - vectors: Final = tuple(await embed(texts)) - except (OpenAIError, ValueError, BudgetExceededError) as exc: - return AgentSearchEmbeddingFailed(reason=f"embedding the search query failed: {exc}") - if len(vectors) != len(texts): - return AgentSearchEmbeddingFailed( - reason=f"embedding model returned {len(vectors)} vectors for {len(texts)} inputs" - ) - return vectors - - -@dataclass(frozen=True, slots=True) -class _Embedded: - query_vector: Vector - vectors: Mapping[str, Vector] - - -def _same_dimension(query_vector: Vector, vectors: Mapping[str, Vector], texts: Sequence[str]) -> bool: - return all(len(vectors[text]) == len(query_vector) for text in texts) - - -async def _embed_query_and_agents( - embed: Embedder, query: str, texts: Sequence[str], cached: Mapping[str, Vector] -) -> _Embedded | AgentSearchEmbeddingFailed: - missing: Final = tuple(dict.fromkeys(text for text in texts if text not in cached)) - embedded: Final = await _embed_all(embed, (query, *missing)) - if isinstance(embedded, AgentSearchEmbeddingFailed): - return embedded - vectors: Final = MappingProxyType(dict(chain(cached.items(), zip(missing, embedded[1:], strict=True)))) - if _same_dimension(embedded[0], vectors, texts): - return _Embedded(query_vector=embedded[0], vectors=vectors) - unique: Final = tuple(dict.fromkeys(texts)) - reembedded: Final = await _embed_all(embed, (query, *unique)) - if isinstance(reembedded, AgentSearchEmbeddingFailed): - return reembedded - return _Embedded( - query_vector=reembedded[0], vectors=MappingProxyType(dict(zip(unique, reembedded[1:], strict=True))) - ) - - class AgentSearchIndex: """Caches one vector per distinct agent text per embedding model, so repeat searches only embed the query.""" def __init__(self) -> None: - self._vectors: Mapping[str, Mapping[str, Vector]] = MappingProxyType({}) - - def _merged(self, embedding_model: str, embedded: _Embedded) -> Mapping[str, Vector]: - kept: Final = { - text: vector - for text, vector in self._vectors.get(embedding_model, _NO_VECTORS).items() - if len(vector) == len(embedded.query_vector) - } - return MappingProxyType({**kept, **embedded.vectors}) + self._index: Final = SemanticTextIndex() async def search( self, query: str, agents: Sequence[AgentResponse], top_k: int, embed: Embedder, embedding_model: str ) -> AgentSearchHits | AgentSearchEmbeddingFailed: - if not agents: - return AgentSearchHits(hits=()) texts: Final = tuple(agent_search_text(agent) for agent in agents) - cached: Final = self._vectors.get(embedding_model, _NO_VECTORS) - embedded: Final = await _embed_query_and_agents(embed, query, texts, cached) - if isinstance(embedded, AgentSearchEmbeddingFailed): - return embedded - if not _same_dimension(embedded.query_vector, embedded.vectors, texts): - return AgentSearchEmbeddingFailed( - reason=f"embedding model {embedding_model} returned vectors of mixed dimensions" - ) - self._vectors = MappingProxyType({**self._vectors, embedding_model: self._merged(embedding_model, embedded)}) + scores: Final = await self._index.scores(query, texts, embed, embedding_model) + if isinstance(scores, EmbeddingFailed): + return AgentSearchEmbeddingFailed(reason=scores.reason) ranked: Final = sorted( - ( - AgentSearchHit(agent=agent, score=cosine_similarity(embedded.query_vector, embedded.vectors[text])) - for agent, text in zip(agents, texts, strict=True) - ), + (AgentSearchHit(agent=agent, score=score) for agent, score in zip(agents, scores, strict=True)), key=lambda hit: hit.score, reverse=True, ) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index b6c41a17503..3e4dc07a521 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -20,7 +20,6 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import ( CommonProxyErrors, @@ -33,6 +32,10 @@ from litellm.proxy.a2a.agent_card import ( merge_agent_card, normalize_protocol_version, ) +from litellm.proxy.agent_endpoints.agent_registry import ( + parse_agent_litellm_params, + redact_sensitive_agent_litellm_params, +) from litellm.proxy.agent_endpoints.agent_search import ( DEFAULT_AGENT_SEARCH_TOP_K, AgentSearchEmbeddingFailed, @@ -139,25 +142,37 @@ async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client) agent.keys = matched_keys or None +def _redact_agent_litellm_params_dict( + litellm_params: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: AgentResponse.litellm_params is declared as a plain dict, not Mapping + """Type-narrowing wrapper: a dict in always yields a dict back from + ``redact_sensitive_agent_litellm_params``, which the function's general + (possible-JSON-string, possibly-None) signature can't express.""" + return dict( # mutable-ok: AgentResponse.litellm_params is declared as a plain dict, not Mapping + parse_agent_litellm_params(redact_sensitive_agent_litellm_params(litellm_params)) + ) + + def _redact_sensitive_agent_fields( agents: Sequence[AgentResponse], + *, + is_admin: bool, ) -> list[AgentResponse]: """ - Return copies of the given agents with sensitive configuration fields - redacted. The original objects are not modified. + Return copies of the given agents with credential-bearing litellm_params + values replaced by a fixed marker (never returned to ANY caller, + admin included) and, for non-admin callers, virtual-key and header + fields stripped entirely. The original objects are not modified. """ redacted: Final[list[AgentResponse]] = [] for agent in agents: copy = agent.model_copy(deep=True) - copy.static_headers = None - copy.extra_headers = None - copy.keys = None + if not is_admin: + copy.static_headers = None + copy.extra_headers = None + copy.keys = None if copy.litellm_params: - copy.litellm_params = _get_masked_values( - copy.litellm_params, - unmasked_length=4, - number_of_asterisks=4, - ) + copy.litellm_params = _redact_agent_litellm_params_dict(copy.litellm_params) redacted.append(copy) return redacted @@ -345,13 +360,13 @@ async def get_agents( global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups) ) - # Redact sensitive fields for non-admin users + # litellm_params secrets are always redacted; keys/headers stay + # admin-only. is_admin: Final = ( user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value ) - if not is_admin: - returned_agents = _redact_sensitive_agent_fields(returned_agents) + returned_agents = _redact_sensitive_agent_fields(returned_agents, is_admin=is_admin) if health_check: agents_with_url: Final = [agent for agent in returned_agents if (agent.agent_card_params or {}).get("url")] @@ -505,7 +520,9 @@ async def create_agent( "Failed to register agent '%s' (ID: %s) in memory: %s", agent_name, agent_id, reg_error ) - return result + # The caller is a proxy admin (enforced above); litellm_params + # secrets are still never echoed back in the response. + return _redact_sensitive_agent_fields((result,), is_admin=True)[0] except HTTPException: raise @@ -578,13 +595,13 @@ async def get_agent_by_id( await _attach_keys_to_agents([agent], prisma_client) - # Redact sensitive fields for non-admin users + # litellm_params secrets are always redacted; keys/headers stay + # admin-only. is_admin = ( user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value ) - if not is_admin: - agent = _redact_sensitive_agent_fields([agent])[0] + agent = _redact_sensitive_agent_fields((agent,), is_admin=is_admin)[0] return agent except HTTPException: @@ -688,7 +705,7 @@ async def update_agent( "Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id ) - return result + return _redact_sensitive_agent_fields((result,), is_admin=True)[0] except HTTPException: raise except Exception as e: @@ -791,7 +808,7 @@ async def patch_agent( "Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id ) - return result + return _redact_sensitive_agent_fields((result,), is_admin=True)[0] except HTTPException: raise except Exception as e: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5703c6cd5e8..0db4cbc3bf2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4705,6 +4705,13 @@ async def is_valid_fallback_model( return True +# The shape abbreviate_api_key writes into LiteLLM_VerificationToken.key_name. The +# last four characters are only barred from being whitespace or a control code, +# because a custom key's can be anything else, punctuation and non-ASCII included; +# a real key is at least MINIMUM_CUSTOM_KEY_LENGTH long, so it never fullmatches. +_MASKED_KEY_NAME_RE: Final = re.compile(r"sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?") + + def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool: """ Throttle an over-budget key instead of blocking it, when the key opted in @@ -4785,10 +4792,15 @@ async def _virtual_key_max_budget_check( if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: if _apply_budget_exceeded_throttle(valid_token): return - # name the key in the error so operators don't have to reverse-map - # spend back to a key; key_name is the masked form (last 4 chars) + # This message is returned to the caller, and key_name has no enforced + # shape (a direct DB write bypasses abbreviate_api_key), so echo it only + # when it still looks masked and fall back to the alias otherwise. key_label: Final = valid_token.key_alias or "key" - key_descriptor: Final = f"{key_label} ({valid_token.key_name})" if valid_token.key_name else key_label + key_descriptor: Final = ( + f"{key_label} ({valid_token.key_name})" + if valid_token.key_name and _MASKED_KEY_NAME_RE.fullmatch(valid_token.key_name) + else key_label + ) raise litellm.BudgetExceededError( current_cost=spend, max_budget=valid_token.max_budget, diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 64878a480a7..b36c8a038fc 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -61,9 +61,7 @@ def _as_proxy_exception(e: Exception) -> ProxyException: return e if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): return ProxyException( - message=( - "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly." - ), + message=PrismaDBExceptionHandler.database_unavailable_message(e), type=ProxyErrorTypes.no_db_connection, param="None", code=status.HTTP_503_SERVICE_UNAVAILABLE, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 89b2c92cdfd..d6007a2d56e 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1,3 +1,4 @@ +import importlib.util import os import re import sys @@ -1402,7 +1403,7 @@ def is_pass_through_provider_route(route: str) -> bool: return False -def _has_user_setup_sso() -> bool: +def has_user_setup_sso() -> bool: """ Check if the user has set up single sign-on (SSO). @@ -1425,6 +1426,63 @@ def _has_user_setup_sso() -> bool: ) +def _is_google_ready() -> bool: + return bool(os.getenv("GOOGLE_CLIENT_ID")) and bool(os.getenv("GOOGLE_CLIENT_SECRET")) + + +def _is_microsoft_ready() -> bool: + return ( + bool(os.getenv("MICROSOFT_CLIENT_ID")) + and bool(os.getenv("MICROSOFT_CLIENT_SECRET")) + and bool(os.getenv("MICROSOFT_TENANT")) + ) + + +def _is_generic_oauth_ready() -> bool: + return ( + bool(os.getenv("GENERIC_CLIENT_ID")) + and bool(os.getenv("GENERIC_CLIENT_SECRET")) + and bool(os.getenv("GENERIC_AUTHORIZATION_ENDPOINT")) + and bool(os.getenv("GENERIC_TOKEN_ENDPOINT")) + and bool(os.getenv("GENERIC_USERINFO_ENDPOINT")) + ) + + +def _is_saml_ready() -> bool: + if not (os.getenv("SAML_IDP_METADATA_URL") or os.getenv("SAML_IDP_METADATA_XML")): + return False + # SAML's runtime (python3-saml) is an optional dependency; the SAML + # handler itself fails closed on every request when it is missing + # (SAMLAuthHandler raises before touching the IdP), so metadata alone + # is not "ready" either. find_spec raises ModuleNotFoundError (rather + # than returning None) when the top-level package is absent entirely, + # so this must not be a bare boolean expression or every password + # login would 500 on a deployment that configured SAML metadata + # without installing the optional extra. + try: + return importlib.util.find_spec("onelogin.saml2.auth") is not None + except ModuleNotFoundError: + return False + + +def is_sso_provider_fully_configured() -> bool: + """Whether ANY configured SSO provider has every companion setting it + needs to actually authenticate a user, not merely a client id. + + A lone ``MICROSOFT_CLIENT_ID`` with no secret or tenant makes + ``has_user_setup_sso()`` return True while every real sign-in attempt + fails, so a gate that BLOCKS the password fallback (unlike the UI + discovery use of ``has_user_setup_sso()``, where a dead login button is + merely confusing) must check readiness here, or it can lock every admin + out with no way to sign in at all. Checks every provider independently + (mirroring ``/sso/readiness``'s per-provider requirements) rather than + stopping at the first one with a client id set, so a stray leftover + client id for an unused provider can never mask a different, fully + configured provider that would otherwise satisfy this gate. + """ + return _is_google_ready() or _is_microsoft_ready() or _is_generic_oauth_ready() or _is_saml_ready() + + def get_customer_user_header_from_mapping(user_id_mapping) -> list | None: """Return the header_name mapped to CUSTOMER role, if any (dict-based).""" if not user_id_mapping: diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index fba95972944..8d4f6f81363 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -7,7 +7,9 @@ login endpoints (e.g., /login and /v2/login). import os import secrets +from collections.abc import Mapping from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import Final, Literal, cast import jwt @@ -24,6 +26,7 @@ from litellm.proxy._types import ( UpdateUserRequest, UserAPIKeyAuth, ) +from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured from litellm.proxy.management_endpoints.internal_user_endpoints import user_update from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, @@ -111,6 +114,7 @@ async def authenticate_user( password: str, master_key: str | None, prisma_client: PrismaClient | None, + general_settings: Mapping[str, object] = MappingProxyType({}), ) -> LoginResult: """ Authenticate a user and generate an API key for UI access. @@ -124,13 +128,40 @@ async def authenticate_user( password: Password from the login form master_key: Master key for the proxy (required) prisma_client: Prisma database client (optional) + general_settings: Proxy general_settings, checked for + `disable_password_login_when_sso_enabled` Returns: LoginResult: Object containing authentication data Raises: - ProxyException: If authentication fails or required configuration is missing + ProxyException: If authentication fails or required configuration is missing, + or if username/password login is disabled while SSO is configured + + Recovery: an admin locked out of the UI by + `disable_password_login_when_sso_enabled` can still administer the proxy over + the API with the master key (Authorization: Bearer ), which never + goes through this function. To restore UI username/password login, unset the + setting in config.yaml (or the DB-persisted general_settings) and restart the + proxy; this is a deliberate, auditable config change rather than a hidden + bypass. + + The gate below requires the SSO provider to be FULLY configured (every + companion secret/endpoint an actual sign-in needs), not merely that a + client id is present, so an incomplete SSO setup can never disable the + only working login path. """ + if general_settings.get("disable_password_login_when_sso_enabled") is True and is_sso_provider_fully_configured(): + raise ProxyException( + message=( + "Username/password login is disabled because SSO is configured " + "and 'disable_password_login_when_sso_enabled' is set. Sign in via SSO." + ), + type=ProxyErrorTypes.auth_error, + param="disable_password_login_when_sso_enabled", + code=403, + ) + if master_key is None: raise ProxyException( message="Master Key not set for Proxy. Please set Master Key to use Admin UI. Set `LITELLM_MASTER_KEY` in .env or set general_settings:master_key in config.yaml. https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.", diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py new file mode 100644 index 00000000000..ab7a565894a --- /dev/null +++ b/litellm/proxy/auth/password_policy.py @@ -0,0 +1,92 @@ +"""Password-strength policy enforcement for locally-managed proxy users. + +Applied at every path that persists a new or changed password for a DB-backed +user (``/user/update``, ``/user/bulk_update``, and the invitation onboarding +claim flow), so the strength bar is configured in one place instead of +per-endpoint. +""" + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +from litellm.proxy._types import ProxyErrorTypes, ProxyException + +DEFAULT_MIN_LENGTH: Final = 12 +MIN_ALLOWED_LENGTH: Final = 8 + + +def _has_uppercase(password: str) -> bool: + return any(ch.isupper() for ch in password) + + +def _has_lowercase(password: str) -> bool: + return any(ch.islower() for ch in password) + + +def _has_digit(password: str) -> bool: + return any(ch.isdigit() for ch in password) + + +def _has_special_character(password: str) -> bool: + """Unicode-aware: a letter or digit from ANY script counts as + alphanumeric, not just ASCII, so an accented letter (e.g. the second + character of "Passwörd1234") cannot be miscounted as the required + special character the way an ASCII-only `[^A-Za-z0-9]` regex would.""" + return any(not ch.isalnum() for ch in password) + + +@dataclass(frozen=True, slots=True) +class PasswordPolicy: + min_length: int + require_uppercase: bool + require_lowercase: bool + require_numbers: bool + require_special_characters: bool + + +def _configured_min_length(general_settings: Mapping[str, object]) -> int: + """The configured minimum, floored at MIN_ALLOWED_LENGTH so a nonpositive + or too-low override (a typo, or `0`/`false` coercing through) cannot + silently disable the length requirement rather than merely relaxing it.""" + min_length_setting: Final = general_settings.get("password_policy_min_length") + if isinstance(min_length_setting, bool) or not isinstance(min_length_setting, (int, float)): + return DEFAULT_MIN_LENGTH + return max(MIN_ALLOWED_LENGTH, int(min_length_setting)) + + +def get_password_policy(general_settings: Mapping[str, object]) -> PasswordPolicy: + return PasswordPolicy( + min_length=_configured_min_length(general_settings), + require_uppercase=general_settings.get("password_policy_require_uppercase", True) is not False, + require_lowercase=general_settings.get("password_policy_require_lowercase", True) is not False, + require_numbers=general_settings.get("password_policy_require_numbers", True) is not False, + require_special_characters=( + general_settings.get("password_policy_require_special_characters", True) is not False + ), + ) + + +def _policy_violations(password: str, policy: PasswordPolicy) -> tuple[str, ...]: + checks: Final = ( + (len(password) < policy.min_length, f"be at least {policy.min_length} characters long"), + (policy.require_uppercase and not _has_uppercase(password), "include an uppercase letter"), + (policy.require_lowercase and not _has_lowercase(password), "include a lowercase letter"), + (policy.require_numbers and not _has_digit(password), "include a number"), + (policy.require_special_characters and not _has_special_character(password), "include a special character"), + ) + return tuple(message for failed, message in checks if failed) + + +def validate_password_policy(password: str, general_settings: Mapping[str, object]) -> None: + """Raise ``ProxyException`` (400) if ``password`` fails the configured policy.""" + policy: Final = get_password_policy(general_settings) + violations: Final = _policy_violations(password, policy) + if not violations: + return + raise ProxyException( + message="Password does not meet the required policy: must " + ", ".join(violations) + ".", + type=ProxyErrorTypes.validation_error, + param="password", + code=400, + ) diff --git a/litellm/proxy/common_utils/semantic_text_index.py b/litellm/proxy/common_utils/semantic_text_index.py new file mode 100644 index 00000000000..0820459af49 --- /dev/null +++ b/litellm/proxy/common_utils/semantic_text_index.py @@ -0,0 +1,142 @@ +"""Embedding-similarity ranking over short texts with a per-model vector cache, shared by agent search and MCP tool search.""" + +from __future__ import annotations + +import math +from collections.abc import Awaitable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias + +from openai import OpenAIError +from pydantic import BaseModel, ConfigDict + +from litellm.exceptions import BudgetExceededError + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.router import Router + +Vector: TypeAlias = tuple[float, ...] + + +class Embedder(Protocol): + def __call__(self, texts: Sequence[str]) -> Awaitable[Sequence[Vector]]: ... + + +@dataclass(frozen=True, slots=True) +class EmbeddingFailed: + reason: str + + +class _EmbeddingItem(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + embedding: tuple[float, ...] + + +class _EmbeddingData(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + data: tuple[_EmbeddingItem, ...] + + +def cosine_similarity(left: Vector, right: Vector) -> float: + dot: Final = sum(a * b for a, b in zip(left, right, strict=True)) + norms: Final = math.sqrt(sum(a * a for a in left)) * math.sqrt(sum(b * b for b in right)) + return dot / norms if norms else 0.0 + + +def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]: # mutable-ok: router mutates it + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + return { # mutable-ok: the router mutates the metadata dict it is handed + **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict), + "user_api_key": user_api_key_dict.api_key, + } + + +def router_embedder(router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth) -> Embedder: + async def embed(texts: Sequence[str]) -> Sequence[Vector]: + batch: Final = list(texts) # mutable-ok: Router.aembedding accepts only str | list input + response: Final = await router.aembedding( + model=embedding_model, input=batch, metadata=embedding_spend_metadata(user_api_key_dict) + ) + return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data) + + return embed + + +_NO_VECTORS: Final[Mapping[str, Vector]] = MappingProxyType({}) + + +async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed: + try: + vectors: Final = tuple(await embed(texts)) + except (OpenAIError, ValueError, BudgetExceededError) as exc: + return EmbeddingFailed(reason=f"embedding the search query failed: {exc}") + if len(vectors) != len(texts): + return EmbeddingFailed(reason=f"embedding model returned {len(vectors)} vectors for {len(texts)} inputs") + return vectors + + +@dataclass(frozen=True, slots=True) +class _Embedded: + query_vector: Vector + vectors: Mapping[str, Vector] + + +def _same_dimension(query_vector: Vector, vectors: Mapping[str, Vector], texts: Sequence[str]) -> bool: + return all(len(vectors[text]) == len(query_vector) for text in texts) + + +async def _embed_query_and_texts( + embed: Embedder, query: str, texts: Sequence[str], cached: Mapping[str, Vector] +) -> _Embedded | EmbeddingFailed: + missing: Final = tuple(dict.fromkeys(text for text in texts if text not in cached)) + embedded: Final = await _embed_all(embed, (query, *missing)) + if isinstance(embedded, EmbeddingFailed): + return embedded + vectors: Final = MappingProxyType(dict(chain(cached.items(), zip(missing, embedded[1:], strict=True)))) + if _same_dimension(embedded[0], vectors, texts): + return _Embedded(query_vector=embedded[0], vectors=vectors) + unique: Final = tuple(dict.fromkeys(texts)) + reembedded: Final = await _embed_all(embed, (query, *unique)) + if isinstance(reembedded, EmbeddingFailed): + return reembedded + return _Embedded( + query_vector=reembedded[0], vectors=MappingProxyType(dict(zip(unique, reembedded[1:], strict=True))) + ) + + +class SemanticTextIndex: + """Caches one vector per distinct text per embedding model, so repeat searches only embed the query.""" + + def __init__(self) -> None: + self._vectors: Mapping[str, Mapping[str, Vector]] = MappingProxyType({}) + + def _merged(self, embedding_model: str, embedded: _Embedded) -> Mapping[str, Vector]: + kept: Final = MappingProxyType( + { + text: vector + for text, vector in self._vectors.get(embedding_model, _NO_VECTORS).items() + if len(vector) == len(embedded.query_vector) + } + ) + return MappingProxyType({**kept, **embedded.vectors}) + + async def scores( + self, query: str, texts: Sequence[str], embed: Embedder, embedding_model: str + ) -> tuple[float, ...] | EmbeddingFailed: + """Cosine similarity of `query` to each entry of `texts`, in the same order.""" + if not texts: + return () + cached: Final = self._vectors.get(embedding_model, _NO_VECTORS) + embedded: Final = await _embed_query_and_texts(embed, query, texts, cached) + if isinstance(embedded, EmbeddingFailed): + return embedded + if not _same_dimension(embedded.query_vector, embedded.vectors, texts): + return EmbeddingFailed(reason=f"embedding model {embedding_model} returned vectors of mixed dimensions") + self._vectors = MappingProxyType({**self._vectors, embedding_model: self._merged(embedding_model, embedded)}) + return tuple(cosine_similarity(embedded.query_vector, embedded.vectors[text]) for text in texts) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 5502543b926..ef1c4a66203 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,4 +1,4 @@ -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Iterator from typing import Any, Final, TypeVar from litellm._logging import verbose_proxy_logger @@ -9,10 +9,43 @@ from litellm.proxy._types import ( ) from litellm.secret_managers.main import str_to_bool -# Bounds the __cause__/__context__ walk in is_database_service_unavailable_error_in_chain. +# Bounds the __cause__/__context__ walk in find_database_service_unavailable_error_in_chain. # Real exception chains are a few links deep; the cap also makes the walk cycle-safe. _MAX_EXCEPTION_CHAIN_DEPTH: Final = 20 +_TRANSIENT_DB_UNAVAILABLE_MESSAGE: Final = ( + "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly." +) + + +def _exception_chain(e: BaseException) -> Iterator[BaseException]: + current = e # rebind-ok: advances one link per iteration of the bounded walk + for _ in range(_MAX_EXCEPTION_CHAIN_DEPTH): + yield current + following = current.__cause__ or current.__context__ + if following is None: + return + current = following + + +def _database_service_unavailable_errors(e: BaseException) -> tuple[Exception, ...]: + return tuple( + link + for link in _exception_chain(e) + if isinstance(link, Exception) and PrismaDBExceptionHandler.is_database_service_unavailable_error(link) + ) + + +def _exception_types(*candidates: object) -> tuple[type[BaseException], ...]: + """Keep only the real exception classes among ``candidates``. + + The predicates below resolve prisma's error classes at call time, so a test + that swaps ``sys.modules["prisma"]`` for a ``MagicMock`` hands them mocks, + and ``isinstance`` against a mock raises ``TypeError`` instead of answering + False. Dropping the non-types lets the call fall through to the other checks. + """ + return tuple(c for c in candidates if isinstance(c, type) and issubclass(c, BaseException)) + class PrismaDBExceptionHandler: """ @@ -59,7 +92,7 @@ class PrismaDBExceptionHandler: if isinstance(e, DB_CONNECTION_ERROR_TYPES): return True - if isinstance(e, prisma.engine.errors.EngineConnectionError): + if isinstance(e, _exception_types(prisma.engine.errors.EngineConnectionError)): return True return isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection @@ -81,7 +114,7 @@ class PrismaDBExceptionHandler: """ import prisma - data_layer_errors: Final = ( + data_layer_errors: Final = _exception_types( prisma.errors.DataError, prisma.errors.UniqueViolationError, prisma.errors.ForeignKeyViolationError, @@ -94,7 +127,7 @@ class PrismaDBExceptionHandler: return False if isinstance(e, DB_CONNECTION_ERROR_TYPES): return True - if isinstance(e, prisma.errors.PrismaError): + if isinstance(e, _exception_types(prisma.errors.PrismaError)): return True if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection: return True @@ -138,13 +171,13 @@ class PrismaDBExceptionHandler: return True if isinstance( e, - ( + _exception_types( prisma.errors.ClientNotConnectedError, prisma.errors.HTTPClientClosedError, ), ): return True - if isinstance(e, prisma.errors.PrismaError): + if isinstance(e, _exception_types(prisma.errors.PrismaError)): error_message: Final = str(e).lower() connection_keywords: Final = ( "can't reach database server", @@ -171,7 +204,7 @@ class PrismaDBExceptionHandler: """True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma.""" import prisma - if not isinstance(e, prisma.errors.PrismaError): + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): return False if getattr(e, "code", None) == "P2034": return True @@ -202,7 +235,7 @@ class PrismaDBExceptionHandler: """ import prisma - if isinstance(e, prisma.errors.PrismaError): + if isinstance(e, _exception_types(prisma.errors.PrismaError)): return False tb = getattr(e, "__traceback__", None) while tb is not None: @@ -268,6 +301,51 @@ class PrismaDBExceptionHandler: ), ) + @staticmethod + def is_permanent_database_fault(e: Exception) -> bool: + """True for a service-unavailable failure that will not clear on its + own: an engine-layer ``PrismaError`` (missing or version-skewed engine + binary, engine error status, misused transaction) that is neither the + transient ``EngineConnectionError`` nor a reconnectable transport failure. + + Picks only the wording of a 503, never whether one is sent; + ``is_database_service_unavailable_error`` stays the status gate. + """ + if PrismaDBExceptionHandler.is_database_connection_error(e): + return False + if PrismaDBExceptionHandler.is_database_transport_error(e): + return False + return PrismaDBExceptionHandler.is_database_infrastructure_error(e) + + @staticmethod + def database_unavailable_message(e: Exception) -> str: + """The 503 detail for a service-unavailable database failure: retry + guidance for a transient outage, a pointer at the deployment for a + fault that retrying cannot fix. A permanent fault anywhere in the + exception chain wins, since the transport error that surfaced it is + not what blocks recovery.""" + fault: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(e) or e + if not PrismaDBExceptionHandler.is_permanent_database_fault(fault): + return _TRANSIENT_DB_UNAVAILABLE_MESSAGE + return ( + "Service Unavailable, the authentication database query engine reported " + f"{type(fault).__name__}, which is not a transient outage and will not clear by retrying. " + "The proxy deployment needs attention." + ) + + @staticmethod + def find_database_service_unavailable_error_in_chain(e: BaseException) -> Exception | None: + """The exception in the ``__cause__`` / ``__context__`` chain that + ``is_database_service_unavailable_error`` accepts, or ``None``. Callers + that word a response by the kind of outage need the wrapped database + error itself, not just the fact that one is present. A permanent fault + outranks a transient one wherever it sits in the chain: a reconnect that + dies on a missing engine binary raises the transport error last, but the + binary is what keeps the database down.""" + outages: Final = _database_service_unavailable_errors(e) + permanent: Final = next(filter(PrismaDBExceptionHandler.is_permanent_database_fault, outages), None) + return permanent if permanent is not None else next(iter(outages), None) + @staticmethod def is_database_service_unavailable_error_in_chain(e: BaseException) -> bool: """Like ``is_database_service_unavailable_error`` but also walks the @@ -285,14 +363,7 @@ class PrismaDBExceptionHandler: The walk is depth-bounded, which also makes it cycle-safe. """ - current: BaseException | None = e - for _ in range(_MAX_EXCEPTION_CHAIN_DEPTH): - if not isinstance(current, Exception): - return False - if PrismaDBExceptionHandler.is_database_service_unavailable_error(current): - return True - current = current.__cause__ or current.__context__ - return False + return PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(e) is not None @staticmethod def handle_db_exception(e: Exception): diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 64414b90ae7..c2053693f2e 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -14,7 +14,7 @@ router: Final = APIRouter() @router.get("/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints) @router.get("/litellm/.well-known/litellm-ui-config", response_model=UiDiscoveryEndpoints) # if mounted at root path async def get_ui_config(): - from litellm.proxy.auth.auth_utils import _has_user_setup_sso + from litellm.proxy.auth.auth_utils import has_user_setup_sso from litellm.proxy.proxy_server import general_settings from litellm.proxy.utils import get_proxy_base_url, get_server_root_path @@ -28,7 +28,7 @@ async def get_ui_config(): or general_settings.get("hide_default_credentials_hint", False) is True ) - sso_configured: Final = _has_user_setup_sso() + sso_configured: Final = has_user_setup_sso() from litellm.proxy.proxy_server import proxy_config diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index da51a905ae3..70ea21320ee 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1633,6 +1633,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Update the guardrails litellm params in memory """ super().update_in_memory_litellm_params(litellm_params) + if self.apply_to_output: + self.output_parse_pii = False if litellm_params.pii_entities_config: self.pii_entities_config = litellm_params.pii_entities_config if litellm_params.presidio_score_thresholds: diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 76dea1b7784..16369abbfb0 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -2,6 +2,7 @@ from typing import Any, Final import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import CommonProxyErrors from litellm.types.guardrails import * @@ -85,7 +86,7 @@ def initialize_lakera_v2(litellm_params: LitellmParams, guardrail: Guardrail): return _lakera_v2_callback -def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): +def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail) -> tuple[CustomGuardrail, ...]: from litellm.proxy.guardrails.guardrail_hooks.presidio import ( _OPTIONAL_PresidioPIIMasking, ) @@ -94,7 +95,7 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): run_input: Final = filter_scope in ("input", "both") run_output: Final = filter_scope in ("output", "both") - def _make_presidio_callback(**overrides): + def _make_presidio_callback(**overrides) -> CustomGuardrail: params: Final = dict( guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, @@ -120,27 +121,27 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): litellm.logging_callback_manager.add_litellm_callback(callback) return callback - primary_callback = None - - if run_input: - primary_callback = _make_presidio_callback() - - if litellm_params.output_parse_pii: - _make_presidio_callback( - output_parse_pii=True, - event_hook=GuardrailEventHooks.post_call.value, - ) - - if run_output: - output_callback: Final = _make_presidio_callback( + input_callback: Final = _make_presidio_callback() if run_input else None + unmask_output_callback: Final = ( + _make_presidio_callback( + output_parse_pii=True, + event_hook=GuardrailEventHooks.post_call.value, + ) + if run_input and litellm_params.output_parse_pii + else None + ) + mask_output_callback: Final = ( + _make_presidio_callback( apply_to_output=True, event_hook=GuardrailEventHooks.post_call.value, output_parse_pii=False, ) - if primary_callback is None: - primary_callback = output_callback - - return primary_callback + if run_output + else None + ) + return tuple( + callback for callback in (input_callback, unmask_output_callback, mask_output_callback) if callback is not None + ) def initialize_hide_secrets(litellm_params: LitellmParams, guardrail: Guardrail): diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index dc13c09dd38..bd35782444b 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -3,10 +3,10 @@ import asyncio import importlib import os -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime, timezone from itertools import chain, count -from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, cast +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeAlias, cast from pydantic import ValidationError @@ -90,6 +90,8 @@ guardrail_initializer_registry: Final = { CONFIG_GUARDRAIL_ID_NAMESPACE: Final = uuid.UUID("625f63f4-935a-50e5-98b5-fbe77babc74a") +GuardrailCallbacks: TypeAlias = tuple[CustomGuardrail, ...] + guardrail_class_registry: Final[dict[str, type[CustomGuardrail]]] = { SupportedGuardrailIntegrations.BEDROCK.value: BedrockGuardrail, SupportedGuardrailIntegrations.GRAYSWAN.value: GraySwanGuardrail, @@ -424,6 +426,41 @@ def _apply_configured_bool_overrides(instance: CustomGuardrail, litellm_params: instance.scan_raw_request = bool(litellm_params.scan_raw_request) +def _as_callback_tuple( + initialized: CustomGuardrail | Sequence[CustomGuardrail] | None, +) -> GuardrailCallbacks: + if initialized is None: + return () + if isinstance(initialized, (list, tuple)): + return tuple(initialized) + return (initialized,) + + +def _configure_callback_scoping( + custom_guardrail_callback: CustomGuardrail, guardrail_name: str, litellm_params: LitellmParams +) -> None: + for scoping_param in ( + "skip_system_message_in_guardrail", + "skip_tool_message_in_guardrail", + "scan_only_tool_results", + ): + setattr(custom_guardrail_callback, scoping_param, getattr(litellm_params, scoping_param, None)) + scan_only_tool_results_enabled: Final = effective_scan_only_tool_results_for_guardrail(custom_guardrail_callback) + if scan_only_tool_results_enabled and not custom_guardrail_callback.supports_scan_only_tool_results(): + raise ValueError( + f"Guardrail {guardrail_name}: scan_only_tool_results is enabled, but this " + "guardrail's role filtering never scans tool results, so no request content would ever " + "be scanned. Remove scan_only_tool_results or the guardrail's role-filtering option." + ) + if scan_only_tool_results_enabled and effective_skip_tool_message_for_guardrail(custom_guardrail_callback): + raise ValueError( + f"Guardrail {guardrail_name}: scan_only_tool_results and " + "skip_tool_message_in_guardrail are enabled together, which excludes every message from " + "scanning, so no request content would ever be scanned. Remove one of the two." + ) + _apply_configured_bool_overrides(custom_guardrail_callback, litellm_params) + + class InMemoryGuardrailHandler: """ Class that handles initializing guardrails and adding them to the CallbackManager @@ -440,6 +477,8 @@ class InMemoryGuardrailHandler: Guardrail id to CustomGuardrail object mapping """ + self.guardrail_id_to_sibling_callbacks: dict[str, GuardrailCallbacks] = {} # mutable-ok: per-id registry + self._sources: dict[str, Literal["db", "config"]] = {} """ Guardrail id to provenance marker. "db" entries are reconciled against @@ -474,7 +513,6 @@ class InMemoryGuardrailHandler: self._sources[guardrail_id] = source return self.IN_MEMORY_GUARDRAILS[guardrail_id] - custom_guardrail_callback: CustomGuardrail | None = None litellm_params_data: Final = guardrail["litellm_params"] verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data) @@ -498,54 +536,15 @@ class InMemoryGuardrailHandler: if guardrail_type is None: raise ValueError("guardrail_type is required") - initializer: Final = guardrail_initializer_registry.get(guardrail_type) - - if initializer: - # Try to call with llm_router first, fall back to without if it fails - import inspect - - sig: Final = inspect.signature(initializer) - if "llm_router" in sig.parameters: - custom_guardrail_callback = initializer( - litellm_params, - guardrail, - llm_router, - ) - else: - custom_guardrail_callback = initializer(litellm_params, guardrail) - elif isinstance(guardrail_type, str) and "." in guardrail_type: - custom_guardrail_callback = self.initialize_custom_guardrail( - guardrail=guardrail, - guardrail_type=guardrail_type, - litellm_params=litellm_params, - config_file_path=config_file_path, - ) - else: - raise ValueError(f"Unsupported guardrail: {guardrail_type}") - - if custom_guardrail_callback is not None: - for scoping_param in ( - "skip_system_message_in_guardrail", - "skip_tool_message_in_guardrail", - "scan_only_tool_results", - ): - setattr(custom_guardrail_callback, scoping_param, getattr(litellm_params, scoping_param, None)) - scan_only_tool_results_enabled: Final = effective_scan_only_tool_results_for_guardrail( - custom_guardrail_callback - ) - if scan_only_tool_results_enabled and not custom_guardrail_callback.supports_scan_only_tool_results(): - raise ValueError( - f"Guardrail {guardrail['guardrail_name']}: scan_only_tool_results is enabled, but this " - "guardrail's role filtering never scans tool results, so no request content would ever " - "be scanned. Remove scan_only_tool_results or the guardrail's role-filtering option." - ) - if scan_only_tool_results_enabled and effective_skip_tool_message_for_guardrail(custom_guardrail_callback): - raise ValueError( - f"Guardrail {guardrail['guardrail_name']}: scan_only_tool_results and " - "skip_tool_message_in_guardrail are enabled together, which excludes every message from " - "scanning, so no request content would ever be scanned. Remove one of the two." - ) - _apply_configured_bool_overrides(custom_guardrail_callback, litellm_params) + created_callbacks: Final = self._create_callbacks( + guardrail=guardrail, + guardrail_type=guardrail_type, + litellm_params=litellm_params, + config_file_path=config_file_path, + llm_router=llm_router, + ) + for custom_guardrail_callback in created_callbacks: + _configure_callback_scoping(custom_guardrail_callback, guardrail["guardrail_name"], litellm_params) parsed_guardrail: Final = Guardrail( guardrail_id=guardrail.get("guardrail_id"), @@ -556,11 +555,44 @@ class InMemoryGuardrailHandler: # store references to the guardrail in memory self.IN_MEMORY_GUARDRAILS[guardrail_id] = parsed_guardrail - self.guardrail_id_to_custom_guardrail[guardrail_id] = custom_guardrail_callback + self.guardrail_id_to_custom_guardrail[guardrail_id] = created_callbacks[0] if created_callbacks else None + self.guardrail_id_to_sibling_callbacks[guardrail_id] = created_callbacks[1:] self._sources[guardrail_id] = source return parsed_guardrail + def _create_callbacks( + self, + guardrail: Guardrail, + guardrail_type: str, + litellm_params: LitellmParams, + config_file_path: str | None, + llm_router: Optional["Router"], + ) -> GuardrailCallbacks: + initializer: Final = guardrail_initializer_registry.get(guardrail_type) + if initializer: + import inspect + + sig: Final = inspect.signature(initializer) + if "llm_router" in sig.parameters: + return _as_callback_tuple(initializer(litellm_params, guardrail, llm_router)) + return _as_callback_tuple(initializer(litellm_params, guardrail)) + if isinstance(guardrail_type, str) and "." in guardrail_type: + return _as_callback_tuple( + self.initialize_custom_guardrail( + guardrail=guardrail, + guardrail_type=guardrail_type, + litellm_params=litellm_params, + config_file_path=config_file_path, + ) + ) + raise ValueError(f"Unsupported guardrail: {guardrail_type}") + + def _tracked_callbacks(self, guardrail_id: str) -> GuardrailCallbacks: + primary: Final = self.guardrail_id_to_custom_guardrail.get(guardrail_id) + siblings: Final = self.guardrail_id_to_sibling_callbacks.get(guardrail_id, ()) + return (() if primary is None else (primary,)) + siblings + def initialize_custom_guardrail( self, guardrail: Guardrail, @@ -630,10 +662,15 @@ class InMemoryGuardrailHandler: self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail self._sources[guardrail_id] = source - custom_guardrail_callback: Final = self.guardrail_id_to_custom_guardrail.get(guardrail_id) - if custom_guardrail_callback: - updated_litellm_params: Final = cast(LitellmParams, guardrail.get("litellm_params", {})) - custom_guardrail_callback.update_in_memory_litellm_params(litellm_params=updated_litellm_params) + tracked_callbacks: Final = self._tracked_callbacks(guardrail_id) + if not tracked_callbacks: + return + updated_litellm_params: Final = cast(LitellmParams, guardrail.get("litellm_params", {})) + tracked_callbacks[0].update_in_memory_litellm_params(litellm_params=updated_litellm_params) + for sibling_callback in tracked_callbacks[1:]: + sibling_stage = sibling_callback.event_hook + sibling_callback.update_in_memory_litellm_params(litellm_params=updated_litellm_params) + sibling_callback.event_hook = sibling_stage def delete_in_memory_guardrail(self, guardrail_id: str) -> None: """ @@ -648,11 +685,11 @@ class InMemoryGuardrailHandler: self.IN_MEMORY_GUARDRAILS.pop(guardrail_id, None) self._sources.pop(guardrail_id, None) - custom_guardrail_callback: Final = self.guardrail_id_to_custom_guardrail.pop(guardrail_id, None) - if custom_guardrail_callback is None: - return - - litellm.logging_callback_manager.remove_callback_from_all_lists(custom_guardrail_callback) + tracked_callbacks: Final = self._tracked_callbacks(guardrail_id) + self.guardrail_id_to_custom_guardrail.pop(guardrail_id, None) + self.guardrail_id_to_sibling_callbacks.pop(guardrail_id, None) + for custom_guardrail_callback in tracked_callbacks: + litellm.logging_callback_manager.remove_callback_from_all_lists(custom_guardrail_callback) def list_in_memory_guardrails(self) -> list[Guardrail]: """ diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index c08ca5b7783..5326074ad3c 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -27,6 +27,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import get_team_object, get_user_object +from litellm.proxy.auth.password_policy import validate_password_policy from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.user_api_key_cache import ( object_permission_cache_key, @@ -154,9 +155,10 @@ def _team_membership_table( return team_membership_table -def _hash_password_in_dict(data: dict) -> None: - """Hash password field in-place if present.""" +def _hash_password_in_dict(data: dict, general_settings: Mapping[str, object]) -> None: + """Validate and hash password field in-place if present.""" if "password" in data and data["password"] is not None: + validate_password_policy(data["password"], general_settings) data["password"] = hash_password(data["password"]) @@ -500,7 +502,7 @@ async def new_user( ``` """ try: - from litellm.proxy.proxy_server import _license_check, prisma_client + from litellm.proxy.proxy_server import _license_check, general_settings, prisma_client if prisma_client is None: raise HTTPException(status_code=400, detail=CommonProxyErrors.db_not_connected_error.value) @@ -548,7 +550,7 @@ async def new_user( # generate_key_helper_fn only forwards object_permission_id, so without this the entitlement # the caller sent would be dropped on the floor. data_json = await _set_object_permission(data_json=data_json, prisma_client=prisma_client) - _hash_password_in_dict(data_json) + _hash_password_in_dict(data_json, general_settings) teams = data.teams if teams is None: teams = check_if_default_team_set() @@ -1405,7 +1407,7 @@ async def _update_single_user_helper( Returns the updated user data or raises an exception on failure. """ - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name, prisma_client if prisma_client is None: raise Exception("Not connected to DB!") @@ -1420,7 +1422,7 @@ async def _update_single_user_helper( data_json: Final[dict] = user_request.model_dump(exclude_unset=True) non_default_values = _update_internal_user_params(data_json=data_json, data=user_request) - _hash_password_in_dict(non_default_values) + _hash_password_in_dict(non_default_values, general_settings) existing_user_row: BaseModel | None = None if user_request.user_id: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c3403cf477c..d7d20d168b5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3785,15 +3785,22 @@ async def info_key_fn_v2( @router.get("/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def info_key_fn( - key: str | None = fastapi.Query(default=None, description="Key in the request parameters"), + key: str | None = fastapi.Query( + default=None, + description=( + "Key to look up. Pass the key's sha256 hash so the raw key stays out of URLs and access " + "logs. Example key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa'" + ), + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Retrieve information about a key. Parameters: - - key: str | None (query parameter) - The key to look up. Accepts the plaintext key or its hash. - Defaults to the key in the Authorization header. + - key: str | None (query parameter) - The key to look up. Accepts the plaintext key or its hash; + prefer the hash, since a query parameter is recorded verbatim by any HTTP access log in front + of the proxy. Defaults to the key in the Authorization header. Returns: - key: str - The key that was looked up, echoed back as it was passed in @@ -3825,7 +3832,7 @@ async def info_key_fn( Example Curl: ``` - curl -X GET "http://0.0.0.0:4000/key/info?key=sk-test-example-key-123" \ + curl -X GET "http://0.0.0.0:4000/key/info?key=d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa" \ -H "Authorization: Bearer sk-1234" ``` diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 606569c5b8b..6b98d9f9a26 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -89,7 +89,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object from litellm.proxy.auth.auth_utils import ( _get_request_ip_address, - _has_user_setup_sso, + has_user_setup_sso, ) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -2617,7 +2617,7 @@ async def get_ui_settings(request: Request): _proxy_base_url: Final = os.getenv("PROXY_BASE_URL", None) _logout_url: Final = os.getenv("PROXY_LOGOUT_URL", None) _api_doc_base_url: Final = os.getenv("LITELLM_UI_API_DOC_BASE_URL", None) - _is_sso_enabled: Final = _has_user_setup_sso() + _is_sso_enabled: Final = has_user_setup_sso() disable_expensive_db_queries: Final = ( proxy_state.get_proxy_state_variable("spend_logs_row_count") > MAX_SPENDLOG_ROWS_TO_QUERY ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 9bc90260de1..bf07f4748ef 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -70,6 +70,15 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( validate_managed_files_requirement, validate_managed_id_requirement, ) +from litellm.proxy.openai_files_endpoints.general_upload_validation import ( + MB, + check_blocked_extension, + check_unsafe_filename, + check_upload_file_size, + coerce_optional_int_setting, + coerce_optional_str_list_setting, + raise_upload_validation_failure, +) from litellm.proxy.utils import ProxyLogging, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository from litellm.router import Router @@ -397,13 +406,23 @@ async def create_file( # descriptor and its disk blocks until the collector runs. spools: Final[list[BinaryIO]] = [] # mutable-ok: filled as the scan opens handles try: + unsafe_filename_failure: Final = check_unsafe_filename(file.filename) + if unsafe_filename_failure is not None: + raise_upload_validation_failure(unsafe_filename_failure) + + max_file_size_mb: Final = coerce_optional_int_setting(general_settings.get("max_file_size_mb")) + # Batch uploads can be gigabytes. Starlette has already spooled the upload # to disk, so stream from that handle instead of reading it into memory. - # Other uploads are small and stay in-memory bytes. + # Other uploads stay in-memory bytes, bounded to max_file_size_mb (plus one + # byte, to still tell "exactly at the limit" from "over it") when it is set, + # so an oversized upload cannot be read to completion before it is rejected. file_source: bytes | BinaryIO if purpose == "batch": await file.seek(0) file_source = file.file + elif max_file_size_mb is not None and max_file_size_mb > 0: + file_source = await file.read(max_file_size_mb * MB + 1) else: file_source = await file.read() custom_llm_provider = ( @@ -442,6 +461,15 @@ async def create_file( # Cast purpose to OpenAIFilesPurpose type purpose = cast(OpenAIFilesPurpose, purpose) + general_size_failure: Final = check_upload_file_size(file_source, max_file_size_mb) + if general_size_failure is not None: + raise_upload_validation_failure(general_size_failure) + + blocked_extensions: Final = coerce_optional_str_list_setting(general_settings.get("blocked_file_extensions")) + blocked_extension_failure: Final = check_blocked_extension(file.filename, blocked_extensions) + if blocked_extension_failure is not None: + raise_upload_validation_failure(blocked_extension_failure) + if purpose == "batch": batch_file_failure: Final = await asyncio.to_thread( check_batch_file_upload, diff --git a/litellm/proxy/openai_files_endpoints/general_upload_validation.py b/litellm/proxy/openai_files_endpoints/general_upload_validation.py new file mode 100644 index 00000000000..9d450cb5b8d --- /dev/null +++ b/litellm/proxy/openai_files_endpoints/general_upload_validation.py @@ -0,0 +1,150 @@ +""" +Upload validation applied to every purpose at POST /v1/files. + +batch_file_validation.py checks the JSONL shape of purpose="batch" uploads; this +module applies the same fast-fail-before-forwarding shape (size cap, blocked +extensions, path-traversal filenames) regardless of purpose. +""" + +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO, Final, NoReturn, assert_never + +from litellm.proxy._types import ProxyException +from litellm.proxy.common_utils.path_utils import safe_filename + +MB: Final = 1024 * 1024 + + +def coerce_optional_int_setting(raw: object) -> int | None: + """A general_settings value declared as an optional integer, e.g. max_file_size_mb. + + bool is an int subclass, so an explicit isinstance(raw, bool) exclusion is needed + or a YAML `true`/`false` would silently pass as 1/0. + """ + if raw is None: + return None + if isinstance(raw, int) and not isinstance(raw, bool): + return raw + raise TypeError(f"expected an integer, got {raw!r}") + + +def coerce_optional_str_list_setting(raw: object) -> tuple[str, ...]: + """A general_settings value declared as an optional list of strings, e.g. blocked_file_extensions.""" + if raw is None: + return () + if not isinstance(raw, list) or not all(isinstance(item, str) for item in raw): + raise TypeError(f"expected a list of strings, got {raw!r}") + return tuple(raw) + + +@dataclass(frozen=True, slots=True) +class UploadedFileTooLarge: + size_bytes: int + limit_mb: int + + +@dataclass(frozen=True, slots=True) +class UploadedFileBlockedExtension: + extension: str + + +@dataclass(frozen=True, slots=True) +class UploadedFileUnsafeFilename: + filename: str + + +UploadValidationFailure = UploadedFileTooLarge | UploadedFileBlockedExtension | UploadedFileUnsafeFilename + + +def _file_size_bytes(file_source: bytes | BinaryIO) -> int: + if isinstance(file_source, bytes): + return len(file_source) + original_position: Final = file_source.tell() + file_source.seek(0, 2) + size: Final = file_source.tell() + file_source.seek(original_position) + return size + + +def check_upload_file_size( + file_source: bytes | BinaryIO, + max_file_size_mb: int | None, +) -> UploadedFileTooLarge | None: + if max_file_size_mb is None or max_file_size_mb <= 0: + return None + size_bytes: Final = _file_size_bytes(file_source) + if size_bytes > max_file_size_mb * MB: + return UploadedFileTooLarge(size_bytes=size_bytes, limit_mb=max_file_size_mb) + return None + + +def check_blocked_extension( + filename: str | None, + blocked_extensions: tuple[str, ...], +) -> UploadedFileBlockedExtension | None: + if not blocked_extensions or not filename: + return None + try: + extension: Final = Path(safe_filename(filename)).suffix.lower() + except ValueError: + return None + # The uploaded name's extension is normalized above; blocked_extensions comes + # straight from config.yaml or the DB and is normalized here too, so a + # differently-cased entry (".EXE") still catches a lowercase upload. + normalized_blocked: Final = frozenset(item.lower() for item in blocked_extensions) + if extension and extension in normalized_blocked: + return UploadedFileBlockedExtension(extension=extension) + return None + + +def check_unsafe_filename(filename: str | None) -> UploadedFileUnsafeFilename | None: + """Reject a filename before it can influence any storage path or backend call. + + Only flags a genuine traversal component ("..") or a null byte, so an ordinary + name like "report.v2.pdf" or ".env" is never rejected. + """ + if not filename: + return None + if "\x00" in filename: + return UploadedFileUnsafeFilename(filename=filename) + normalized: Final = filename.replace("\\", "/") + if any(part == ".." for part in normalized.split("/")): + return UploadedFileUnsafeFilename(filename=filename) + return None + + +def raise_upload_validation_failure(failure: UploadValidationFailure) -> NoReturn: + match failure: + case UploadedFileTooLarge(size_bytes=size_bytes, limit_mb=limit_mb): + raise ProxyException( + message=( + f"Uploaded file exceeds the configured max_file_size_mb of {limit_mb} MB " + f"(read stopped at {size_bytes / MB:.1f} MB). The file was not forwarded to the provider." + ), + type="invalid_request_error", + param="file", + code=413, + ) + case UploadedFileBlockedExtension(extension=extension): + raise ProxyException( + message=( + f"File extension '{extension}' is blocked by this proxy's blocked_file_extensions " + "setting. The file was not forwarded to the provider." + ), + type="invalid_request_error", + param="file", + code=400, + ) + case UploadedFileUnsafeFilename(filename=filename): + raise ProxyException( + message=( + f"Filename '{filename}' is not allowed: directory traversal sequences are not " + "permitted in uploaded file names. The file was not forwarded to the provider." + ), + type="invalid_request_error", + param="file", + code=400, + ) + case _: + assert_never(failure) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d52f05f6166..1e0792cb8a5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -310,6 +310,7 @@ from litellm.proxy.auth.model_checks import ( get_mcp_server_ids, get_team_models, ) +from litellm.proxy.auth.password_policy import validate_password_policy from litellm.proxy.auth.user_api_key_auth import ( _fetch_global_spend_with_event_coordination, user_api_key_auth, @@ -4743,6 +4744,63 @@ class ProxyConfig: ) return coordination_redis_cache + @staticmethod + async def _init_coordination_redis_env_fallback(litellm_settings: Mapping[str, object]) -> RedisCache | None: + """ + Last-resort coordination Redis, tried after an explicit + `general_settings.coordination_redis` block and `litellm_settings.cache` + have both had a chance to resolve one. Without this, a deployment that + only exports REDIS_HOST/REDIS_PORT (no cache block, no coordination_redis + block) gets NO cross-pod coordination at all: spend counters, budget-window + enforcement, and the reset_spend cache-eviction broadcast all silently stay + per-pod local, so a key reset on one pod never clears another pod's stale + enforcement. + + Unlike the explicit block and cache-backend paths (a deliberate opt-in, so a + bad connection target or a malformed REDIS_CLUSTER_NODES/REDIS_SENTINEL_NODES + value should fail loudly), this one is inferred from bare env vars that may be + set for an unrelated reason -- e.g. a REDIS_HOST left over from a different + job/service, or a REDIS_CLUSTER_NODES value nothing here ever asked to be + parsed. Wrongly guessing "coordination available" must not turn a previously + harmless in-memory-only proxy into one that fails to boot or raises on every + cache write, so a malformed value or a failed/slow ping are both treated the + same as no REDIS_* vars at all. + """ + try: + env_coordination_redis_cache: Final = _build_redis_usage_cache_from_environment() + except Exception as e: # noqa: BLE001 # a malformed inferred Redis env var must not block startup + verbose_proxy_logger.warning( + "coordination_redis: could not build a Redis client from REDIS_* environment variables " + "(%s); cross-pod coordination stays in-memory. Set general_settings.coordination_redis " + "explicitly to require it.", + e, + ) + return None + if env_coordination_redis_cache is None: + return None + try: + reachable: Final = await asyncio.wait_for(env_coordination_redis_cache.ping(), timeout=2.0) + except Exception as e: # noqa: BLE001 # an unreachable inferred Redis must not block startup or writes + verbose_proxy_logger.warning( + "coordination_redis: REDIS_* environment variables named a Redis that is not reachable " + "(%s); cross-pod coordination stays in-memory. Set general_settings.coordination_redis " + "explicitly to require it.", + e, + ) + return None + if not reachable: + return None + _attach_redis_usage_cache( + env_coordination_redis_cache, + enable_redis_auth_cache=litellm_settings.get("enable_redis_auth_cache", False) is True, + ) + verbose_proxy_logger.info( + "coordination_redis: using a standalone Redis built from REDIS_* " + "environment variables for usage tracking, rate limiting, and " + "cross-pod coordination." + ) + return env_coordination_redis_cache + def _init_cache( self, cache_params: dict, @@ -5411,6 +5469,13 @@ class ProxyConfig: reset_audit_log_callback_cache() _in_memory_loggers[:] = [cb for cb in _in_memory_loggers if not isinstance(cb, S3V2Logger)] + if redis_usage_cache is None: + env_coordination_redis_cache: Final = await self._init_coordination_redis_env_fallback( + litellm_settings=litellm_settings + ) + if env_coordination_redis_cache is not None: + _set_redis_usage_cache(env_coordination_redis_cache) + ## GENERAL SERVER SETTINGS (e.g. master key,..) # do this after initializing litellm, to ensure sentry logging works for proxylogging general_settings = config.get("general_settings", {}) if general_settings is None: @@ -6644,6 +6709,12 @@ class ProxyConfig: if "max_batch_file_size_mb" not in self._yaml_general_settings_keys: general_settings["max_batch_file_size_mb"] = _general_settings.get("max_batch_file_size_mb") + if "max_file_size_mb" not in self._yaml_general_settings_keys: + general_settings["max_file_size_mb"] = _general_settings.get("max_file_size_mb") + + if "blocked_file_extensions" not in self._yaml_general_settings_keys: + general_settings["blocked_file_extensions"] = _general_settings.get("blocked_file_extensions") + ## ALERTING ARGS ## if "alerting_args" in _general_settings: general_settings["alerting_args"] = _general_settings["alerting_args"] @@ -15236,6 +15307,7 @@ async def login(request: Request): password=password, master_key=master_key, prisma_client=prisma_client, + general_settings=general_settings, ) # Create UI token object @@ -15310,6 +15382,7 @@ async def login_v2(request: Request): password=password, master_key=master_key, prisma_client=prisma_client, + general_settings=general_settings, ) returned_ui_token_object: Final = create_ui_token_object( @@ -15380,6 +15453,7 @@ async def login_v3(request: Request): password=password, master_key=master_key, prisma_client=prisma_client, + general_settings=general_settings, ) returned_ui_token_object: Final = create_ui_token_object( @@ -15749,6 +15823,7 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): detail={"error": "Invalid onboarding session for invitation link."}, ) + validate_password_policy(data.password, general_settings) hashed_pw: Final = hash_password(data.password) current_time = litellm.utils.get_utc_datetime() async with prisma_client.db.tx() as tx: @@ -16412,6 +16487,8 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "global_max_parallel_requests": "Integer", "max_request_size_mb": "Integer", "max_batch_file_size_mb": "Integer", + "max_file_size_mb": "Integer", + "blocked_file_extensions": "List", "max_response_size_mb": "Integer", "proxy_config_reload_interval_seconds": "Integer", "pass_through_endpoints": "PydanticModel", diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 41c65b1d5c5..c48750cea72 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1204,7 +1204,10 @@ async def get_global_spend_report( ), api_key: str | None = fastapi.Query( default=None, - description="View spend for a specific api_key. Example api_key='sk-1234", + description=( + "View spend for a specific api_key. Pass the key's sha256 hash so the raw key stays " + "out of URLs and access logs. Example api_key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa'" + ), ), internal_user_id: str | None = fastapi.Query( default=None, @@ -1685,7 +1688,11 @@ async def get_key_spend_report( api_key: Annotated[ str | None, fastapi.Query( - description="View spend for a specific api_key. Proxy admin only; other callers are scoped to their own key." + description=( + "View spend for a specific api_key. Proxy admin only; other callers are scoped to their " + "own key. Pass the key's sha256 hash so the raw key stays out of URLs and access logs. " + "Example api_key='d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa'" + ) ), ] = None, ) -> Sequence[Mapping[str, object]]: @@ -2945,7 +2952,7 @@ async def view_spend_logs( Example Request for specific api_key ``` - curl -X GET "http://0.0.0.0:8000/spend/logs?api_key=sk-test-example-key-123" \ + curl -X GET "http://0.0.0.0:8000/spend/logs?api_key=d5345c0ecc68ae6295c69f91926b2bd379e25481a40c34b5884d157a9f65d8fa" \ -H "Authorization: Bearer sk-1234" ``` diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 52258602581..c12d071dd36 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -21,6 +21,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys +from litellm.proxy._experimental.mcp_server.tool_search import MCP_TOOL_SEARCH_SETTINGS_KEY from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.config_resolvers.sso import ( @@ -38,6 +39,7 @@ from litellm.repositories.table_repositories import ( UISettingsRepository, ) from litellm.repositories.team_repository import TeamRepository +from litellm.types.mcp import MCPToolSearchSettings from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, SSOConfig, @@ -448,6 +450,10 @@ class MCPSemanticFilterSettingsResponse(SettingsResponse): """Response model for MCP semantic filter settings""" +class MCPToolSearchSettingsResponse(SettingsResponse): + """Response model for native MCP tool search settings""" + + @router.get( "/get/allowed_ips", tags=["Budget & Spend Tracking"], @@ -835,7 +841,7 @@ async def update_default_team_member_budget(teams: list[NewUserRequestTeam], use async def _update_litellm_setting( - settings: DefaultInternalUserParams | DefaultTeamSSOParams | MCPSemanticFilterSettings, + settings: DefaultInternalUserParams | DefaultTeamSSOParams | MCPSemanticFilterSettings | MCPToolSearchSettings, settings_key: str, success_message: str, user_api_key_dict: UserAPIKeyAuth, @@ -861,7 +867,7 @@ async def _update_litellm_setting( detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) - in_memory_var: Final = settings.model_dump(exclude_none=True) + in_memory_var: Final = settings.model_dump(mode="json", exclude_none=True) # Load existing config first, then set in-memory value after, # because get_config() may overwrite litellm. with stale DB values @@ -1359,6 +1365,59 @@ async def update_mcp_semantic_filter_settings( return result +@router.get( + "/get/mcp_tool_search_settings", + tags=["Settings"], # mutable-ok: FastAPI's route decorator only accepts a list + dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI's route decorator only accepts a list + response_model=MCPToolSearchSettingsResponse, +) +async def get_mcp_tool_search_settings( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> Mapping[str, object]: + """ + Get the `litellm_settings.mcp_tool_search` configuration used by the native `mcp_tool_search` virtual tool. + """ + from litellm.proxy.proxy_server import prisma_client, proxy_config + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected. Please connect a database.") + + config: Final = await proxy_config.get_config() + + return await _get_settings_with_schema( + settings_key=MCP_TOOL_SEARCH_SETTINGS_KEY, + settings_class=MCPToolSearchSettings, + config=config, + ) + + +@router.patch( + "/update/mcp_tool_search_settings", + tags=["Settings"], # mutable-ok: FastAPI's route decorator only accepts a list + dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI's route decorator only accepts a list +) +async def update_mcp_tool_search_settings( + settings: MCPToolSearchSettings, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> Mapping[str, object]: + """ + Update `litellm_settings.mcp_tool_search` in the database. + Settings will be picked up by all pods within approximately 10 seconds via background polling. + """ + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Only proxy admins can update MCP tool search settings.", + ) + + return await _update_litellm_setting( + settings=settings, + settings_key=MCP_TOOL_SEARCH_SETTINGS_KEY, + success_message="MCP tool search settings updated successfully. Changes will be applied across all pods within 10 seconds.", + user_api_key_dict=user_api_key_dict, + ) + + UI_SETTINGS_CACHE_KEY: Final = "ui_settings:settings_dict" UI_SETTINGS_CACHE_TTL: Final = 600 # 10 minutes @@ -1594,7 +1653,10 @@ async def update_ui_settings( tags=["UI Theme Settings"], dependencies=[Depends(user_api_key_auth)], ) -async def upload_logo(file: UploadFile = File(...)): +async def upload_logo( + file: UploadFile = File(...), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Upload a custom logo for the admin UI. Accepts image files (PNG, JPG, JPEG, SVG) and stores them for use in the UI. @@ -1602,6 +1664,12 @@ async def upload_logo(file: UploadFile = File(...)): import os from pathlib import Path + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Only proxy admins can upload a UI logo.", + ) + # Validate file type allowed_extensions: Final = {".png", ".jpg", ".jpeg", ".svg"} file_extension: Final = Path(file.filename or "").suffix.lower() @@ -1612,9 +1680,11 @@ async def upload_logo(file: UploadFile = File(...)): detail=f"Invalid file type. Allowed types: {', '.join(allowed_extensions)}", ) - # Validate file size (max 5MB) - file_content: Final = await file.read() - if len(file_content) > 5 * 1024 * 1024: # 5MB + # Read bounded to one byte past the limit, so an oversized upload is never + # fully buffered in memory before being rejected. + max_logo_size_bytes: Final = 5 * 1024 * 1024 + file_content: Final = await file.read(max_logo_size_bytes + 1) + if len(file_content) > max_logo_size_bytes: raise HTTPException(status_code=400, detail="File size too large. Maximum size is 5MB.") # Create uploads directory if it doesn't exist diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 051d36c4d0f..cab2bd6d9db 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1524,6 +1524,7 @@ class ProxyLogging: prompt_label=data.pop("prompt_label", None) or {}, prompt_version=data.pop("prompt_version", None) or {}, request_kwargs=data, + injected_for_every_deployment=True, ) data.update(optional_params) diff --git a/litellm/router.py b/litellm/router.py index b1285f9e55e..9353e391e76 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -30,7 +30,7 @@ import anyio import httpx import openai from openai import AsyncOpenAI -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import overload import litellm @@ -383,6 +383,28 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) return False +_NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) +_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]: + """ + Realtime client-secret requests carry the model inside ``session`` as well, and the caller's copy of it still + holds the pre-routing model group name, so it has to follow the deployment the router just picked. + + Returns kwargs to merge into the downstream call, empty when there is no session model to resolve. + """ + try: + typed_session: Final = _SESSION_ADAPTER.validate_python(session) + except ValidationError: + return _NO_SESSION_KWARGS + if "model" not in typed_session: + return _NO_SESSION_KWARGS + return MappingProxyType( + {"session": {**typed_session, "model": model_name}} # mutable-ok: callees deepcopy and JSON-dump session + ) + + # Router._aanthropic_messages_streaming_iterator buffers lifecycle chunks # until real content commits the primary stream; a hostile or slow-starting # upstream that never emits content or an error could otherwise grow that @@ -4043,6 +4065,7 @@ class Router: prompt_variables=prompt_variables, prompt_label=prompt_label, request_kwargs=kwargs, + injected_for_every_deployment=True, ) # Filter out prompt management specific parameters from data before merging @@ -4930,6 +4953,7 @@ class Router: "caching": self.cache_responses, **kwargs, "model": model_name, + **_with_router_resolved_session_model(kwargs.get("session"), model_name), } # Only set custom_llm_provider if it's not None if custom_llm_provider is not None: diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 1b8baf2da09..a59fcb1bcb5 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal from urllib.parse import urlsplit import httpx -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, Field from typing_extensions import TypedDict from litellm.types.llms.base import HiddenParams @@ -91,6 +91,33 @@ class MCPPublicServer(BaseModel): mcp_info: dict[str, Any] | None = None +class MCPToolSearchSettings(BaseModel): + """`litellm_settings.mcp_tool_search`: how the native `mcp_tool_search` virtual tool ranks the caller's tools.""" + + model_config = ConfigDict(frozen=True) + + embedding_model: str | None = Field( + default=None, + description="Embedding model from model_list used to rank tools by meaning. Unset keeps keyword matching.", + ) + top_k: int = Field( + default=5, + ge=1, + le=100, + description="Most ranked tools a search returns. A smaller top_k in the tool call wins. Core tools do not count.", + ) + similarity_threshold: float = Field( + default=0.0, + ge=0.0, + le=1.0, + description="Lowest cosine similarity a tool needs to appear in semantic results (0.0 = no cutoff).", + ) + core_tools: tuple[str, ...] = Field( + default=(), + description="Tool names always returned first when the caller can access them, e.g. `my_server-get_rates`.", + ) + + # OAuth 2.0 token-endpoint client authentication method (RFC 6749 section 2.3.1). MCPTokenEndpointAuthMethod = Literal["client_secret_basic", "client_secret_post"] diff --git a/litellm/utils.py b/litellm/utils.py index 252b6756937..ba456fc353b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5106,49 +5106,6 @@ def get_response_string(response_obj: ModelResponse | ModelResponseStream) -> st return "".join(response_parts) -def get_api_key(llm_provider: str, dynamic_api_key: str | None): - api_key = dynamic_api_key or litellm.api_key - # openai - if llm_provider == "openai" or llm_provider == "text-completion-openai": - api_key = api_key or litellm.openai_key or get_secret("OPENAI_API_KEY") - # anthropic - elif llm_provider == "anthropic" or llm_provider == "anthropic_text": - api_key = api_key or litellm.anthropic_key or get_secret("ANTHROPIC_API_KEY") - # ai21 - elif llm_provider == "ai21": - api_key = api_key or litellm.ai21_key or get_secret("AI21_API_KEY") - # aleph_alpha - elif llm_provider == "aleph_alpha": - api_key = api_key or litellm.aleph_alpha_key or get_secret("ALEPH_ALPHA_API_KEY") - # baseten - elif llm_provider == "baseten": - api_key = api_key or litellm.baseten_key or get_secret("BASETEN_API_KEY") - # cohere - elif llm_provider == "cohere" or llm_provider == "cohere_chat": - api_key = api_key or litellm.cohere_key or get_secret("COHERE_API_KEY") - # huggingface - elif llm_provider == "huggingface": - api_key = api_key or litellm.huggingface_key or get_secret("HUGGINGFACE_API_KEY") - # nlp_cloud - elif llm_provider == "nlp_cloud": - api_key = api_key or litellm.nlp_cloud_key or get_secret("NLP_CLOUD_API_KEY") - # replicate - elif llm_provider == "replicate": - api_key = api_key or litellm.replicate_key or get_secret("REPLICATE_API_KEY") - # together_ai - elif llm_provider == "together_ai": - api_key = ( - api_key or litellm.togetherai_api_key or get_secret("TOGETHERAI_API_KEY") or get_secret("TOGETHER_AI_TOKEN") - ) - # nebius - elif llm_provider == "nebius": - api_key = api_key or litellm.nebius_key or get_secret("NEBIUS_API_KEY") - # wandb - elif llm_provider == "wandb": - api_key = api_key or litellm.wandb_key or get_secret("WANDB_API_KEY") - return api_key - - def get_utc_datetime(): import datetime as dt from datetime import datetime diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2846d12db6e..d8a8f84b032 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -33093,6 +33093,88 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "meta/muse-spark-1.3": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "meta/muse-spark-1.3-contributor": { + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 1e-07, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "meta_llama/Llama-3.3-70B-Instruct": { "litellm_provider": "meta_llama", "max_input_tokens": 128000, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 9b1cc977a64..ae91b711e13 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -9,7 +9,7 @@ "limit": 809 }, "ANN201": { - "limit": 2001 + "limit": 2000 }, "ANN202": { "limit": 835 @@ -108,7 +108,7 @@ "limit": 3 }, "F401": { - "limit": 13 + "limit": 12 }, "LOG015": { "limit": 5 @@ -147,7 +147,7 @@ "limit": 3 }, "PLR1714": { - "limit": 256 + "limit": 253 }, "PLW0127": { "limit": 57 diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 790956156b0..0578dc60119 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -64,6 +64,8 @@ IGNORE_FUNCTIONS = [ "_flatten_form_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible). "_flatten_form_data_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible). "_json_safe", # max depth set (_MAX_DEPTH) plus a seen-ids cycle guard for self-referential input. + "_redact_agent_params_tree", # max depth set (default 10), same shape as _redact_sensitive_litellm_params. + "_restore_redacted_nested_value", # max depth set (default 10), mirrors _redact_agent_params_tree on the write side. ] diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 3f9f2bacdd3..0c362db8853 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -333,6 +333,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters( timeout=None, client=None, _is_async=False, + router: "litellm.Router | None" = None, ): litellm_params_dict = ( litellm_params.model_dump(exclude_none=False) diff --git a/tests/rust-python-harness/README.md b/tests/rust-python-harness/README.md new file mode 100644 index 00000000000..e94ac87c3b3 --- /dev/null +++ b/tests/rust-python-harness/README.md @@ -0,0 +1,145 @@ +# Rust ↔ Python SDK parity harness + +This folder is the operator-facing harness for the Rust migration test plan. It runs pytest normally, listens to test events in-process, and redraws a live matrix grouped by testing strategy and SDK-level function. + +The matrix always has these SDK columns: + +- `ocr / aocr` +- `messages / amessages` +- `responses / aresponses` +- `count_tokens` + +The harness has three deliberately broad test-strategy folders: + +| Strategy | Folder | +| --- | --- | +| Public SDK parity over generated and recorded inputs | [`e2e_fuzz_tests/`](e2e_fuzz_tests/) | +| Focused tests of Rust-owned behavior | [`unit_tests_rust/`](unit_tests_rust/) | +| Isolated transform and Python-to-Rust helper coverage | [`validate_sub_methods/`](validate_sub_methods/) | + +## Run it + +From the repository root: + +```bash +poetry run python -m tests.rust-python-harness +``` + +The default runs every configured test once and updates all matching cells in real time. Narrow a run by strategy, SDK function, or both: + +```bash +poetry run python -m tests.rust-python-harness --strategy e2e_fuzz_tests +poetry run python -m tests.rust-python-harness --function messages +poetry run python -m tests.rust-python-harness --strategy validate_sub_methods --function ocr +``` + +For a guided run, use the interactive picker. It asks which strategy rows and SDK +function columns to include, then hands the terminal to the live dashboard. It never +captures keys while tests are running, so Ctrl-C and pytest debugging remain safe. + +```bash +poetry run python -m tests.rust-python-harness --interactive +``` + +Useful operator options: + +```bash +# Inspect coverage and pytest selectors without running anything. +poetry run python -m tests.rust-python-harness --list + +# Stable line-oriented output for CI logs or redirected output. +poetry run python -m tests.rust-python-harness --plain + +# Measure Python reference lines exercised by this parity run and build an HTML heatmap. +poetry run python -m tests.rust-python-harness --coverage + +# Forward pytest options. Use the equals form when the value begins with a dash. +poetry run python -m tests.rust-python-harness --pytest-arg=-x +``` + +The process returns pytest's exit code. A configured selector that collects no test is also a failure. A planned cell has no selector yet and does not fail the run. + +The dashboard adapts to narrow terminals, shows elapsed time and unique-test progress, +and prints the three slowest tests when the run ends. Each failure includes a focused +`poetry run pytest ... -q` command. Redirected output and CI automatically use the +line-oriented plain renderer; `--plain` lets you opt into it locally. + +The final screen includes a confidence score for every SDK section. It is the direct +ratio of required strategy rows with passing evidence, such as `1/3 = 33%`; High means +all required strategies passed, Medium means some passed, and Low means none passed. +This behavioral score is intentionally shown separately from Python and Rust LOC. + +Coverage reports are written outside the three strategy folders at +`target/rust-python-harness/`. Open `python-html/index.html` to inspect executed and +missing Python lines; `python.json` and `python.xml` are available for automation. +Coverage is finalized after pytest exits, because worker processes must flush their +data first. + +## Port coverage and confidence + +Treat these as separate signals instead of one ambiguous coverage percentage: + +| Signal | Tool | What it proves | +| --- | --- | --- | +| Python reference LOC | `coverage.py` / `pytest-cov` via `--coverage` | The mapped Python behavior ran | +| Rust port LOC | `cargo-llvm-cov` | The mapped Rust implementation ran | +| Parity contracts | This harness matrix | Python and Rust had the same observable behavior | + +`validate_sub_methods/` owns the future source-section inventory that maps a stable +Python qualified symbol to its Rust symbol. That inventory is the denominator for +per-function rollups; raw coverage for the entire LiteLLM repository would obscure +the port's real gaps. `unit_tests_rust/` owns direct `cargo-llvm-cov` runs, while +`e2e_fuzz_tests/` owns behavioral parity and fuzz-case counts. Keep Python, Rust, and +parity percentages visible side by side and label section confidence High only when +the mapped implementation exists, every required strategy passes, and both sides meet +their LOC thresholds. Generated Rust LCOV/HTML and the combined index also belong in +`target/rust-python-harness/`, not in a fourth strategy folder. + +## Read the matrix + +| Mark | Meaning | +| --- | --- | +| `✓` | All collected tests passed | +| `✗` | At least one test failed | +| `!` | Test setup or teardown failed | +| `↷` | All collected tests skipped | +| `?` | A configured selector did not collect a test | +| `—` | Strategy is planned but has no test yet | +| `n/a` | Strategy does not apply to this SDK function | +| `◐` | The configured tests cover only part of the TDD's parity contract | + +The initial end-to-end entries deliberately show `◐`: the repository has Rust bridge tests for OCR, Messages, and Responses websocket plumbing, but those are not yet frozen-Python-oracle comparisons. The remaining TDD cells stay visible as planned work instead of disappearing from a green summary. + +## Attach parity tests + +Each of the three folders contains a concise `README.md` and a `strategy.json`. Add a pytest file or node ID to the appropriate SDK function's `selectors` list: + +```json +{ + "coverage": "complete", + "selectors": [ + "tests/rust-python-harness/validate_sub_methods/test_messages.py" + ] +} +``` + +Selectors use the same syntax as pytest. A file selector aggregates every test in the file; a node selector can target one test or parametrized family. The runner deduplicates selectors, so one test may intentionally prove more than one cell without executing twice. + +Use these coverage values: + +- `complete`: implements the full strategy contract for that SDK function. +- `partial`: useful coverage exists, but the TDD contract is not fully proven. +- `planned`: no runnable parity test exists yet. +- `not_applicable`: the strategy cannot apply, such as streaming for OCR. + +Keep comparison mechanics in shared harness modules and provider/function facts in the owning strategy folder. A Python/Rust mismatch is a test failure; do not normalize away observable return types, exception classes, private response fields, chunk ordering, or callback payload differences merely to make a cell green. + +## Architecture + +- `catalog.py` validates and loads every strategy manifest. +- `models.py` owns typed strategy, case, coverage, and run-state models. +- `runner.py` maps live pytest events back to one or more matrix cells. +- `ui.py` renders the interactive Rich dashboard and a dependency-free plain fallback. +- `cli.py` handles filtering and preserves pytest exit semantics. + +The harness is driven from Python, matching the SDK surface and existing test tooling. Rust remains responsible for the implementation under comparison; the harness does not move provider semantics into the PyO3 bridge. diff --git a/tests/rust-python-harness/__init__.py b/tests/rust-python-harness/__init__.py new file mode 100644 index 00000000000..70362674d2b --- /dev/null +++ b/tests/rust-python-harness/__init__.py @@ -0,0 +1,5 @@ +"""Interactive Rust/Python SDK parity test harness.""" + +from .catalog import load_catalog + +__all__ = ["load_catalog"] diff --git a/tests/rust-python-harness/__main__.py b/tests/rust-python-harness/__main__.py new file mode 100644 index 00000000000..bfdcd0c1158 --- /dev/null +++ b/tests/rust-python-harness/__main__.py @@ -0,0 +1,4 @@ +from .cli import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/rust-python-harness/catalog.py b/tests/rust-python-harness/catalog.py new file mode 100644 index 00000000000..e23b9b125f0 --- /dev/null +++ b/tests/rust-python-harness/catalog.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from .models import Coverage, HarnessCase, SDK_FUNCTIONS, Strategy + +STRATEGIES_ROOT = Path(__file__).parent + + +def _require_string(value: Any, field: str, source: Path) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{source}: {field} must be a non-empty string") + return value + + +def _load_strategy(source: Path) -> Strategy: + with source.open(encoding="utf-8") as stream: + data = json.load(stream) + + strategy_id = _require_string(data.get("id"), "id", source) + label = _require_string(data.get("label"), "label", source) + description = _require_string(data.get("description"), "description", source) + order = data.get("order") + if not isinstance(order, int): + raise ValueError(f"{source}: order must be an integer") + function_data = data.get("functions") + if not isinstance(function_data, dict): + raise ValueError(f"{source}: functions must be an object") + + missing = set(SDK_FUNCTIONS) - set(function_data) + extra = set(function_data) - set(SDK_FUNCTIONS) + if missing or extra: + raise ValueError( + f"{source}: functions must exactly match {SDK_FUNCTIONS}; missing={missing}, extra={extra}" + ) + + cases: list[HarnessCase] = [] + for sdk_function in SDK_FUNCTIONS: + case_data = function_data[sdk_function] + if not isinstance(case_data, dict): + raise ValueError(f"{source}: functions.{sdk_function} must be an object") + try: + coverage = Coverage(case_data.get("coverage")) + except ValueError as exc: + raise ValueError(f"{source}: invalid coverage for {sdk_function}") from exc + selectors = case_data.get("selectors", []) + if not isinstance(selectors, list) or not all( + isinstance(item, str) and item for item in selectors + ): + raise ValueError( + f"{source}: selectors for {sdk_function} must be a list of strings" + ) + if coverage is Coverage.NOT_APPLICABLE and selectors: + raise ValueError( + f"{source}: not_applicable case {sdk_function} cannot have selectors" + ) + cases.append( + HarnessCase( + strategy_id=strategy_id, + strategy_label=label, + sdk_function=sdk_function, + coverage=coverage, + selectors=tuple(selectors), + note=str(case_data.get("note", "")), + ) + ) + + return Strategy( + order=order, + id=strategy_id, + label=label, + description=description, + directory=source.parent, + cases=tuple(cases), + ) + + +def load_catalog(root: Path = STRATEGIES_ROOT) -> tuple[Strategy, ...]: + sources = sorted(root.glob("*/strategy.json")) + if not sources: + raise ValueError(f"No strategy manifests found below {root}") + strategies = tuple( + sorted( + (_load_strategy(source) for source in sources), + key=lambda strategy: strategy.order, + ) + ) + ids = [strategy.id for strategy in strategies] + if len(ids) != len(set(ids)): + raise ValueError(f"Duplicate strategy id in {root}") + return strategies diff --git a/tests/rust-python-harness/cli.py b/tests/rust-python-harness/cli.py new file mode 100644 index 00000000000..41c46f9613a --- /dev/null +++ b/tests/rust-python-harness/cli.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import argparse +import importlib.util +from collections.abc import Sequence +from pathlib import Path + +from .catalog import load_catalog +from .models import HarnessCase, Strategy +from .runner import run_pytest +from .ui import make_dashboard + +REPO_ROOT = Path(__file__).resolve().parents[2] +COVERAGE_ROOT = REPO_ROOT / "target" / "rust-python-harness" + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="rust-python-harness", + description="Run Rust/Python parity tests with a live strategy-by-SDK-function dashboard.", + ) + parser.add_argument( + "-i", + "--interactive", + action="store_true", + help="pick strategies and SDK functions in a guided terminal menu", + ) + parser.add_argument( + "--list", action="store_true", help="show the catalog without running tests" + ) + parser.add_argument( + "--strategy", + action="append", + default=[], + metavar="ID", + help="run only this strategy", + ) + parser.add_argument( + "--function", + action="append", + default=[], + dest="sdk_functions", + choices=("ocr", "messages", "responses", "count_tokens"), + help="run only this SDK function", + ) + parser.add_argument( + "--plain", + action="store_true", + help="disable the interactive terminal dashboard", + ) + parser.add_argument( + "--coverage", + action="store_true", + help="write Python reference LOC reports (HTML, JSON, and XML)", + ) + parser.add_argument( + "--pytest-arg", + action="append", + default=[], + metavar="ARG", + help="append an argument to pytest (repeatable, for example --pytest-arg=-x)", + ) + return parser + + +def _coverage_pytest_args(output_root: Path = COVERAGE_ROOT) -> tuple[str, ...]: + output_root.mkdir(parents=True, exist_ok=True) + return ( + "--cov=litellm", + "--cov-context=test", + f"--cov-report=json:{output_root / 'python.json'}", + f"--cov-report=xml:{output_root / 'python.xml'}", + f"--cov-report=html:{output_root / 'python-html'}", + ) + + +def _pick_values( + title: str, options: Sequence[tuple[str, str]], input_fn=input +) -> set[str]: + print(f"\n{title} (Enter = all)") + for index, (value, label) in enumerate(options, start=1): + print(f" {index:>2}. {label} [{value}]") + while True: + answer = input_fn("Choose numbers, comma-separated: ").strip() + if not answer: + return set() + try: + indexes = {int(part.strip()) for part in answer.split(",")} + except ValueError: + print("Please enter numbers separated by commas.") + continue + if indexes and all(1 <= index <= len(options) for index in indexes): + return {options[index - 1][0] for index in indexes} + print(f"Choose values from 1 to {len(options)}.") + + +def _interactive_filters(strategies: Sequence[Strategy]) -> tuple[set[str], set[str]]: + strategy_ids = _pick_values( + "Testing strategies", [(strategy.id, strategy.label) for strategy in strategies] + ) + sdk_functions = _pick_values( + "SDK functions", + [(name, name) for name in ("ocr", "messages", "responses", "count_tokens")], + ) + return strategy_ids, sdk_functions + + +def _select( + strategies: Sequence[Strategy], strategy_ids: set[str], sdk_functions: set[str] +) -> tuple[HarnessCase, ...]: + known_ids = {strategy.id for strategy in strategies} + unknown = strategy_ids - known_ids + if unknown: + raise ValueError(f"Unknown strategy: {', '.join(sorted(unknown))}") + return tuple( + case + for strategy in strategies + if not strategy_ids or strategy.id in strategy_ids + for case in strategy.cases + if not sdk_functions or case.sdk_function in sdk_functions + ) + + +def _print_catalog(strategies: Sequence[Strategy]) -> None: + for strategy in strategies: + print(f"{strategy.id:20} {strategy.label}") + for case in strategy.cases: + selectors = ( + ", ".join(case.selectors) if case.selectors else "no test configured" + ) + print(f" {case.sdk_function:12} {case.coverage.value:14} {selectors}") + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parser().parse_args(argv) + if args.coverage and importlib.util.find_spec("pytest_cov") is None: + _parser().error( + "--coverage requires the project's pytest-cov dependency; run with " + "`poetry run python -m tests.rust-python-harness --coverage`" + ) + strategies = load_catalog() + if args.list: + _print_catalog(strategies) + return 0 + + strategy_ids = set(args.strategy) + sdk_functions = set(args.sdk_functions) + if args.interactive: + picked_strategies, picked_functions = _interactive_filters(strategies) + strategy_ids = strategy_ids or picked_strategies + sdk_functions = sdk_functions or picked_functions + + try: + cases = _select(strategies, strategy_ids, sdk_functions) + except ValueError as exc: + _parser().error(str(exc)) + selected_strategy_ids = {case.strategy_id for case in cases} + visible_strategies = tuple( + strategy for strategy in strategies if strategy.id in selected_strategy_ids + ) + dashboard = make_dashboard( + visible_strategies, + plain=args.plain, + confidence_strategies=strategies, + ) + pytest_args = [*args.pytest_arg] + if args.coverage: + pytest_args.extend(_coverage_pytest_args()) + with dashboard: + exit_code, run = run_pytest( + cases=cases, + repo_root=REPO_ROOT, + on_update=dashboard.update, + pytest_args=pytest_args, + ) + dashboard.finish(run, exit_code) + if args.coverage and (COVERAGE_ROOT / "python.json").exists(): + print(f"Python LOC heatmap: {COVERAGE_ROOT / 'python-html' / 'index.html'}") + print(f"Machine-readable coverage: {COVERAGE_ROOT / 'python.json'}") + return exit_code diff --git a/tests/rust-python-harness/e2e_fuzz_tests/README.md b/tests/rust-python-harness/e2e_fuzz_tests/README.md new file mode 100644 index 00000000000..34b12050ff9 --- /dev/null +++ b/tests/rust-python-harness/e2e_fuzz_tests/README.md @@ -0,0 +1,3 @@ +# End-to-end fuzz tests + +Runs the same SDK call through the Python and Rust paths using generated inputs and recorded provider responses. It compares public results, streams, callbacks, and exceptions to catch behavior differences a unit test can miss. diff --git a/tests/rust-python-harness/e2e_fuzz_tests/strategy.json b/tests/rust-python-harness/e2e_fuzz_tests/strategy.json new file mode 100644 index 00000000000..abeea01d9b5 --- /dev/null +++ b/tests/rust-python-harness/e2e_fuzz_tests/strategy.json @@ -0,0 +1,12 @@ +{ + "order": 10, + "id": "e2e_fuzz_tests", + "label": "End-to-end fuzz tests", + "description": "Compare observable Python and Rust SDK behavior over generated and recorded inputs.", + "functions": { + "ocr": {"coverage": "partial", "selectors": ["tests/test_litellm/ocr/test_rust_bridge.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."}, + "messages": {"coverage": "partial", "selectors": ["tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py"], "note": "Bridge coverage exists; frozen-oracle fuzz parity is still being added."}, + "responses": {"coverage": "partial", "selectors": ["tests/test_litellm/responses/test_rust_bridge_websocket.py"], "note": "Covers the websocket bridge; full responses parity is still being added."}, + "count_tokens": {"coverage": "planned", "selectors": [], "note": "No Rust count_tokens parity test is present yet."} + } +} diff --git a/tests/rust-python-harness/models.py b/tests/rust-python-harness/models.py new file mode 100644 index 00000000000..21097e0f7d0 --- /dev/null +++ b/tests/rust-python-harness/models.py @@ -0,0 +1,219 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from time import monotonic +from typing import Iterable + + +class Coverage(str, Enum): + COMPLETE = "complete" + PARTIAL = "partial" + PLANNED = "planned" + NOT_APPLICABLE = "not_applicable" + + +class RunStatus(str, Enum): + NOT_RUN = "not_run" + QUEUED = "queued" + RUNNING = "running" + PASSED = "passed" + FAILED = "failed" + SKIPPED = "skipped" + ERROR = "error" + MISSING = "missing" + PLANNED = "planned" + NOT_APPLICABLE = "not_applicable" + + +class ConfidenceLevel(str, Enum): + HIGH = "HIGH" + MEDIUM = "MEDIUM" + LOW = "LOW" + + +SDK_FUNCTIONS = ("ocr", "messages", "responses", "count_tokens") + + +@dataclass(frozen=True) +class HarnessCase: + strategy_id: str + strategy_label: str + sdk_function: str + coverage: Coverage + selectors: tuple[str, ...] + note: str = "" + + @property + def key(self) -> str: + return f"{self.strategy_id}:{self.sdk_function}" + + +@dataclass(frozen=True) +class Strategy: + order: int + id: str + label: str + description: str + directory: Path + cases: tuple[HarnessCase, ...] + + +@dataclass +class CaseResult: + case: HarnessCase + status: RunStatus = RunStatus.NOT_RUN + collected: set[str] = field(default_factory=set) + completed: set[str] = field(default_factory=set) + passed: int = 0 + failed: int = 0 + skipped: int = 0 + errors: int = 0 + outcomes: dict[str, RunStatus] = field(default_factory=dict) + durations: dict[str, float] = field(default_factory=dict) + + @property + def total(self) -> int: + return len(self.collected) + + @property + def duration(self) -> float: + return sum(self.durations.values()) + + def record(self, nodeid: str, status: RunStatus, duration: float = 0.0) -> None: + """Record a terminal outcome, allowing teardown errors to replace a pass.""" + self.outcomes[nodeid] = status + self.durations[nodeid] = self.durations.get(nodeid, 0.0) + duration + self.completed = set(self.outcomes) + values = tuple(self.outcomes.values()) + self.passed = values.count(RunStatus.PASSED) + self.failed = values.count(RunStatus.FAILED) + self.skipped = values.count(RunStatus.SKIPPED) + self.errors = values.count(RunStatus.ERROR) + self.finalize() + + def set_initial_status(self) -> None: + if self.case.coverage is Coverage.NOT_APPLICABLE: + self.status = RunStatus.NOT_APPLICABLE + elif not self.case.selectors: + self.status = RunStatus.PLANNED + else: + self.status = RunStatus.QUEUED + + def finalize(self) -> None: + if self.status in {RunStatus.NOT_APPLICABLE, RunStatus.PLANNED}: + return + if not self.collected: + self.status = RunStatus.MISSING + elif self.errors: + self.status = RunStatus.ERROR + elif self.failed: + self.status = RunStatus.FAILED + elif self.passed and len(self.completed) == len(self.collected): + self.status = RunStatus.PASSED + elif self.skipped and len(self.completed) == len(self.collected): + self.status = RunStatus.SKIPPED + + +@dataclass +class HarnessRun: + results: dict[str, CaseResult] + current_nodeid: str | None = None + failures: list[tuple[str, str]] = field(default_factory=list) + started_at: float = field(default_factory=monotonic) + finished_at: float | None = None + + @property + def duration(self) -> float: + return (self.finished_at or monotonic()) - self.started_at + + @property + def unique_tests(self) -> int: + return len( + {nodeid for result in self.results.values() for nodeid in result.collected} + ) + + @property + def completed_tests(self) -> int: + return len( + {nodeid for result in self.results.values() for nodeid in result.completed} + ) + + @classmethod + def from_cases(cls, cases: Iterable[HarnessCase]) -> "HarnessRun": + results = {case.key: CaseResult(case=case) for case in cases} + for result in results.values(): + result.set_initial_status() + return cls(results=results) + + +@dataclass(frozen=True) +class SectionConfidence: + sdk_function: str + verified_strategies: int + required_strategies: int + level: ConfidenceLevel + details: tuple[str, ...] + + @property + def percentage(self) -> int: + if not self.required_strategies: + return 0 + return round(100 * self.verified_strategies / self.required_strategies) + + +def section_confidence( + run: HarnessRun, strategies: Iterable[Strategy] +) -> tuple[SectionConfidence, ...]: + strategy_list = tuple(strategies) + scores: list[SectionConfidence] = [] + for sdk_function in SDK_FUNCTIONS: + cases = tuple( + case + for strategy in strategy_list + for case in strategy.cases + if case.sdk_function == sdk_function + and case.coverage is not Coverage.NOT_APPLICABLE + ) + verified = 0 + details: list[str] = [] + for case in cases: + result = run.results.get(case.key) + status = result.status if result is not None else RunStatus.NOT_RUN + if status is RunStatus.PASSED: + verified += 1 + details.append( + f"{STATUS_LABELS[status]} {case.strategy_id} ({case.coverage.value})" + ) + required = len(cases) + if required and verified == required: + level = ConfidenceLevel.HIGH + elif verified: + level = ConfidenceLevel.MEDIUM + else: + level = ConfidenceLevel.LOW + scores.append( + SectionConfidence( + sdk_function=sdk_function, + verified_strategies=verified, + required_strategies=required, + level=level, + details=tuple(details), + ) + ) + return tuple(scores) + + +STATUS_LABELS = { + RunStatus.NOT_RUN: "·", + RunStatus.QUEUED: "○", + RunStatus.RUNNING: "◉", + RunStatus.PASSED: "✓", + RunStatus.FAILED: "✗", + RunStatus.SKIPPED: "↷", + RunStatus.ERROR: "!", + RunStatus.MISSING: "?", + RunStatus.PLANNED: "—", + RunStatus.NOT_APPLICABLE: "n/a", +} diff --git a/tests/rust-python-harness/runner.py b/tests/rust-python-harness/runner.py new file mode 100644 index 00000000000..82393ef234e --- /dev/null +++ b/tests/rust-python-harness/runner.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import os +from collections.abc import Callable, Sequence +from pathlib import Path +from time import monotonic + +import pytest + +from .models import CaseResult, HarnessCase, HarnessRun, RunStatus + +UpdateCallback = Callable[[HarnessRun], None] + + +def selector_matches_node(selector: str, nodeid: str) -> bool: + normalized_selector = selector.replace("\\", "/") + normalized_nodeid = nodeid.replace("\\", "/") + if "::" in normalized_selector: + return normalized_nodeid == normalized_selector or normalized_nodeid.startswith( + f"{normalized_selector}[" + ) + return normalized_nodeid == normalized_selector or normalized_nodeid.startswith( + f"{normalized_selector}::" + ) + + +def selector_path(selector: str) -> Path: + return Path(selector.split("::", 1)[0]) + + +def runnable_selectors( + cases: Sequence[HarnessCase], repo_root: Path +) -> tuple[str, ...]: + selectors = { + selector + for case in cases + for selector in case.selectors + if (repo_root / selector_path(selector)).exists() + } + return tuple(sorted(selectors)) + + +class HarnessPytestPlugin: + def __init__(self, run: HarnessRun, on_update: UpdateCallback) -> None: + self.run = run + self.on_update = on_update + self.node_to_results: dict[str, list[CaseResult]] = {} + + def _notify(self) -> None: + self.on_update(self.run) + + def pytest_collection_modifyitems(self, items: list[pytest.Item]) -> None: + for item in items: + matched_results: list[CaseResult] = [] + for result in self.run.results.values(): + if any( + selector_matches_node(selector, item.nodeid) + for selector in result.case.selectors + ): + result.collected.add(item.nodeid) + matched_results.append(result) + if matched_results: + self.node_to_results[item.nodeid] = matched_results + for result in self.run.results.values(): + if result.status is RunStatus.QUEUED and not result.collected: + result.status = RunStatus.MISSING + self._notify() + + def pytest_runtest_logstart( + self, nodeid: str, location: tuple[str, int | None, str] + ) -> None: + del location + self.run.current_nodeid = nodeid + for result in self.node_to_results.get(nodeid, []): + if result.status not in {RunStatus.FAILED, RunStatus.ERROR}: + result.status = RunStatus.RUNNING + self._notify() + + def pytest_runtest_logreport(self, report: pytest.TestReport) -> None: + if report.when not in {"setup", "call", "teardown"}: + return + results = self.node_to_results.get(report.nodeid, []) + if not results: + return + + terminal = report.when == "call" or report.failed or report.skipped + if not terminal: + for result in results: + result.durations[report.nodeid] = ( + result.durations.get(report.nodeid, 0.0) + report.duration + ) + return + for result in results: + if report.when == "teardown" and not report.failed: + result.durations[report.nodeid] = ( + result.durations.get(report.nodeid, 0.0) + report.duration + ) + continue + if report.skipped: + status = RunStatus.SKIPPED + elif report.failed and report.when in {"setup", "teardown"}: + status = RunStatus.ERROR + elif report.failed: + status = RunStatus.FAILED + else: + status = RunStatus.PASSED + result.record(report.nodeid, status, report.duration) + if report.failed: + failure = (report.nodeid, str(report.longrepr)) + if failure not in self.run.failures: + self.run.failures.append(failure) + self._notify() + + def pytest_sessionfinish( + self, session: pytest.Session, exitstatus: int | pytest.ExitCode + ) -> None: + del session, exitstatus + self.run.current_nodeid = None + self.run.finished_at = monotonic() + for result in self.run.results.values(): + result.finalize() + self._notify() + + +def run_pytest( + cases: Sequence[HarnessCase], + repo_root: Path, + on_update: UpdateCallback, + pytest_args: Sequence[str] = (), +) -> tuple[int, HarnessRun]: + run = HarnessRun.from_cases(cases) + selectors = runnable_selectors(cases, repo_root) + if not selectors: + for result in run.results.values(): + result.finalize() + run.finished_at = monotonic() + on_update(run) + has_missing_test = any( + result.status is RunStatus.MISSING for result in run.results.values() + ) + exit_code = ( + int(pytest.ExitCode.TESTS_FAILED) + if has_missing_test + else int(pytest.ExitCode.OK) + ) + return exit_code, run + + plugin = HarnessPytestPlugin(run=run, on_update=on_update) + args = [*selectors, "-p", "no:terminal", *pytest_args] + previous_directory = Path.cwd() + try: + os.chdir(repo_root) + exit_code = int(pytest.main(args, plugins=[plugin])) + finally: + os.chdir(previous_directory) + if exit_code == 0 and any( + result.status is RunStatus.MISSING for result in run.results.values() + ): + exit_code = int(pytest.ExitCode.TESTS_FAILED) + return exit_code, run diff --git a/tests/rust-python-harness/ui.py b/tests/rust-python-harness/ui.py new file mode 100644 index 00000000000..57fedd17fa6 --- /dev/null +++ b/tests/rust-python-harness/ui.py @@ -0,0 +1,295 @@ +from __future__ import annotations + +import os +import shlex +import sys +from collections.abc import Sequence +from contextlib import AbstractContextManager +from pathlib import Path +from typing import Any + +from .models import ( + Coverage, + HarnessRun, + RunStatus, + SDK_FUNCTIONS, + Strategy, + section_confidence, +) + +STATUS_GLYPHS = { + RunStatus.NOT_RUN: "·", + RunStatus.QUEUED: "○", + RunStatus.RUNNING: "◉", + RunStatus.PASSED: "✓", + RunStatus.FAILED: "✗", + RunStatus.SKIPPED: "↷", + RunStatus.ERROR: "!", + RunStatus.MISSING: "?", + RunStatus.PLANNED: "—", + RunStatus.NOT_APPLICABLE: "n/a", +} + +STATUS_STYLES = { + RunStatus.QUEUED: "dim", + RunStatus.RUNNING: "bold cyan", + RunStatus.PASSED: "bold green", + RunStatus.FAILED: "bold red", + RunStatus.SKIPPED: "yellow", + RunStatus.ERROR: "bold red", + RunStatus.MISSING: "magenta", + RunStatus.PLANNED: "dim", + RunStatus.NOT_APPLICABLE: "dim", +} + + +def _format_duration(seconds: float) -> str: + if seconds < 1: + return f"{seconds * 1000:.0f}ms" + if seconds < 60: + return f"{seconds:.1f}s" + return f"{int(seconds // 60)}m {seconds % 60:.0f}s" + + +def _rerun_command(nodeid: str) -> str: + return f"poetry run pytest {shlex.quote(nodeid)} -q" + + +def _summary(run: HarnessRun) -> tuple[int, int, int, int]: + outcomes: dict[str, RunStatus] = {} + for result in run.results.values(): + outcomes.update(result.outcomes) + return ( + list(outcomes.values()).count(RunStatus.PASSED), + list(outcomes.values()).count(RunStatus.FAILED), + list(outcomes.values()).count(RunStatus.ERROR), + list(outcomes.values()).count(RunStatus.SKIPPED), + ) + + +def _cell_text(run: HarnessRun, strategy_id: str, sdk_function: str) -> tuple[str, str]: + result = run.results.get(f"{strategy_id}:{sdk_function}") + if result is None: + return "", "" + counts = "" + if result.total: + counts = f" {len(result.completed)}/{result.total}" + coverage = " ◐" if result.case.coverage is Coverage.PARTIAL else "" + return f"{STATUS_GLYPHS[result.status]}{counts}{coverage}", STATUS_STYLES.get( + result.status, "" + ) + + +class RichDashboard(AbstractContextManager["RichDashboard"]): + def __init__( + self, + strategies: Sequence[Strategy], + confidence_strategies: Sequence[Strategy], + ) -> None: + from rich.console import Console + from rich.live import Live + + self.strategies = strategies + self.confidence_strategies = confidence_strategies + self.console = Console() + self.live: Any = Live( + console=self.console, refresh_per_second=12, transient=False + ) + + def _table(self, run: HarnessRun) -> Any: + from rich import box + from rich.table import Table + from rich.text import Text + + narrow = self.console.width < 96 + if narrow: + table = Table(box=box.SIMPLE_HEAVY, expand=True, show_header=False) + table.add_column("Strategy", ratio=3) + table.add_column("Results", ratio=5) + for strategy in self.strategies: + values = [] + for sdk_function in SDK_FUNCTIONS: + value, style = _cell_text(run, strategy.id, sdk_function) + if value: + values.append( + Text.assemble((f"{sdk_function} ", "dim"), (value, style)) + ) + table.add_row(strategy.label, Text(" ").join(values)) + return table + + table = Table(box=box.ROUNDED, expand=True, title="Strategy × SDK function") + table.add_column("Strategy", ratio=3) + for label in ("ocr/aocr", "messages", "responses", "count_tokens"): + table.add_column(label, justify="center", ratio=1) + for strategy in self.strategies: + cells = [] + for sdk_function in SDK_FUNCTIONS: + value, style = _cell_text(run, strategy.id, sdk_function) + cells.append(Text(value, style=style)) + table.add_row(strategy.label, *cells) + return table + + def __enter__(self) -> "RichDashboard": + self.live.__enter__() + return self + + def __exit__(self, *args: object) -> None: + self.live.__exit__(*args) + + def update(self, run: HarnessRun) -> None: + from rich.markup import escape + from rich.panel import Panel + + active = run.current_nodeid or "Waiting for test events…" + if len(active) > max(40, self.console.width - 16): + active = f"…{active[-(self.console.width - 17):]}" + passed, failed, errors, skipped = _summary(run) + progress = ( + f"[bold]{run.completed_tests}/{run.unique_tests}[/bold] tests " + f"[green]{passed} passed[/green] [red]{failed + errors} failed[/red] " + f"[yellow]{skipped} skipped[/yellow] [dim]{_format_duration(run.duration)}[/dim]" + ) + legend = "✓ pass ✗ fail ! error ↷ skip\n? configured test missing — planned ◐ partial coverage" + self.live.update( + Panel( + self._table(run), + title="⚡ Rust ↔ Python parity lab", + subtitle=f"{progress}\n[dim]{escape(active)}[/dim]\n{legend}", + border_style="cyan", + ) + ) + + def finish(self, run: HarnessRun, exit_code: int) -> None: + self.update(run) + if run.failures: + from rich.markup import escape + from rich.panel import Panel + + for nodeid, detail in run.failures[:5]: + rerun = _rerun_command(nodeid) + self.console.print( + Panel( + f"{escape(detail)}\n\n[bold]Rerun just this test[/bold]\n" + f"[cyan]{escape(rerun)}[/cyan]", + title=f"✗ {escape(nodeid)}", + border_style="red", + ) + ) + durations: dict[str, float] = {} + for result in run.results.values(): + for nodeid, duration in result.durations.items(): + durations[nodeid] = max(duration, durations.get(nodeid, 0.0)) + if durations: + slow = sorted(durations.items(), key=lambda item: item[1], reverse=True)[:3] + self.console.print( + "[bold]Slowest tests[/bold] " + + " • ".join( + f"{Path(nodeid).name} [dim]{_format_duration(duration)}[/dim]" + for nodeid, duration in slow + ) + ) + from rich import box + from rich.table import Table + + confidence_table = Table( + title="Port confidence by SDK section", box=box.ROUNDED, expand=True + ) + confidence_table.add_column("SDK section") + confidence_table.add_column("Score", justify="right") + confidence_table.add_column("Confidence") + confidence_table.add_column("Strategy evidence", ratio=4) + confidence_styles = {"HIGH": "green", "MEDIUM": "yellow", "LOW": "red"} + for score in section_confidence(run, self.confidence_strategies): + confidence_table.add_row( + score.sdk_function, + f"{score.verified_strategies}/{score.required_strategies} {score.percentage}%", + f"[{confidence_styles[score.level.value]}]{score.level.value}[/]", + " ".join(score.details), + ) + self.console.print(confidence_table) + self.console.print( + "[dim]Score = required strategies with passing evidence. " + "LOC coverage remains a separate report.[/dim]" + ) + style = "green" if exit_code == 0 else "red" + self.console.print( + f"[{style}]Harness finished in {_format_duration(run.duration)} " + f"(exit {exit_code})[/{style}]" + ) + + +class PlainDashboard(AbstractContextManager["PlainDashboard"]): + def __init__( + self, + strategies: Sequence[Strategy], + confidence_strategies: Sequence[Strategy], + ) -> None: + self.strategies = strategies + self.confidence_strategies = confidence_strategies + self._seen: dict[str, tuple[RunStatus, int]] = {} + + def __enter__(self) -> "PlainDashboard": + print("Rust <-> Python SDK parity harness", flush=True) + return self + + def __exit__(self, *args: object) -> None: + return None + + def update(self, run: HarnessRun) -> None: + for key, result in run.results.items(): + state = (result.status, len(result.completed)) + if self._seen.get(key) != state: + self._seen[key] = state + progress = ( + f" {len(result.completed)}/{result.total}" if result.total else "" + ) + print( + f"{STATUS_GLYPHS[result.status]} {key}: {result.status.value}{progress}", + flush=True, + ) + + def finish(self, run: HarnessRun, exit_code: int) -> None: + self.update(run) + passed, failed, errors, skipped = _summary(run) + print( + f"Summary: {passed} passed, {failed} failed, {errors} errors, " + f"{skipped} skipped in {_format_duration(run.duration)}", + flush=True, + ) + for nodeid, _ in run.failures[:5]: + print(f"Rerun: {_rerun_command(nodeid)}", flush=True) + print("Port confidence by SDK section", flush=True) + for score in section_confidence(run, self.confidence_strategies): + print( + f" {score.sdk_function:12} " + f"{score.verified_strategies}/{score.required_strategies} " + f"{score.percentage:3}% {score.level.value:6} " + f"{' | '.join(score.details)}", + flush=True, + ) + print( + " Score = required strategies with passing evidence; LOC is reported separately.", + flush=True, + ) + print(f"Harness finished with exit code {exit_code}", flush=True) + + +def make_dashboard( + strategies: Sequence[Strategy], + plain: bool = False, + confidence_strategies: Sequence[Strategy] | None = None, +) -> RichDashboard | PlainDashboard: + confidence_strategies = confidence_strategies or strategies + interactive_terminal = ( + sys.stdout.isatty() + and not os.environ.get("CI") + and os.environ.get("TERM") != "dumb" + ) + if not plain and interactive_terminal: + try: + import rich # noqa: F401 + + return RichDashboard(strategies, confidence_strategies) + except ImportError: + pass + return PlainDashboard(strategies, confidence_strategies) diff --git a/tests/rust-python-harness/unit_tests_rust/README.md b/tests/rust-python-harness/unit_tests_rust/README.md new file mode 100644 index 00000000000..12c7eb0089c --- /dev/null +++ b/tests/rust-python-harness/unit_tests_rust/README.md @@ -0,0 +1,3 @@ +# Rust unit tests + +Holds focused Cargo tests for Rust-owned parsing, transforms, errors, and streaming behavior. These tests make failures fast to diagnose before the Python bridge or full SDK path is involved. diff --git a/tests/rust-python-harness/unit_tests_rust/strategy.json b/tests/rust-python-harness/unit_tests_rust/strategy.json new file mode 100644 index 00000000000..89e897c872d --- /dev/null +++ b/tests/rust-python-harness/unit_tests_rust/strategy.json @@ -0,0 +1,12 @@ +{ + "order": 20, + "id": "unit_tests_rust", + "label": "Rust unit tests", + "description": "Exercise Rust-owned behavior directly with focused unit tests.", + "functions": { + "ocr": {"coverage": "planned", "selectors": []}, + "messages": {"coverage": "planned", "selectors": []}, + "responses": {"coverage": "planned", "selectors": []}, + "count_tokens": {"coverage": "planned", "selectors": []} + } +} diff --git a/tests/rust-python-harness/validate_sub_methods/README.md b/tests/rust-python-harness/validate_sub_methods/README.md new file mode 100644 index 00000000000..24894366f23 --- /dev/null +++ b/tests/rust-python-harness/validate_sub_methods/README.md @@ -0,0 +1,3 @@ +# Validate sub-methods + +Checks each request, response, stream, and error-mapping sub-method independently across Python and Rust. It also validates that traced Python helpers have an explicit Rust implementation and parity test. diff --git a/tests/rust-python-harness/validate_sub_methods/strategy.json b/tests/rust-python-harness/validate_sub_methods/strategy.json new file mode 100644 index 00000000000..6e6381678e0 --- /dev/null +++ b/tests/rust-python-harness/validate_sub_methods/strategy.json @@ -0,0 +1,12 @@ +{ + "order": 30, + "id": "validate_sub_methods", + "label": "Validate sub-methods", + "description": "Compare isolated transforms and verify Python-to-Rust helper coverage.", + "functions": { + "ocr": {"coverage": "planned", "selectors": []}, + "messages": {"coverage": "planned", "selectors": []}, + "responses": {"coverage": "planned", "selectors": []}, + "count_tokens": {"coverage": "planned", "selectors": []} + } +} diff --git a/tests/test_litellm/integrations/otel/test_langfuse_logger.py b/tests/test_litellm/integrations/otel/test_langfuse_logger.py new file mode 100644 index 00000000000..8f93a9a564f --- /dev/null +++ b/tests/test_litellm/integrations/otel/test_langfuse_logger.py @@ -0,0 +1,359 @@ +"""Tests for ``LangfuseOpenTelemetryV2``: the root observation's input and output are stamped from the +request-task hooks, while the root span is still recording, so Langfuse can show them on the trace.""" + +import asyncio +import json +from collections.abc import AsyncIterator, Sequence +from typing import Final + +import pytest + +pytest.importorskip("opentelemetry") + +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter # noqa: E402 + +import litellm # noqa: E402 +from litellm.caching.dual_cache import DualCache # noqa: E402 +from litellm.integrations.otel.logger import build_otel_v2_logger # noqa: E402 +from litellm.integrations.otel.model.config import OpenTelemetryV2Config, is_otel_v2_enabled # noqa: E402 +from litellm.integrations.otel.model.spans import LITELLM_PROXY_REQUEST_SPAN_NAME, SpanRole # noqa: E402 +from litellm.integrations.otel.plumbing import context as otel_context # noqa: E402 +from litellm.integrations.otel.plumbing import providers # noqa: E402 +from litellm.integrations.otel.plumbing.context import set_request_root_span # noqa: E402 +from litellm.litellm_core_utils.litellm_logging import _maybe_construct_otel_v2 # noqa: E402 +from litellm.proxy._types import UserAPIKeyAuth # noqa: E402 +from litellm.proxy.utils import ProxyLogging # noqa: E402 +from litellm.types.llms.openai import ( # noqa: E402 + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, +) +from litellm.types.utils import ( # noqa: E402 + Choices, + Delta, + Embedding, + EmbeddingResponse, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) + +INPUT_ATTR: Final = "langfuse.observation.input" +OUTPUT_ATTR: Final = "langfuse.observation.output" +CHAT_DATA: Final = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "ping"}]} + + +@pytest.fixture(autouse=True) +def _reset_request_root_span(): + otel_context._request_root_span.set(None) + yield + otel_context._request_root_span.set(None) + + +def _logger(*, capture: str = "span_only", mappers: Sequence[str] = ("genai", "langfuse")): + cfg = OpenTelemetryV2Config(exporter="in_memory", mapper_names=list(mappers), capture_message_content=capture) + exporter = InMemorySpanExporter() + tracer_provider = providers.build_tracer_provider(cfg, exporter=exporter) + return build_otel_v2_logger(config=cfg, tracer_provider=tracer_provider), exporter + + +def _start_root(logger): + root = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + set_request_root_span(root) + return root + + +def _root_attrs(exporter): + by_name = {span.name: span for span in exporter.get_finished_spans()} + return dict(by_name[LITELLM_PROXY_REQUEST_SPAN_NAME].attributes or {}) + + +def _run_request(logger, data: dict, call_type: str, response: object): + root = _start_root(logger) + asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, call_type)) + asyncio.run(logger.async_post_call_success_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), response=response)) + root.end() + + +async def _relay(logger, chunks: Sequence[object], data: dict) -> list[object]: + async def source() -> AsyncIterator[object]: + for chunk in chunks: + yield chunk + + return [chunk async for chunk in logger.async_post_call_streaming_iterator_hook(UserAPIKeyAuth(), source(), data)] + + +def _run_stream(logger, data: dict, chunks: Sequence[object]) -> list[object]: + root = _start_root(logger) + asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, "acompletion")) + relayed = asyncio.run(_relay(logger, chunks, data)) + root.end() + return relayed + + +def _chat_chunk(content: str | None, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-1", + created=1, + model="gpt-5.4-mini", + choices=[StreamingChoices(index=0, delta=Delta(content=content), finish_reason=finish_reason)], + ) + + +def _responses_api_response() -> ResponsesAPIResponse: + return ResponsesAPIResponse( + id="resp_1", + created_at=1, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "pong", "annotations": []}], + } + ], + ) + + +def _anthropic_sse_frames() -> tuple[bytes, ...]: + events = ( + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "po"}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ng"}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}}, + {"type": "message_stop"}, + ) + return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events) + + +def test_chat_request_stamps_root_observation_input_and_output(): + logger, exporter = _logger() + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))]) + + _run_request(logger, CHAT_DATA, "acompletion", response) + + attrs = _root_attrs(exporter) + assert json.loads(attrs[INPUT_ATTR]) == [{"role": "user", "content": "ping"}] + output = json.loads(attrs[OUTPUT_ATTR]) + assert [(turn["role"], turn["content"]) for turn in output] == [("assistant", "pong")] + + +def test_responses_request_folds_instructions_into_input_and_stamps_output_items(): + logger, exporter = _logger() + data = {"model": "gpt-5.4-mini", "instructions": "be terse", "input": "ping"} + + _run_request(logger, data, "aresponses", _responses_api_response()) + + attrs = _root_attrs(exporter) + assert json.loads(attrs[INPUT_ATTR]) == [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "ping"}, + ] + output = json.loads(attrs[OUTPUT_ATTR]) + assert output[0]["role"] == "assistant" + assert output[0]["content"][0]["text"] == "pong" + + +def test_anthropic_messages_request_folds_system_into_input_and_stamps_content_blocks(): + logger, exporter = _logger() + data = {"model": "claude-sonnet-4-5", "system": "be terse", "messages": [{"role": "user", "content": "ping"}]} + response = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "pong"}]} + + _run_request(logger, data, "aanthropic_messages", response) + + attrs = _root_attrs(exporter) + assert json.loads(attrs[INPUT_ATTR]) == [ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "ping"}, + ] + assert json.loads(attrs[OUTPUT_ATTR]) == [{"role": "assistant", "content": [{"type": "text", "text": "pong"}]}] + + +def test_chat_stream_relays_chunks_untouched_and_stamps_assembled_output(): + logger, exporter = _logger() + chunks = (_chat_chunk("po"), _chat_chunk("ng"), _chat_chunk(None, finish_reason="stop")) + + relayed = _run_stream(logger, CHAT_DATA, chunks) + + assert [id(chunk) for chunk in relayed] == [id(chunk) for chunk in chunks] + output = json.loads(_root_attrs(exporter)[OUTPUT_ATTR]) + assert [(turn["role"], turn["content"]) for turn in output] == [("assistant", "pong")] + + +def test_responses_stream_stamps_output_from_the_completed_event(): + logger, exporter = _logger() + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_responses_api_response() + ) + chunks = ({"type": "response.created"}, {"type": "response.output_text.delta", "delta": "pong"}, completed) + + relayed = _run_stream(logger, {"model": "gpt-5.4-mini", "input": "ping"}, chunks) + + assert relayed == list(chunks) + output = json.loads(_root_attrs(exporter)[OUTPUT_ATTR]) + assert output[0]["content"][0]["text"] == "pong" + + +def test_anthropic_sse_stream_stamps_output_from_the_assembled_frames(): + logger, exporter = _logger() + frames = _anthropic_sse_frames() + + relayed = _run_stream( + logger, {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "ping"}]}, frames + ) + + assert relayed == list(frames) + output = json.loads(_root_attrs(exporter)[OUTPUT_ATTR]) + assert [(turn["role"], turn["content"]) for turn in output] == [("assistant", "pong")] + + +def test_root_observation_io_survives_the_root_ending_before_the_success_callback(): + logger, exporter = _logger() + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))]) + root = _start_root(logger) + asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), CHAT_DATA, "acompletion")) + logger.log_pre_api_call( + model="gpt-5.4-mini", + messages=[], + kwargs={"litellm_call_id": "call_1", "litellm_params": {"metadata": {}}}, + ) + asyncio.run( + logger.async_post_call_success_hook(data=CHAT_DATA, user_api_key_dict=UserAPIKeyAuth(), response=response) + ) + root.end() + + payload = { + "call_type": "acompletion", + "custom_llm_provider": "openai", + "model": "gpt-5.4-mini", + "messages": CHAT_DATA["messages"], + "response": response.model_dump(), + "status": "success", + "litellm_call_id": "call_1", + "metadata": {}, + "hidden_params": {}, + } + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": payload, "litellm_params": {"metadata": {}}}, response, None, None + ) + ) + + attrs = _root_attrs(exporter) + assert INPUT_ATTR in attrs and OUTPUT_ATTR in attrs + generation = next(span for span in exporter.get_finished_spans() if span.name != LITELLM_PROXY_REQUEST_SPAN_NAME) + assert OUTPUT_ATTR in dict(generation.attributes or {}) + + +def test_root_input_is_the_request_as_the_pre_call_chain_left_it(): + logger, exporter = _logger() + raw = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}]} + masked = {"model": "gpt-5.4-mini", "messages": [{"role": "user", "content": "my ssn is [REDACTED]"}]} + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="noted"))]) + root = _start_root(logger) + asyncio.run(logger.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), raw, "acompletion")) + asyncio.run(logger.async_post_call_success_hook(data=masked, user_api_key_dict=UserAPIKeyAuth(), response=response)) + root.end() + + assert json.loads(_root_attrs(exporter)[INPUT_ATTR]) == masked["messages"] + + +def test_root_already_ended_is_left_alone(): + logger, exporter = _logger() + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))]) + root = _start_root(logger) + root.end() + + asyncio.run( + logger.async_post_call_success_hook(data=CHAT_DATA, user_api_key_dict=UserAPIKeyAuth(), response=response) + ) + + attrs = _root_attrs(exporter) + assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs + + +def test_responses_without_a_message_body_stamp_neither_input_nor_output(): + logger, exporter = _logger() + embedding = EmbeddingResponse(model="e", data=[Embedding(embedding=[0.1], index=0, object="embedding")]) + + _run_request(logger, {"model": "e", "input": "ping"}, "aembedding", embedding) + + attrs = _root_attrs(exporter) + assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs + + +def test_unrenderable_output_never_raises_into_the_request(): + logger, exporter = _logger() + + _run_request(logger, CHAT_DATA, "acompletion", object()) + + attrs = _root_attrs(exporter) + assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs + + +@pytest.mark.parametrize( + ("capture", "mappers"), + [("no_content", ("genai", "langfuse")), ("span_only", ("genai",))], +) +def test_factory_keeps_the_base_logger_unless_langfuse_content_capture_is_on(capture, mappers): + logger, exporter = _logger(capture=capture, mappers=mappers) + + _run_request(logger, CHAT_DATA, "acompletion", ModelResponse()) + attrs = _root_attrs(exporter) + assert INPUT_ATTR not in attrs and OUTPUT_ATTR not in attrs + + +@pytest.mark.parametrize( + ("capture", "mappers", "relays_streams"), + [ + ("span_only", ("genai", "langfuse"), True), + ("no_content", ("genai", "langfuse"), False), + ("span_only", ("genai",), False), + ], +) +def test_only_langfuse_content_capture_takes_proxy_streams_off_the_fast_path( + monkeypatch, capture, mappers, relays_streams +): + logger, _ = _logger(capture=capture, mappers=mappers) + monkeypatch.setattr(litellm, "callbacks", [logger]) + + assert ProxyLogging._callback_capabilities().has_iterator_override is relays_streams + + +def test_langfuse_otel_preset_builds_a_logger_that_stamps_the_root(monkeypatch): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk") + monkeypatch.setenv("LANGFUSE_HOST", "https://cloud.langfuse.com") + monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "span_only") + is_otel_v2_enabled.cache_clear() + + loggers: list = [] + try: + built = _maybe_construct_otel_v2("langfuse_otel", loggers) + assert built is not None + assert _maybe_construct_otel_v2("langfuse_otel", loggers) is built + root = _start_root(built) + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="pong"))]) + asyncio.run( + built.async_post_call_success_hook(data=CHAT_DATA, user_api_key_dict=UserAPIKeyAuth(), response=response) + ) + attrs = dict(root.attributes or {}) + assert INPUT_ATTR in attrs and OUTPUT_ATTR in attrs + finally: + is_otel_v2_enabled.cache_clear() diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index e995cbae782..de8b654987b 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -2858,6 +2858,27 @@ class TestRecordGatewayInjection: AnthropicCacheControlHook.record_gateway_injection(kwargs, 0) assert kwargs["litellm_metadata"][self.KEY] == self.DEPLOYMENT + def test_an_every_deployment_mark_survives_a_later_per_deployment_stamp(self): + """A per-leg stamp like the Bedrock converse tool_config one describes one leg of + a payload every leg sends, so narrowing an every-deployment mark to that leg's + deployment would uncredit whichever leg gets billed after a failover.""" + kwargs: dict = {"litellm_metadata": {self.KEY: ""}, "model_info": {"id": self.DEPLOYMENT}} + AnthropicCacheControlHook.record_gateway_injection(kwargs, 1) + assert kwargs["litellm_metadata"][self.KEY] == "" + + def test_a_pre_choice_pass_stamps_the_sentinel_over_a_provisional_deployment(self): + """The router's prompt-management factory stamps a provisional deployment's + model_info into kwargs before the prompt pass runs, and any other deployment can + end up billed, so the pass declares every-deployment scope explicitly.""" + kwargs: dict = {"litellm_metadata": {}, "model_info": {"id": self.DEPLOYMENT}} + AnthropicCacheControlHook.record_gateway_injection(kwargs, 1, injected_for_every_deployment=True) + assert kwargs["litellm_metadata"][self.KEY] == "" + + def test_a_per_deployment_mark_still_follows_the_latest_leg(self): + kwargs: dict = {"litellm_metadata": {self.KEY: "dep-old"}, "model_info": {"id": self.DEPLOYMENT}} + AnthropicCacheControlHook.record_gateway_injection(kwargs, 1) + assert kwargs["litellm_metadata"][self.KEY] == self.DEPLOYMENT + def test_v1_messages_auto_injection_stamps_the_marker(self, monkeypatch): monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) kwargs: dict = {"litellm_metadata": {}, "model_info": {"id": self.DEPLOYMENT}} diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 366f61ded49..f1de7390b5b 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -6002,7 +6002,9 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o """The savings gate reads litellm_gateway_injected_cache from the request's metadata bucket. Recording lives in the shared prompt-hook wrappers, so chat, /v1/responses, router prompt deployments, and proxy prompt templates all mark - injected requests the same way; a hook that injects nothing leaves no marker.""" + injected requests the same way; a hook that injects nothing leaves no marker. + A pass that runs before deployment choice declares it and gets the every-deployment + sentinel, which a later per-deployment pass never narrows.""" from litellm.integrations.custom_prompt_management import CustomPromptManagement class _InjectingHook(CustomPromptManagement): @@ -6086,6 +6088,28 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o ) assert "litellm_gateway_injected_cache" not in untouched["metadata"] + pre_choice = {"metadata": {}, "model_info": {"id": "dep-of-this-attempt"}} + logging_obj.get_chat_completion_prompt( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + non_default_params={}, + prompt_variables=None, + prompt_management_logger=_InjectingHook(), + request_kwargs=pre_choice, + injected_for_every_deployment=True, + ) + assert pre_choice["metadata"]["litellm_gateway_injected_cache"] == "" + + await logging_obj.async_get_chat_completion_prompt( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "a fresh turn"}], + non_default_params={}, + prompt_variables=None, + prompt_management_logger=_InjectingHook(), + request_kwargs=pre_choice, + ) + assert pre_choice["metadata"]["litellm_gateway_injected_cache"] == "" + def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj): """LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead diff --git a/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py b/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py index b924ea8f93f..7e1139ca79a 100644 --- a/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py +++ b/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py @@ -184,6 +184,50 @@ async def test_download_file_drops_query_string_from_the_stored_url(mock_env_var ) +@pytest.mark.parametrize( + "malicious_filename", + [ + "report.jsonl/../../etc/cron.d/evil", + "a.b/../../../root/.ssh/authorized_keys", + ], +) +@pytest.mark.parametrize("strategy", ["uuid", "timestamp"]) +@pytest.mark.asyncio +async def test_generate_file_name_strips_path_traversal_from_extension(mock_env_vars, malicious_filename, strategy): + """ + original_filename.split(".")[-1] does not parse path structure, so a filename whose + last "." is followed by a directory traversal sequence used to put that sequence + straight into the blob path built from this name. The mutant this pins is reverting + _safe_extension() back to that bare split. + """ + backend = _make_backend() + generated = backend._generate_file_name(malicious_filename, strategy) + assert "/" not in generated + assert ".." not in generated + + +@pytest.mark.asyncio +async def test_generate_file_name_uuid_strategy_preserves_ordinary_extension(mock_env_vars): + backend = _make_backend() + generated = backend._generate_file_name("data.jsonl", "uuid") + assert generated.endswith(".jsonl") + + +@pytest.mark.asyncio +async def test_generate_file_name_original_filename_strategy_strips_directory_components(mock_env_vars): + """The blob name must never carry a directory the caller supplied, traversal or not.""" + backend = _make_backend() + generated = backend._generate_file_name("../../etc/passwd", "original_filename") + assert generated == "passwd" + + +@pytest.mark.asyncio +async def test_generate_file_name_null_byte_filename_falls_back_to_safe_default(mock_env_vars): + backend = _make_backend() + generated = backend._generate_file_name("report.pdf\x00.exe", "uuid") + assert "\x00" not in generated + + @pytest.mark.parametrize( "env_fixture, expected_suffix", [("mock_env_vars", "core.windows.net"), ("mock_gov_env_vars", GOV_SUFFIX)], diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 0144fbb17dd..c8ea4867f2c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -6005,6 +6005,44 @@ class TestMCPDcrBridgeDelegateAdmission: await MCPRequestHandler.process_mcp_request(scope) assert exc_info.value.status_code == 503 + assert exc_info.value.detail == ( + "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly." + ) + + async def test_user_subject_envelope_permanent_db_fault_is_503_not_worded_as_transient(self): + """A query engine fault that never heals (a missing engine binary) still fails admission with 503, + but the detail must not call the database "temporarily unreachable" or ask the client to retry: the + DCR client would loop on a retry that can never succeed. The fault reaches the handler wrapped in + get_user_object's bare ValueError, so the wording has to be picked off the wrapped cause.""" + from prisma.engine.errors import BinaryNotFoundError + + envelope = self._mint_bridge_envelope(user_id="sso-user-7") + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/bridge_delegate_server", + "headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))], + } + with ( + patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling admission tests + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + patch( # test-quality-ok: the envelope opener reads master_key off the proxy module, no injection seam + "litellm.proxy.proxy_server.master_key", self._MASTER_KEY + ), + self._patch_user_reload( + side_effect=self._wrapped_user_lookup_error(BinaryNotFoundError("query engine binary not found")) + ), + ): + mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server() + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 503 + assert "temporarily unreachable" not in exc_info.value.detail + assert "retry shortly" not in exc_info.value.detail.lower() + assert "BinaryNotFoundError" in exc_info.value.detail + assert "will not clear by retrying" in exc_info.value.detail async def test_user_subject_envelope_scim_deactivated_user_fails_closed_401(self): """SCIM-deactivating the envelope's user revokes it immediately: the reloaded user carries 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 598e9276423..588eba4adb7 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 @@ -6432,6 +6432,23 @@ async def test_bridge_mint_db_outage_is_503_before_upstream(): response, post = await _prepare_only_bridge_exchange("unavailable") assert response.status_code == 503 assert json.loads(response.body)["error"] == "temporarily_unavailable" + assert "retry shortly" in json.loads(response.body)["error_description"] + post.assert_not_called() + + +@pytest.mark.asyncio +async def test_bridge_mint_permanent_db_fault_is_503_without_retry_advice(): + """A query engine fault that never heals is still a 503 (the gateway is at fault, not the client), but + the description must not tell the client the database is temporarily unreachable and to retry: that + sends an operator to wait out an outage that is not one. The code stays temporarily_unavailable, the + only RFC 6749 error a client treats as a server-side 503.""" + response, post = await _prepare_only_bridge_exchange("faulted") + assert response.status_code == 503 + body = json.loads(response.body) + assert body["error"] == "temporarily_unavailable" + assert "temporarily unreachable" not in body["error_description"] + assert "retry shortly" not in body["error_description"] + assert "not a transient outage" in body["error_description"] post.assert_not_called() @@ -7143,6 +7160,56 @@ async def test_resolve_active_litellm_key_db_outage_is_unavailable(proxy_globals assert await _resolve_active_litellm_key(request) == "unavailable" +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_permanent_engine_fault_is_faulted(proxy_globals): + """A query engine that is missing or version-skewed cannot resolve any key until the deployment is + repaired, so the resolver reports "faulted" (still statused 503 by the mint) rather than "unavailable", + whose wording promises the outage is transient and asks the client to retry.""" + from prisma.engine.errors import BinaryNotFoundError + + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _FaultedPrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + raise BinaryNotFoundError("query engine binary not found") + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _FaultedPrisma() + + request = _token_request({"x-litellm-api-key": "sk-during-engine-fault"}) + assert await _resolve_active_litellm_key(request) == "faulted" + + +@pytest.mark.asyncio +async def test_resolve_active_litellm_key_transport_error_over_permanent_fault_is_faulted(proxy_globals): + """A reconnect that dies on a missing engine binary raises the transport error last, with the + BinaryNotFoundError as __context__. The binary is what blocks recovery, so the key read is "faulted", + not the "unavailable" that the outer ConnectError alone would suggest.""" + import httpx + from prisma.engine.errors import BinaryNotFoundError + + from litellm.proxy._experimental.mcp_server.bridge_token_flow import ( + _resolve_active_litellm_key, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _ReconnectFailedPrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + try: + raise BinaryNotFoundError("query engine binary not found") + except BinaryNotFoundError: + raise httpx.ConnectError("All connection attempts failed") + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _ReconnectFailedPrisma() + + request = _token_request({"x-litellm-api-key": "sk-during-failed-reconnect"}) + assert await _resolve_active_litellm_key(request) == "faulted" + + @pytest.mark.asyncio async def test_resolve_active_litellm_key_no_database_is_unresolvable(proxy_globals): """With no database connection configured the gateway cannot verify the presented key at all, so @@ -7214,6 +7281,26 @@ async def test_reload_active_user_by_id_db_outage_is_unavailable(proxy_globals): assert await _reload_active_user_by_id("sso-user-7") == "unavailable" +@pytest.mark.asyncio +async def test_reload_active_user_by_id_permanent_engine_fault_is_faulted(proxy_globals): + """A permanent query engine fault while re-validating the user on refresh is "faulted", not + "unavailable": both are 503s, but only the transient one may tell the client to retry. get_user_object + wraps the fault in a bare ValueError, so the classification has to read the wrapped cause.""" + from prisma.engine.errors import MismatchedVersionsError + + from litellm.proxy._experimental.mcp_server.bridge_token_flow import _reload_active_user_by_id + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + with patch( # test-quality-ok: get_user_object is the DB seam that wraps the fault; same patch as the outage sibling + "litellm.proxy.auth.auth_checks.get_user_object", + new=AsyncMock(side_effect=_wrapped_user_lookup_error(MismatchedVersionsError(expected="1", got="2"))), + ): + assert await _reload_active_user_by_id("sso-user-7") == "faulted" + + @pytest.mark.asyncio async def test_token_endpoint_uses_client_secret_basic_when_configured(): """LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 32a3f70c357..1670370f082 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -426,6 +426,7 @@ async def test_token_rejects_expired_code_and_missing_configuration(): [ ("no_active_key", 400, "invalid_grant"), ("unavailable", 503, "temporarily_unavailable"), + ("faulted", 503, "temporarily_unavailable"), ("unresolvable", 500, "server_error"), ], ) @@ -460,6 +461,25 @@ async def test_token_gates_on_live_user_revalidation(failure, expected_status, e assert json.loads(response.body)["error"] == expected_error +def test_permanent_db_fault_503_does_not_promise_a_retry_will_help(): + """Both DB failures are 503 temporarily_unavailable (the only OAuth error a client reads as a + server-side outage), so the description is the one place the two are told apart: a transient outage + says retry, a fault that never heals must say retrying will not help and point at the deployment.""" + from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + _consent_lookup_failure_response, + _mint_failure_response, + _reload_failure_response, + ) + + for render in (_reload_failure_response, _consent_lookup_failure_response, _mint_failure_response): + transient = json.loads(render("unavailable").body)["error_description"] + faulted = json.loads(render("faulted").body)["error_description"] + assert transient == "the gateway database is unavailable; retry" + assert "retry" not in faulted.replace("retrying will not help", "") + assert "not a transient outage" in faulted + assert "retrying will not help" in faulted + + @pytest.mark.asyncio async def test_flow_is_single_use_shared_cache_rejects_second_complete(): """A double-submit of the finish step mints only ONE code: the second complete over the @@ -1250,6 +1270,7 @@ async def test_native_authorize_refuses_a_hosted_redirect_for_the_proxy_api(): "failure, status, error", [ ("unavailable", 503, "temporarily_unavailable"), + ("faulted", 503, "temporarily_unavailable"), ("unresolvable", 500, "server_error"), ("no_active_key", 403, "access_denied"), ], @@ -1424,6 +1445,7 @@ async def test_native_code_without_a_minter_is_refused_server_side(): ("team_required", 400, "invalid_grant"), ("no_active_key", 400, "invalid_grant"), ("unavailable", 503, "temporarily_unavailable"), + ("faulted", 503, "temporarily_unavailable"), ("unresolvable", 500, "server_error"), ], ) @@ -1805,5 +1827,12 @@ async def test_introspect_fails_closed_on_dead_user_and_503s_on_outage(): status, body = await _introspect(minted.token.get_secret_value(), reload_user=_reload_user_outage) assert (status, body["error"]) == (503, "temporarily_unavailable") + async def _reload_user_faulted(user_id: str): + return "faulted" + + status, body = await _introspect(minted.token.get_secret_value(), reload_user=_reload_user_faulted) + assert (status, body["error"]) == (503, "temporarily_unavailable") + assert "not a transient outage" in body["error_description"] + status, body = await _introspect(minted.token.get_secret_value(), master_key=None) assert (status, body["error"]) == (500, "server_error") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 16221f44efe..239f89ebd90 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -10,33 +10,36 @@ Covers: """ import json +from collections.abc import Sequence from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest +from mcp.types import Tool +import litellm from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing from litellm.proxy._experimental.mcp_server.tool_search import ( AGENT_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, MCP_TOOL_SEARCH_TOOL_NAME, + SemanticToolRanker, + ToolSearchResult, coerce_top_k, get_virtual_tool_definitions, + search_mcp_tools, search_tools, ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.common_utils.semantic_text_index import EmbeddingFailed, SemanticTextIndex, Vector +from litellm.types.mcp import MCPToolSearchSettings -def _make_tools(specs: list[tuple[str, str]]) -> list[dict[str, Any]]: - return [ - { - "name": name, - "description": desc, - "inputSchema": {"type": "object", "properties": {}}, - } - for name, desc in specs - ] +def _make_tools(specs: list[tuple[str, str]]) -> tuple[Tool, ...]: + return tuple( + Tool(name=name, description=desc, inputSchema={"type": "object", "properties": {}}) for name, desc in specs + ) def _make_perm(**kwargs: Any) -> LiteLLM_ObjectPermissionTable: @@ -54,6 +57,160 @@ SAMPLE_TOOLS = _make_tools( ) +FX_TOOL = Tool( + name="treasury-get_rates", + description="Get foreign exchange rates for a currency pair", + inputSchema={"type": "object", "properties": {}}, +) +WEATHER_TOOL = Tool( + name="weather-forecast", + description="Get the weather forecast for a city", + inputSchema={"type": "object", "properties": {}}, +) +CALENDAR_TOOL = Tool( + name="calendar-create_event", + description="Create a calendar event", + inputSchema={"type": "object", "properties": {}}, +) +CATALOG = (FX_TOOL, WEATHER_TOOL, CALENDAR_TOOL) + +# A stand-in embedding space: "FX" sits next to the foreign-exchange tool and far from the rest. +FAKE_VECTORS: dict[str, Vector] = { + "FX": (1.0, 0.0), + f"{FX_TOOL.name}\n{FX_TOOL.description}": (0.9, 0.1), + f"{WEATHER_TOOL.name}\n{WEATHER_TOOL.description}": (0.3, 1.0), + f"{CALENDAR_TOOL.name}\n{CALENDAR_TOOL.description}": (0.0, 1.0), +} + + +class RecordingEmbedder: + def __init__(self) -> None: + self.calls: list[tuple[str, ...]] = [] + + async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]: + self.calls.append(tuple(texts)) + return tuple(FAKE_VECTORS[text] for text in texts) + + +def _ranker(embedder: RecordingEmbedder | None = None) -> SemanticToolRanker: + return SemanticToolRanker(embed=embedder or RecordingEmbedder(), embedding_model="emb", index=SemanticTextIndex()) + + +def _names(results: Sequence[ToolSearchResult] | EmbeddingFailed) -> list[str]: + assert not isinstance(results, EmbeddingFailed) + return [tool["name"] for tool in results] + + +class TestSearchMcpTools: + @pytest.mark.asyncio + async def test_semantic_mode_finds_foreign_exchange_tool_for_fx(self) -> None: + keyword_only = await search_mcp_tools("FX", CATALOG, 5, MCPToolSearchSettings(), ranker=None) + assert _names(keyword_only) == [] + + results = await search_mcp_tools("FX", CATALOG, 5, MCPToolSearchSettings(embedding_model="emb"), _ranker()) + assert _names(results) == [FX_TOOL.name, WEATHER_TOOL.name, CALENDAR_TOOL.name] + assert not isinstance(results, EmbeddingFailed) + assert results[0]["score"] > results[1]["score"] > results[2]["score"] + assert results[0]["inputSchema"] == FX_TOOL.inputSchema + + @pytest.mark.asyncio + async def test_similarity_threshold_drops_weak_matches(self) -> None: + settings = MCPToolSearchSettings(embedding_model="emb", similarity_threshold=0.5) + results = await search_mcp_tools("FX", CATALOG, 5, settings, _ranker()) + assert _names(results) == [FX_TOOL.name] + + @pytest.mark.asyncio + async def test_request_top_k_limits_semantic_results(self) -> None: + results = await search_mcp_tools("FX", CATALOG, 2, MCPToolSearchSettings(embedding_model="emb"), _ranker()) + assert _names(results) == [FX_TOOL.name, WEATHER_TOOL.name] + + @pytest.mark.asyncio + async def test_configured_top_k_caps_request_top_k(self) -> None: + settings = MCPToolSearchSettings(embedding_model="emb", top_k=1) + assert _names(await search_mcp_tools("FX", CATALOG, 50, settings, _ranker())) == [FX_TOOL.name] + assert _names(await search_mcp_tools("weather", CATALOG, 50, MCPToolSearchSettings(top_k=1), None)) == [ + WEATHER_TOOL.name + ] + + @pytest.mark.asyncio + async def test_core_tools_lead_and_do_not_consume_top_k(self) -> None: + settings = MCPToolSearchSettings(embedding_model="emb", top_k=1, core_tools=(CALENDAR_TOOL.name,)) + results = await search_mcp_tools("FX", CATALOG, 1, settings, _ranker()) + assert _names(results) == [CALENDAR_TOOL.name, FX_TOOL.name] + assert not isinstance(results, EmbeddingFailed) + assert "score" not in results[0] + + @pytest.mark.asyncio + async def test_core_tools_apply_in_keyword_mode_too(self) -> None: + settings = MCPToolSearchSettings(core_tools=(CALENDAR_TOOL.name,)) + assert _names(await search_mcp_tools("weather", CATALOG, 5, settings, None)) == [ + CALENDAR_TOOL.name, + WEATHER_TOOL.name, + ] + + @pytest.mark.asyncio + async def test_core_tools_outside_the_callers_catalog_are_not_returned(self) -> None: + settings = MCPToolSearchSettings(embedding_model="emb", core_tools=("payroll-run", CALENDAR_TOOL.name)) + results = await search_mcp_tools("FX", (FX_TOOL, WEATHER_TOOL), 5, settings, _ranker()) + assert _names(results) == [FX_TOOL.name, WEATHER_TOOL.name] + + @pytest.mark.asyncio + async def test_core_tools_are_listed_once_and_never_embedded(self) -> None: + embedder = RecordingEmbedder() + settings = MCPToolSearchSettings(embedding_model="emb", core_tools=(FX_TOOL.name, FX_TOOL.name)) + results = await search_mcp_tools("FX", CATALOG, 5, settings, _ranker(embedder)) + assert _names(results) == [FX_TOOL.name, WEATHER_TOOL.name, CALENDAR_TOOL.name] + assert all(FX_TOOL.description not in text for call in embedder.calls for text in call) + + @pytest.mark.asyncio + async def test_empty_query_returns_only_core_tools_without_embedding(self) -> None: + embedder = RecordingEmbedder() + settings = MCPToolSearchSettings(embedding_model="emb", core_tools=(CALENDAR_TOOL.name,)) + assert _names(await search_mcp_tools("", CATALOG, 5, settings, _ranker(embedder))) == [CALENDAR_TOOL.name] + assert embedder.calls == [] + + @pytest.mark.asyncio + async def test_repeat_searches_only_embed_the_query(self) -> None: + embedder = RecordingEmbedder() + ranker = _ranker(embedder) + settings = MCPToolSearchSettings(embedding_model="emb") + await search_mcp_tools("FX", CATALOG, 5, settings, ranker) + await search_mcp_tools("FX", CATALOG, 5, settings, ranker) + assert [len(call) for call in embedder.calls] == [4, 1] + + @pytest.mark.asyncio + async def test_embedding_failure_is_reported_not_raised(self) -> None: + async def failing(texts: Sequence[str]) -> Sequence[Vector]: + raise ValueError("embedding model is down") + + ranker = SemanticToolRanker(embed=failing, embedding_model="emb", index=SemanticTextIndex()) + result = await search_mcp_tools("FX", CATALOG, 5, MCPToolSearchSettings(embedding_model="emb"), ranker) + assert isinstance(result, EmbeddingFailed) + assert "embedding model is down" in result.reason + + +class TestMcpToolSearchSettings: + def test_rejects_out_of_range_values(self) -> None: + from pydantic import ValidationError + + with pytest.raises(ValidationError): + MCPToolSearchSettings(top_k=0) + with pytest.raises(ValidationError): + MCPToolSearchSettings(similarity_threshold=1.5) + + def test_yaml_shape_round_trips(self) -> None: + settings = MCPToolSearchSettings.model_validate( + {"embedding_model": "emb", "top_k": 3, "similarity_threshold": 0.2, "core_tools": ["a", "b"]} + ) + assert settings.core_tools == ("a", "b") + assert settings.model_dump() == { + "embedding_model": "emb", + "top_k": 3, + "similarity_threshold": 0.2, + "core_tools": ("a", "b"), + } + + class TestCoerceTopK: def test_int_passthrough(self) -> None: assert coerce_top_k(3) == 3 @@ -92,10 +249,10 @@ class TestSearchTools: assert len(results) <= 2 def test_empty_query_returns_empty(self) -> None: - assert search_tools("", SAMPLE_TOOLS) == [] + assert search_tools("", SAMPLE_TOOLS) == () def test_no_match_returns_empty(self) -> None: - assert search_tools("xyzzy_nonexistent_zzz", SAMPLE_TOOLS) == [] + assert search_tools("xyzzy_nonexistent_zzz", SAMPLE_TOOLS) == () def test_matches_description_not_just_name(self) -> None: results = search_tools("channel", SAMPLE_TOOLS) @@ -603,6 +760,63 @@ class TestCallToolRestApiVirtualTools: assert result.isError is True assert result.content[0].text == "set agent_search_embedding_model" + def _semantic_request(self, query: str = "FX") -> MagicMock: + return self._make_request({"name": MCP_TOOL_SEARCH_TOOL_NAME, "arguments": {"query": query}}) + + @pytest.mark.asyncio + async def test_mcp_tool_search_ranks_the_callers_catalog_with_the_configured_embedding_model( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "mcp_tool_search", {"embedding_model": "emb", "similarity_threshold": 0.5}) + user_api_key_dict = UserAPIKeyAuth( + api_key="k", team_id="team-1", object_permission=_make_perm(mcp_tool_search_enabled=True) + ) + + async def fake_aembedding(model: str, input: list[str], metadata: dict[str, Any]) -> MagicMock: + assert model == "emb" + assert metadata["user_api_key"] == "k" + assert metadata["user_api_key_team_id"] == "team-1" + response = MagicMock() + response.model_dump.return_value = {"data": [{"embedding": list(FAKE_VECTORS[t])} for t in input]} + return response + + router = MagicMock() + router.aembedding = AsyncMock(side_effect=fake_aembedding) + with ( + patch( # test-quality-ok: the proxy's router is a module global; the handler reaches it the way production does + "litellm.proxy.proxy_server.llm_router", router + ), + patch( # test-quality-ok: the authorized catalog is the seam every virtual tool shares; the ranking under test stays real + "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + new_callable=AsyncMock, + return_value=AggregateToolListing(tools=list(CATALOG), outcomes={}), + ) as mock_list, + ): + result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) + + assert mock_list.await_args.kwargs["user_api_key_auth"] is user_api_key_dict + assert result.isError is False + assert [t["name"] for t in json.loads(result.content[0].text)] == [FX_TOOL.name] + + @pytest.mark.asyncio + async def test_mcp_tool_search_reports_missing_router_as_tool_error(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "mcp_tool_search", {"embedding_model": "emb"}) + user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) + with patch( # test-quality-ok: the proxy's router is a module global; the handler reaches it the way production does + "litellm.proxy.proxy_server.llm_router", None + ): + result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) + assert result.isError is True + assert "mcp_tool_search.embedding_model" in result.content[0].text + + @pytest.mark.asyncio + async def test_mcp_tool_search_reports_invalid_settings_as_tool_error(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "mcp_tool_search", {"top_k": 0}) + user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True)) + result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) + assert result.isError is True + assert "top_k" in result.content[0].text + @pytest.mark.asyncio async def test_agent_search_requires_flag_enabled(self) -> None: from fastapi import HTTPException diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index 7ce62fdf648..231626c7eb5 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -8,7 +8,18 @@ from unittest.mock import AsyncMock, MagicMock import pytest -from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry, GrantMigrationResult +from litellm.constants import REDACTED_BY_LITELM_STRING +from litellm.proxy.agent_endpoints.agent_registry import ( + AgentRegistry, + GrantMigrationResult, + _restore_redacted_litellm_params, + redact_sensitive_agent_litellm_params, +) + +# Obviously-fake stand-ins for a real AWS credential pair (LIT-6736 regression +# fixtures) -- never a real key shape, and must never appear in any response. +SENTINEL_AWS_ACCESS_KEY_ID: Final = "AKIATESTSENTINEL0000" +SENTINEL_AWS_SECRET_ACCESS_KEY: Final = "test-sentinel-do-not-use-secret-value" def _sample_agent_card_params() -> dict: @@ -49,6 +60,7 @@ async def test_update_agent_in_db_clears_static_headers_and_extra_headers_when_o mock_update = AsyncMock(return_value=updated_agent) mock_prisma.db.litellm_agentstable.update = mock_update + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) # Agent config WITHOUT static_headers or extra_headers (omitted) agent_config = { @@ -95,6 +107,7 @@ async def test_update_agent_in_db_preserves_explicit_static_headers_and_extra_he mock_update = AsyncMock(return_value=updated_agent) mock_prisma.db.litellm_agentstable.update = mock_update + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) agent_config = { "agent_name": "Updated Agent", @@ -436,6 +449,9 @@ async def test_update_agent_in_db_raises_when_row_deleted_mid_update(): guard the code dereferences None and reports an opaque AttributeError instead of the id.""" registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace(litellm_params={}, object_permission_id=None) + ) mock_prisma.db.litellm_agentstable.update = AsyncMock(return_value=None) with pytest.raises(Exception, match="Error updating agent in DB") as exc_info: @@ -485,3 +501,492 @@ async def test_delete_agent_from_db_raises_when_row_already_gone(): await registry.delete_agent_from_db(agent_id="agent-123", prisma_client=mock_prisma) assert str(exc_info.value) == "Error deleting agent from DB: Agent not found, passed agent_id=agent-123" + + +# ---------- LIT-6736: agent litellm_params secret redaction ---------- + + +def test_redact_sensitive_agent_litellm_params_masks_secrets_keeps_the_rest(): + """The sentinel secret must never appear in the redacted output; non-secret + keys (model reference, is_public) must survive untouched.""" + redacted = redact_sensitive_agent_litellm_params( + { + "aws_access_key_id": SENTINEL_AWS_ACCESS_KEY_ID, + "aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, + "model": "bedrock/agentcore/my-agent", + "is_public": True, + } + ) + + assert SENTINEL_AWS_ACCESS_KEY_ID not in json.dumps(redacted) + assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted) + assert redacted["aws_access_key_id"] == REDACTED_BY_LITELM_STRING + assert redacted["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert redacted["model"] == "bedrock/agentcore/my-agent" + assert redacted["is_public"] is True + + +def test_redact_sensitive_agent_litellm_params_recurses_into_nested_dicts(): + """A secret nested one level down (e.g. a per-provider sub-config) must + also be redacted, not just top-level keys.""" + redacted = redact_sensitive_agent_litellm_params( + {"provider_config": {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-east-1"}} + ) + + assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted) + assert redacted["provider_config"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert redacted["provider_config"]["region"] == "us-east-1" + + +def test_redact_sensitive_agent_litellm_params_handles_none_and_json_string(): + assert redact_sensitive_agent_litellm_params(None) is None + + serialized = json.dumps({"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "model": "gpt-4"}) + redacted = redact_sensitive_agent_litellm_params(serialized) + + assert SENTINEL_AWS_SECRET_ACCESS_KEY not in redacted + assert json.loads(redacted)["api_key"] == REDACTED_BY_LITELM_STRING + assert json.loads(redacted)["model"] == "gpt-4" + + +def test_redact_sensitive_agent_litellm_params_recurses_into_lists_of_dicts(): + """A secret nested inside a list of provider sub-configs (a shape a + non-sensitively-named key can legitimately hold) must also be redacted, + not silently returned as-is.""" + redacted = redact_sensitive_agent_litellm_params( + { + "provider_configs": [ + {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-east-1"}, + {"aws_secret_access_key": "other-" + SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-west-2"}, + ] + } + ) + + assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted) + assert redacted["provider_configs"][0]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert redacted["provider_configs"][0]["region"] == "us-east-1" + assert redacted["provider_configs"][1]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert redacted["provider_configs"][1]["region"] == "us-west-2" + + +def test_redact_sensitive_agent_litellm_params_redacts_secrets_inside_model_list(): + """The exact shape flagged in review: litellm_params.model_list, where each + entry carries its own nested litellm_params with a provider credential.""" + redacted = redact_sensitive_agent_litellm_params( + { + "model_list": [ + { + "model_name": "gpt-4", + "litellm_params": {"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "model": "gpt-4"}, + }, + { + "model_name": "claude", + "litellm_params": { + "aws_secret_access_key": "other-" + SENTINEL_AWS_SECRET_ACCESS_KEY, + "model": "bedrock/claude", + }, + }, + ] + } + ) + + assert SENTINEL_AWS_SECRET_ACCESS_KEY not in json.dumps(redacted) + assert redacted["model_list"][0]["litellm_params"]["api_key"] == REDACTED_BY_LITELM_STRING + assert redacted["model_list"][0]["litellm_params"]["model"] == "gpt-4" + assert redacted["model_list"][1]["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert redacted["model_list"][1]["litellm_params"]["model"] == "bedrock/claude" + + +def test_restore_redacted_litellm_params_preserves_secret_inside_model_list(): + """The write-side counterpart: a caller editing a model_list entry's own + non-secret field (renaming it) while leaving that same entry's nested + secret masked must not corrupt the stored per-deployment credential. + List entries correspond by position (see the module docstring on + ``_restore_redacted_nested_value``), so this -- the common "edit this + entry, keep its secret" pattern -- must keep working.""" + existing = { + "agent_name": "my-agent", + "model_list": [ + { + "model_name": "gpt-4", + "litellm_params": {"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "model": "gpt-4"}, + }, + ], + } + incoming = { + "agent_name": "my-agent-renamed", + "model_list": [ + { + "model_name": "gpt-4-renamed", + "litellm_params": {"api_key": REDACTED_BY_LITELM_STRING, "model": "gpt-4"}, + }, + ], + } + + restored = _restore_redacted_litellm_params(incoming, existing) + + assert SENTINEL_AWS_SECRET_ACCESS_KEY == restored["model_list"][0]["litellm_params"]["api_key"] + assert restored["model_list"][0]["model_name"] == "gpt-4-renamed" + assert restored["agent_name"] == "my-agent-renamed" + + +def test_restore_redacted_litellm_params_matches_list_entries_by_position(): + """Documents the accepted trade-off: a list has no stable per-element + identity in a plain ``dict[str, object]`` schema, so restoration matches + entries by index, the same correspondence every other part of this merge + (and the endpoints' full-replace-on-PUT semantics) already assumes. If a + caller both reorders the list AND echoes back a masked marker in the same + request, a credential can end up attached to a different logical entry. + That is a known, narrow limitation -- not a leak between different + agents or tenants, since it only reshuffles one agent's own stored + values -- and this test pins the current, deliberate behavior rather + than asserting it away.""" + existing = { + "model_list": [ + {"model_name": "gpt-4", "litellm_params": {"api_key": SENTINEL_AWS_SECRET_ACCESS_KEY}}, + {"model_name": "claude", "litellm_params": {"api_key": "other-" + SENTINEL_AWS_SECRET_ACCESS_KEY}}, + ], + } + incoming = { + "model_list": [ + # Same index (0) now holds what used to be at index 1's entry. + {"model_name": "claude", "litellm_params": {"api_key": REDACTED_BY_LITELM_STRING}}, + ], + } + + restored = _restore_redacted_litellm_params(incoming, existing) + + assert restored["model_list"][0]["litellm_params"]["api_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY + + +def test_restore_redacted_litellm_params_recovers_a_whole_subtree_collapsed_by_the_depth_cap(): + """Past the read-side recursion depth cap, a whole nested subtree is + collapsed to the flat REDACTED_BY_LITELM marker rather than a dict/list. + If the caller echoes that flat marker back unchanged, the whole + subtree -- not just the literal marker string -- must be restored.""" + existing_subtree = {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, "region": "us-east-1"} + incoming = {"provider_config": REDACTED_BY_LITELM_STRING} + existing = {"provider_config": existing_subtree} + + restored = _restore_redacted_litellm_params(incoming, existing) + + assert restored["provider_config"] == existing_subtree + + +def test_redact_sensitive_agent_litellm_params_does_not_reinterpret_plain_string_values_as_json(): + """A plain non-JSON string value (most string leaves) must pass through + unchanged rather than failing to parse and getting redacted.""" + redacted = redact_sensitive_agent_litellm_params({"model": "bedrock/agentcore/my-agent", "is_public": True}) + + assert redacted["model"] == "bedrock/agentcore/my-agent" + assert redacted["is_public"] is True + + +@pytest.mark.asyncio +async def test_add_agent_to_db_drops_a_sentinel_value_instead_of_storing_the_placeholder(): + """A create has nothing stored to restore behind a redaction marker, so a + sensitive key submitted as the literal marker is dropped rather than + persisted as the placeholder string itself.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + created_agent = MagicMock() + created_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + created_agent.object_permission = None + mock_create = AsyncMock(return_value=created_agent) + mock_prisma.db.litellm_agentstable.create = mock_create + + await registry.add_agent_to_db( + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": { + "aws_secret_access_key": REDACTED_BY_LITELM_STRING, + "model": "bedrock/agentcore/my-agent", + }, + }, + prisma_client=mock_prisma, + created_by="test-user", + ) + + stored_params: Final = json.loads(mock_create.call_args.kwargs["data"]["litellm_params"]) + assert "aws_secret_access_key" not in stored_params + assert stored_params["model"] == "bedrock/agentcore/my-agent" + + +@pytest.mark.asyncio +async def test_update_agent_in_db_preserves_secret_when_echoed_back_redacted(): + """PUT round-trips the GET response, which shows the secret redacted. Saving + an unrelated field change must not overwrite the real stored credential + with the redaction marker.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace( + litellm_params={ + "aws_access_key_id": SENTINEL_AWS_ACCESS_KEY_ID, + "aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, + "model": "bedrock/agentcore/my-agent", + }, + object_permission_id=None, + ) + ) + updated_agent = MagicMock() + updated_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Renamed Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + updated_agent.object_permission = None + mock_update = AsyncMock(return_value=updated_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.update_agent_in_db( + agent_id="agent-123", + agent={ + "agent_name": "Renamed Agent", + "agent_card_params": _sample_agent_card_params(), + # The UI round-tripped the redacted secret and the untouched + # access key id verbatim; only agent_name actually changed. + "litellm_params": { + "aws_access_key_id": SENTINEL_AWS_ACCESS_KEY_ID, + "aws_secret_access_key": REDACTED_BY_LITELM_STRING, + "model": "bedrock/agentcore/my-agent", + }, + }, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) + assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY + assert stored_params["aws_access_key_id"] == SENTINEL_AWS_ACCESS_KEY_ID + assert stored_params["model"] == "bedrock/agentcore/my-agent" + + +@pytest.mark.asyncio +async def test_update_agent_in_db_preserves_secret_when_key_omitted_entirely(): + """Omitting the sensitive key altogether must fall back to the stored + value too, not just an explicit redaction-marker round-trip.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace( + litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, + object_permission_id=None, + ) + ) + updated_agent = MagicMock() + updated_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + updated_agent.object_permission = None + mock_update = AsyncMock(return_value=updated_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.update_agent_in_db( + agent_id="agent-123", + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"model": "bedrock/agentcore/my-agent"}, + }, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) + assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY + + +@pytest.mark.asyncio +async def test_update_agent_in_db_preserves_secret_nested_under_a_non_sensitive_key(): + """A secret nested inside a dict held by a non-sensitively-named key + (e.g. a per-provider sub-config) must also survive an echoed-back + redaction marker, not just top-level secret keys.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace( + litellm_params={ + "provider_config": { + "aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, + "region": "us-east-1", + } + }, + object_permission_id=None, + ) + ) + updated_agent = MagicMock() + updated_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + updated_agent.object_permission = None + mock_update = AsyncMock(return_value=updated_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.update_agent_in_db( + agent_id="agent-123", + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": { + # The GET response redacted the nested secret; the caller + # round-trips it verbatim while changing nothing. + "provider_config": { + "aws_secret_access_key": REDACTED_BY_LITELM_STRING, + "region": "us-west-2", + } + }, + }, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) + assert stored_params["provider_config"]["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY + assert stored_params["provider_config"]["region"] == "us-west-2" + + +@pytest.mark.asyncio +async def test_update_agent_in_db_clears_secret_on_explicit_empty_value(): + """An explicit empty string is a deliberate clear, distinct from an omitted + key or the redaction marker, and must actually clear the stored secret.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace( + litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, + object_permission_id=None, + ) + ) + updated_agent = MagicMock() + updated_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + updated_agent.object_permission = None + mock_update = AsyncMock(return_value=updated_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.update_agent_in_db( + agent_id="agent-123", + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"aws_secret_access_key": ""}, + }, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) + assert stored_params["aws_secret_access_key"] == "" + + +@pytest.mark.asyncio +async def test_patch_agent_in_db_preserves_secret_when_litellm_params_omitted(): + """A PATCH that only renames the agent must not touch (let alone drop) the + stored litellm_params secret.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Old Name", + "litellm_params": {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, + "object_permission_id": None, + } + ) + patched_agent = MagicMock() + patched_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "New Name", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, + "object_permission": None, + } + patched_agent.object_permission = None + mock_update = AsyncMock(return_value=patched_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", + agent={"agent_name": "New Name"}, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + update_data: Final = mock_update.call_args.kwargs["data"] + assert "litellm_params" not in update_data + + +@pytest.mark.asyncio +async def test_patch_agent_in_db_preserves_secret_when_echoed_back_redacted(): + """A PATCH that includes litellm_params (e.g. to flip an unrelated flag) + with the secret round-tripped as the redaction marker must not clobber + the stored credential.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Test Agent", + "litellm_params": { + "aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY, + "is_public": False, + }, + "object_permission_id": None, + } + ) + patched_agent = MagicMock() + patched_agent.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + } + patched_agent.object_permission = None + mock_update = AsyncMock(return_value=patched_agent) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", + agent={ + "litellm_params": { + "aws_secret_access_key": REDACTED_BY_LITELM_STRING, + "is_public": True, + } + }, + prisma_client=mock_prisma, + updated_by="test-user", + ) + + stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) + assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY + assert stored_params["is_public"] is True diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py index 3fb09076e5f..daca244c0a1 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py @@ -16,13 +16,12 @@ from litellm.proxy.agent_endpoints.agent_search import ( AgentSearchHits, AgentSearchIndex, AgentSearchNotConfigured, - Vector, agent_search_text, - cosine_similarity, search_agents, ) from litellm.proxy.agent_endpoints.auth.agent_permission_handler import RestrictedAgentAccess from litellm.proxy.agent_endpoints.endpoints import router, user_api_key_auth +from litellm.proxy.common_utils.semantic_text_index import Vector, cosine_similarity from litellm.types.agents import AgentResponse CALLER: Final = UserAPIKeyAuth(api_key="hashed-caller-key", team_id="team-1", user_id="user-1") diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index ea196bda529..a78b3238a9a 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -4,6 +4,7 @@ import pytest from fastapi import FastAPI from fastapi.testclient import TestClient +from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints import endpoints as agent_endpoints from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( @@ -484,11 +485,15 @@ class TestAgentRBACInternalUserViewOnly: assert resp.status_code == 403 +SENTINEL_AGENT_API_KEY = "sk-test-sentinel-do-not-use" + + class TestAgentRBACProxyAdminViewOnly: """Read-only proxy admins go through the object-permission scoped branch on GET /v1/agents (the admin fast path stays full PROXY_ADMIN only, so viewers - cannot fan out health checks beyond their allowlist), and secret unredaction - also stays gated on full PROXY_ADMIN.""" + cannot fan out health checks beyond their allowlist). litellm_params + secrets are redacted for every caller, admin included (LIT-6736); only the + virtual-key/header visibility stays gated on full PROXY_ADMIN.""" @pytest.fixture(autouse=True) def _setup(self, monkeypatch): @@ -501,7 +506,7 @@ class TestAgentRBACProxyAdminViewOnly: agent_id=f"agent-{index}", agent_name=f"Agent {index}", agent_card_params=_sample_agent_card_params(), - litellm_params={"api_key": "sk-super-secret-agent-key"}, + litellm_params={"api_key": SENTINEL_AGENT_API_KEY}, ) for index in (1, 2) ] @@ -544,7 +549,7 @@ class TestAgentRBACProxyAdminViewOnly: def test_should_still_redact_secrets_for_view_only_admin(self): """An unrestricted viewer sees the same agents as an admin but with keys - stripped and litellm_params masked.""" + stripped; litellm_params secrets never appear in either response.""" self.allowed_agents_spy.return_value = UnrestrictedAgentAccess() viewer_resp = self._list_agents(self.viewer_client) admin_resp = self._list_agents(self.admin_client) @@ -553,14 +558,12 @@ class TestAgentRBACProxyAdminViewOnly: viewer_by_id = {agent["agent_id"]: agent for agent in viewer_resp.json()} assert set(viewer_by_id) == {"agent-1", "agent-2"} assert viewer_by_id["agent-1"]["keys"] is None - assert "sk-super-secret-agent-key" not in viewer_resp.text + assert SENTINEL_AGENT_API_KEY not in viewer_resp.text admin_by_id = {agent["agent_id"]: agent for agent in admin_resp.json()} assert admin_by_id["agent-1"]["keys"][0]["token"] == "hash-aaa" - assert ( - admin_by_id["agent-1"]["litellm_params"]["api_key"] - == "sk-super-secret-agent-key" - ) + assert SENTINEL_AGENT_API_KEY not in admin_resp.text + assert admin_by_id["agent-1"]["litellm_params"]["api_key"] == REDACTED_BY_LITELM_STRING class TestAgentRBACProxyAdmin: @@ -616,6 +619,109 @@ class TestAgentRBACProxyAdmin: # Security scheme is the LiteLLM scheme. assert "LiteLLMKey" in stored_card["securitySchemes"] + def test_create_agent_response_never_echoes_secret(self): + """LIT-6736: POST /v1/agents must not echo the stored secret back, even + though it's the caller's own value and even for a proxy admin.""" + with patch("litellm.proxy.proxy_server.prisma_client"): # test-quality-ok: proxy_server module global is the endpoint's only injection point + self.mock_registry.get_agent_by_name = MagicMock(return_value=None) + self.mock_registry.add_agent_to_db = AsyncMock( + return_value=AgentResponse( + agent_id="agent-123", + agent_name="Test Agent", + agent_card_params=_sample_agent_card_params(), + litellm_params={ + "aws_secret_access_key": SENTINEL_AGENT_API_KEY, + "model": "bedrock/agentcore/my-agent", + }, + ) + ) + self.mock_registry.register_agent = MagicMock() + + resp = self.admin_client.post( + "/v1/agents", + json={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": { + "aws_secret_access_key": SENTINEL_AGENT_API_KEY, + "model": "bedrock/agentcore/my-agent", + }, + }, + headers={"Authorization": "Bearer k"}, + ) + + assert resp.status_code == 200 + assert SENTINEL_AGENT_API_KEY not in resp.text + body = resp.json() + assert body["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + assert body["litellm_params"]["model"] == "bedrock/agentcore/my-agent" + + def test_update_agent_response_never_echoes_secret(self): + """LIT-6736: PUT /v1/agents/{id} must not echo the stored secret back.""" + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: # test-quality-ok: proxy_server module global is the endpoint's only injection point + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Existing Agent", + "agent_card_params": _sample_agent_card_params(), + } + ) + self.mock_registry.update_agent_in_db = AsyncMock( + return_value=AgentResponse( + agent_id="agent-123", + agent_name="Test Agent", + agent_card_params=_sample_agent_card_params(), + litellm_params={"aws_secret_access_key": SENTINEL_AGENT_API_KEY}, + ) + ) + self.mock_registry.deregister_agent = MagicMock() + self.mock_registry.register_agent = MagicMock() + + resp = self.admin_client.put( + "/v1/agents/agent-123", + json={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"aws_secret_access_key": REDACTED_BY_LITELM_STRING}, + }, + headers={"Authorization": "Bearer k"}, + ) + + assert resp.status_code == 200 + assert SENTINEL_AGENT_API_KEY not in resp.text + assert resp.json()["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + + def test_patch_agent_response_never_echoes_secret(self): + """LIT-6736: PATCH /v1/agents/{id} must not echo the stored secret back.""" + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: # test-quality-ok: proxy_server module global is the endpoint's only injection point + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Existing Agent", + "agent_card_params": _sample_agent_card_params(), + } + ) + self.mock_registry.patch_agent_in_db = AsyncMock( + return_value=AgentResponse( + agent_id="agent-123", + agent_name="Renamed Agent", + agent_card_params=_sample_agent_card_params(), + litellm_params={"aws_secret_access_key": SENTINEL_AGENT_API_KEY}, + ) + ) + self.mock_registry.deregister_agent = MagicMock() + self.mock_registry.register_agent = MagicMock() + + resp = self.admin_client.patch( + "/v1/agents/agent-123", + json={"agent_name": "Renamed Agent"}, + headers={"Authorization": "Bearer k"}, + ) + + assert resp.status_code == 200 + assert SENTINEL_AGENT_API_KEY not in resp.text + assert resp.json()["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING + def test_should_allow_admin_to_delete_agent(self): existing = { "agent_id": "agent-123", diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 3dea89ed67b..ffe3e0fee12 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7448,3 +7448,64 @@ async def test_delete_cache_key_object_is_best_effort_when_the_cache_backend_fai healthy_cache.delete_cache.assert_called_once_with(key=hashed_token) healthy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache.assert_awaited_once_with(key=hashed_token) assert caplog.records == [], "a healthy eviction must stay silent, and must still reach both caches" + + +# --------------------------------------------------------------------------- +# Budget-exceeded error text must not carry a raw virtual key (LIT-5909) +# --------------------------------------------------------------------------- + + +class _BudgetAlertRecorder: + async def budget_alerts(self, type, user_info): + return None + + +async def _run_key_budget_check(key_name: str) -> str: + """Drive the over-budget key path and return the raised message.""" + valid_token = UserAPIKeyAuth( + token="hashed-token", + key_name=key_name, + key_alias="prod-key", + spend=10.0, + max_budget=1.0, + ) + with pytest.raises(litellm.BudgetExceededError, match="Budget has been exceeded") as exc_info: + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=_BudgetAlertRecorder(), + ) + await asyncio.sleep(0) + return exc_info.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_name", + [ + "sk-mx5ous1o9Iezz5fj3pkLuA", + "my-company-key-2026", + "sk-...5LuA-but-longer", + # /key/generate takes a custom key ending in an escape sequence, and this + # message reaches a terminal and a log viewer + "sk-...\x1b[2J", + "sk-...a\x9bm", + ], +) +async def test_key_budget_error_does_not_carry_a_raw_key_name(key_name): + """key_name is written masked, but the column has no enforced shape (a direct DB + write bypasses abbreviate_api_key) and this message is returned to the caller.""" + message = await _run_key_budget_check(key_name) + assert key_name not in message + assert "Key=prod-key Current cost" in message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key_name", ["sk-...5LuA", "sk-...", "sk-...ke.!", "sk-...café"]) +async def test_key_budget_error_keeps_the_masked_key_name(key_name): + """The masked form is the whole point of naming the key, so it must survive. + + abbreviate_api_key takes the last four characters of the key verbatim, and a + custom key may end in punctuation or a non-ASCII character, so those masked + names are just as valid as the alphanumeric ones.""" + message = await _run_key_budget_check(key_name) + assert f"Key=prod-key ({key_name}) Current cost" in message diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 90be51cfa5b..21e0b83791f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -9,6 +9,7 @@ from prisma import errors as prisma_errors from prisma.engine.errors import ( BinaryNotFoundError, EngineConnectionError, + EngineRequestError, MismatchedVersionsError, ) from prisma.errors import ( @@ -32,6 +33,12 @@ from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler +class _EngineHttp500: + """The response half of an EngineRequestError: the query engine answered a request with HTTP 500.""" + + status = 500 + + @pytest.mark.asyncio @pytest.mark.parametrize( "db_error", @@ -113,6 +120,90 @@ async def test_handle_authentication_error_permanent_fault_gets_no_fallback_iden assert exc_info.value.code == str(status.HTTP_503_SERVICE_UNAVAILABLE) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prisma_error", + [ + pytest.param(BinaryNotFoundError("query engine binary not found"), id="BinaryNotFoundError"), + pytest.param(MismatchedVersionsError(expected="1", got="2"), id="MismatchedVersionsError"), + pytest.param(EngineRequestError(_EngineHttp500(), "query engine crashed"), id="EngineRequestError"), + pytest.param(PrismaError(), id="bare_PrismaError"), + ], +) +async def test_handle_authentication_error_permanent_fault_503_is_not_worded_as_transient(prisma_error): + """The 503 for a fault that never heals must not say the database is + "temporarily unreachable" and ask the caller to retry. The status is right + (the service is at fault) but that wording sends the operator to wait out an + outage that is not one, so the message has to say retrying will not help and + name the engine fault.""" + handler = UserAPIKeyAuthExceptionHandler() + + with patch( # test-quality-ok: the handler reads general_settings off the proxy module, no injection seam + "litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": False} + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error(prisma_error, MagicMock(), {}, "/test", None, "test-key") + + assert exc_info.value.code == str(status.HTTP_503_SERVICE_UNAVAILABLE) + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "temporarily unreachable" not in exc_info.value.message + assert "retry shortly" not in exc_info.value.message.lower() + assert "will not clear by retrying" in exc_info.value.message + assert type(prisma_error).__name__ in exc_info.value.message + + +@pytest.mark.asyncio +async def test_handle_authentication_error_transport_error_raised_over_a_permanent_fault_names_the_fault(): + """A reconnect attempt that fails because the engine binary is missing surfaces as a transport + error with the BinaryNotFoundError as __context__. The response must describe the binary, which is + what keeps the database down, rather than promise the connection will come back.""" + try: + raise BinaryNotFoundError("query engine binary not found") + except BinaryNotFoundError: + try: + raise httpx.ConnectError("All connection attempts failed") + except httpx.ConnectError as surfaced: + transport_over_fault = surfaced + handler = UserAPIKeyAuthExceptionHandler() + + with patch( # test-quality-ok: the handler reads general_settings off the proxy module, no injection seam + "litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": False} + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error(transport_over_fault, MagicMock(), {}, "/test", None, "k") + + assert exc_info.value.code == str(status.HTTP_503_SERVICE_UNAVAILABLE) + assert "temporarily unreachable" not in exc_info.value.message + assert "BinaryNotFoundError" in exc_info.value.message + assert "will not clear by retrying" in exc_info.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "db_error", + [ + pytest.param(httpx.ConnectError("All connection attempts failed"), id="ConnectError"), + pytest.param(EngineConnectionError(), id="EngineConnectionError"), + pytest.param(PrismaError("can't reach database server"), id="P1001_text"), + ], +) +async def test_handle_authentication_error_transient_outage_503_keeps_retry_wording(db_error): + """A genuine outage is expected to come back, so its 503 keeps telling the + caller the database is temporarily unreachable and to retry.""" + handler = UserAPIKeyAuthExceptionHandler() + + with patch( # test-quality-ok: the handler reads general_settings off the proxy module, no injection seam + "litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": False} + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error(db_error, MagicMock(), {}, "/test", None, "test-key") + + assert exc_info.value.code == str(status.HTTP_503_SERVICE_UNAVAILABLE) + assert exc_info.value.message == ( + "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly." + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( "prisma_error", diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 9301176f3ed..f513f397b64 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -3169,7 +3169,7 @@ class TestIsRequestBodySafeChecksBracketNotationMetadata: class TestHasUserSetupSso: - """_has_user_setup_sso must treat SAML IdP metadata as SSO configured. + """has_user_setup_sso must treat SAML IdP metadata as SSO configured. Regression: UI discovery used this helper for sso_configured, but it only checked OAuth client IDs, so SAML-only setups left the login button gray. @@ -3187,29 +3187,167 @@ class TestHasUserSetupSso: monkeypatch.delenv(key, raising=False) def test_false_when_no_sso_env(self): - from litellm.proxy.auth.auth_utils import _has_user_setup_sso + from litellm.proxy.auth.auth_utils import has_user_setup_sso - assert _has_user_setup_sso() is False + assert has_user_setup_sso() is False def test_true_for_oauth_client_ids(self, monkeypatch): - from litellm.proxy.auth.auth_utils import _has_user_setup_sso + from litellm.proxy.auth.auth_utils import has_user_setup_sso monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-client") - assert _has_user_setup_sso() is True + assert has_user_setup_sso() is True def test_true_for_saml_metadata_url(self, monkeypatch): - from litellm.proxy.auth.auth_utils import _has_user_setup_sso + from litellm.proxy.auth.auth_utils import has_user_setup_sso monkeypatch.setenv( "SAML_IDP_METADATA_URL", "https://idp.example.com/metadata.xml" ) - assert _has_user_setup_sso() is True + assert has_user_setup_sso() is True def test_true_for_saml_metadata_xml(self, monkeypatch): - from litellm.proxy.auth.auth_utils import _has_user_setup_sso + from litellm.proxy.auth.auth_utils import has_user_setup_sso monkeypatch.setenv("SAML_IDP_METADATA_XML", "") - assert _has_user_setup_sso() is True + assert has_user_setup_sso() is True + + +class TestIsSsoProviderFullyConfigured: + """A lone client id must not read as ready: `has_user_setup_sso()` only + checks the client id (correct for a UI-discovery "show the login button" + decision), but a gate that BLOCKS the password fallback needs every + companion setting the provider requires, or an incomplete setup locks + every admin out with no working login path at all.""" + + @pytest.fixture(autouse=True) + def _clear_sso_env(self, monkeypatch): + for key in ( + "GOOGLE_CLIENT_ID", + "GOOGLE_CLIENT_SECRET", + "MICROSOFT_CLIENT_ID", + "MICROSOFT_CLIENT_SECRET", + "MICROSOFT_TENANT", + "GENERIC_CLIENT_ID", + "GENERIC_CLIENT_SECRET", + "GENERIC_AUTHORIZATION_ENDPOINT", + "GENERIC_TOKEN_ENDPOINT", + "GENERIC_USERINFO_ENDPOINT", + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", + ): + monkeypatch.delenv(key, raising=False) + + def test_false_when_nothing_configured(self): + from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured + + assert is_sso_provider_fully_configured() is False + + def test_google_client_id_alone_is_not_ready(self, monkeypatch): + from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured + + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-client") + assert is_sso_provider_fully_configured() is False + + def test_google_with_secret_is_ready(self, monkeypatch): + from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured + + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-client") + monkeypatch.setenv("GOOGLE_CLIENT_SECRET", "google-secret") + assert is_sso_provider_fully_configured() is True + + def test_microsoft_client_id_alone_is_not_ready(self, monkeypatch): + from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured + + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "ms-client") + assert is_sso_provider_fully_configured() is False + + def test_microsoft_missing_tenant_is_not_ready(self, monkeypatch): + from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured + + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "ms-client") + monkeypatch.setenv("MICROSOFT_CLIENT_SECRET", "ms-secret") + assert is_sso_provider_fully_configured() is False + + def test_microsoft_with_secret_and_tenant_is_ready(self, monkeypatch): + from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured + + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "ms-client") + monkeypatch.setenv("MICROSOFT_CLIENT_SECRET", "ms-secret") + monkeypatch.setenv("MICROSOFT_TENANT", "ms-tenant") + assert is_sso_provider_fully_configured() is True + + def test_generic_client_id_alone_is_not_ready(self, monkeypatch): + from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured + + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-client") + assert is_sso_provider_fully_configured() is False + + def test_generic_missing_one_endpoint_is_not_ready(self, monkeypatch): + from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured + + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-client") + monkeypatch.setenv("GENERIC_CLIENT_SECRET", "generic-secret") + monkeypatch.setenv("GENERIC_AUTHORIZATION_ENDPOINT", "https://idp.example.com/authorize") + monkeypatch.setenv("GENERIC_TOKEN_ENDPOINT", "https://idp.example.com/token") + # GENERIC_USERINFO_ENDPOINT deliberately left unset. + assert is_sso_provider_fully_configured() is False + + def test_generic_with_every_endpoint_is_ready(self, monkeypatch): + from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured + + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-client") + monkeypatch.setenv("GENERIC_CLIENT_SECRET", "generic-secret") + monkeypatch.setenv("GENERIC_AUTHORIZATION_ENDPOINT", "https://idp.example.com/authorize") + monkeypatch.setenv("GENERIC_TOKEN_ENDPOINT", "https://idp.example.com/token") + monkeypatch.setenv("GENERIC_USERINFO_ENDPOINT", "https://idp.example.com/userinfo") + assert is_sso_provider_fully_configured() is True + + def test_saml_metadata_url_is_ready_when_runtime_installed(self, monkeypatch): + from litellm.proxy.auth import auth_utils + + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata.xml") + monkeypatch.setattr(auth_utils.importlib.util, "find_spec", lambda name: object()) + assert auth_utils.is_sso_provider_fully_configured() is True + + def test_saml_metadata_url_is_not_ready_without_runtime(self, monkeypatch): + """Regression: python3-saml (``onelogin.saml2``) is an optional + dependency; SAMLAuthHandler fails closed on every request when it is + not installed, so IdP metadata alone must not read as ready.""" + from litellm.proxy.auth import auth_utils + + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata.xml") + monkeypatch.setattr(auth_utils.importlib.util, "find_spec", lambda name: None) + assert auth_utils.is_sso_provider_fully_configured() is False + + def test_saml_check_does_not_raise_when_package_entirely_absent(self, monkeypatch): + """Regression: `importlib.util.find_spec("onelogin.saml2.auth")` + raises ModuleNotFoundError (not merely returns None) when the + TOP-LEVEL `onelogin` package is not installed at all, which is + exactly the real-world "optional extra not installed" case. If the + gate does not catch this, every password login 500s instead of + falling back, on a deployment that configured SAML metadata but + skipped the extra.""" + from litellm.proxy.auth import auth_utils + + def _raise(name: str): + raise ModuleNotFoundError("No module named 'onelogin'") + + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata.xml") + monkeypatch.setattr(auth_utils.importlib.util, "find_spec", _raise) + assert auth_utils.is_sso_provider_fully_configured() is False + + def test_incomplete_earlier_provider_does_not_mask_a_ready_later_one(self, monkeypatch): + """Regression: a stray GOOGLE_CLIENT_ID with no secret (e.g. a + leftover from a migration) must not stop the check from reaching a + fully configured Microsoft provider set alongside it — every + provider is evaluated independently, not in a first-match order.""" + from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured + + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-client") + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "ms-client") + monkeypatch.setenv("MICROSOFT_CLIENT_SECRET", "ms-secret") + monkeypatch.setenv("MICROSOFT_TENANT", "ms-tenant") + assert is_sso_provider_fully_configured() is True class TestIsRequestBodySafeBlocksAwsIdentitySelectors: diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 1c66acf8678..8d93d801bfd 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -6,6 +6,7 @@ to login_utils.py for better reusability. """ import os +from contextlib import ExitStack from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -598,3 +599,203 @@ class TestEncodeUiSessionJwt: request.cookies = {"token": token} with patch("litellm.proxy.proxy_server.master_key", "sk-master-for-tests"): assert _user_id_from_session_cookie(request) == "cornell-user" + + +def _patch_sso_configured(stack: ExitStack, *, configured: bool) -> None: + stack.enter_context( + patch( # test-quality-ok: no HTTP boundary here; same internal the pre-existing tests above already mock + "litellm.proxy.auth.login_utils.is_sso_provider_fully_configured", return_value=configured + ) + ) + + +def _patch_successful_admin_login_deps(stack: ExitStack) -> None: + """The collaborators a real admin login exercises past the SSO gate: + generating the UI session key, syncing the admin role, and reading the + experimental-login flag. Shared so the two "still allowed" tests below + don't each repeat the same three-mock wiring.""" + stack.enter_context( + patch( # test-quality-ok: internal orchestration, no HTTP boundary; matches pre-existing tests + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"token": "test-token", "user_id": LITELLM_PROXY_ADMIN_NAME}, + ) + ) + stack.enter_context( + patch( # test-quality-ok: internal orchestration, no HTTP boundary; matches pre-existing tests + "litellm.proxy.auth.login_utils.user_update", + new_callable=AsyncMock, + return_value=None, + ) + ) + stack.enter_context( + patch( # test-quality-ok: internal orchestration, no HTTP boundary; matches pre-existing tests + "litellm.proxy.auth.login_utils.get_secret_bool", + return_value=False, + ) + ) + + +class TestDisablePasswordLoginWhenSSOEnabled: + """`disable_password_login_when_sso_enabled` must reject every + username/password login attempt (including the UI_USERNAME/UI_PASSWORD + admin fallback) once SSO is configured, so SSO becomes the only way to + reach the Admin UI. It must not affect logins when SSO is unconfigured, + so admins can never lock themselves out with no SSO to fall back to.""" + + @pytest.mark.asyncio + async def test_rejects_correct_admin_credentials_when_sso_configured(self): + master_key = "sk-1234" + ui_username = "admin" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": master_key}): + with ExitStack() as stack: + _patch_sso_configured(stack, configured=True) + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username=ui_username, + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_password_login_when_sso_enabled": True}, + ) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == "403" + # The credential comparison must never even run. + mock_prisma_client.db.litellm_usertable.find_first.assert_not_called() + + @pytest.mark.asyncio + async def test_rejects_correct_db_user_credentials_when_sso_configured(self): + master_key = "sk-1234" + user_email = "test@example.com" + password = "correct-password" + + mock_user = LiteLLM_UserTable( + user_id="test-user-123", + user_email=user_email, + password=hash_token(token=password), + user_role=LitellmUserRoles.INTERNAL_USER, + ) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=mock_user) + + with patch.dict(os.environ, {"UI_USERNAME": "admin", "UI_PASSWORD": "unrelated"}): + with ExitStack() as stack: + _patch_sso_configured(stack, configured=True) + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username=user_email, + password=password, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_password_login_when_sso_enabled": True}, + ) + + assert exc_info.value.code == "403" + mock_prisma_client.db.litellm_usertable.find_first.assert_not_called() + + @pytest.mark.asyncio + async def test_allows_password_login_when_setting_enabled_but_sso_not_configured(self): + """The setting alone must not lock out an admin who has not actually + configured SSO — there would be no fallback left.""" + master_key = "sk-1234" + ui_username = "admin" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict( + os.environ, + { + "UI_USERNAME": ui_username, + "UI_PASSWORD": master_key, + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + clear=True, + ): + with ExitStack() as stack: + _patch_sso_configured(stack, configured=False) + _patch_successful_admin_login_deps(stack) + result = await authenticate_user( + username=ui_username, + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_password_login_when_sso_enabled": True}, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + + @pytest.mark.asyncio + async def test_allows_password_login_when_sso_env_is_incomplete(self): + """Regression: a lone MICROSOFT_CLIENT_ID with no client secret or + tenant makes has_user_setup_sso() True, but a real SSO sign-in would + fail. The gate must read the real env (no is_sso_provider_fully_configured + mock here) and still let password login through, or an admin who set + one env var by mistake is locked out with no way in.""" + master_key = "sk-1234" + ui_username = "admin" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict( + os.environ, + { + "UI_USERNAME": ui_username, + "UI_PASSWORD": master_key, + "DATABASE_URL": "postgresql://test:test@localhost/test", + "MICROSOFT_CLIENT_ID": "ms-client-id-only", + }, + clear=True, + ): + with ExitStack() as stack: + _patch_successful_admin_login_deps(stack) + result = await authenticate_user( + username=ui_username, + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={"disable_password_login_when_sso_enabled": True}, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + + @pytest.mark.asyncio + async def test_allows_password_login_when_sso_configured_but_setting_not_enabled(self): + """SSO being configured must not, by itself, disable the password + fallback: the setting is opt-in.""" + master_key = "sk-1234" + ui_username = "admin" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict( + os.environ, + { + "UI_USERNAME": ui_username, + "UI_PASSWORD": master_key, + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + clear=True, + ): + with ExitStack() as stack: + _patch_sso_configured(stack, configured=True) + _patch_successful_admin_login_deps(stack) + result = await authenticate_user( + username=ui_username, + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + general_settings={}, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME diff --git a/tests/test_litellm/proxy/auth/test_onboarding.py b/tests/test_litellm/proxy/auth/test_onboarding.py index d55a5472af1..524b655b465 100644 --- a/tests/test_litellm/proxy/auth/test_onboarding.py +++ b/tests/test_litellm/proxy/auth/test_onboarding.py @@ -231,7 +231,7 @@ async def test_claim_token_rejects_already_used_link(): data = InvitationClaim( invitation_link="invite-abc", user_id="user-123", - password="NewP@ssw0rd", + password="NewP@ssw0rd123", ) with patch("litellm.proxy.proxy_server.prisma_client", prisma): @@ -254,7 +254,7 @@ async def test_claim_token_rejects_expired_link(): data = InvitationClaim( invitation_link="invite-abc", user_id="user-123", - password="NewP@ssw0rd", + password="NewP@ssw0rd123", ) with patch("litellm.proxy.proxy_server.prisma_client", prisma): @@ -275,7 +275,7 @@ async def test_claim_token_rejects_mismatched_user_id(): data = InvitationClaim( invitation_link="invite-abc", user_id="wrong-user", - password="NewP@ssw0rd", + password="NewP@ssw0rd123", ) with patch("litellm.proxy.proxy_server.prisma_client", prisma): @@ -296,7 +296,7 @@ async def test_claim_token_rejects_missing_onboarding_token(): data = InvitationClaim( invitation_link="invite-abc", user_id="user-123", - password="NewP@ssw0rd", + password="NewP@ssw0rd123", ) with ( @@ -322,7 +322,7 @@ async def test_claim_token_rejects_wrong_onboarding_session(): data = InvitationClaim( invitation_link="invite-abc", user_id="user-123", - password="NewP@ssw0rd", + password="NewP@ssw0rd123", ) request = _make_claim_request( _make_onboarding_token(invitation_link="other-invite") @@ -351,7 +351,7 @@ async def test_claim_token_rejects_invalid_bearer_token(): data = InvitationClaim( invitation_link="invite-abc", user_id="user-123", - password="NewP@ssw0rd", + password="NewP@ssw0rd123", ) request = _make_claim_request("sk-regular-key") @@ -380,7 +380,7 @@ async def test_claim_token_rejects_concurrent_reuse_before_password_write(): data = InvitationClaim( invitation_link="invite-abc", user_id="user-123", - password="NewP@ssw0rd", + password="NewP@ssw0rd123", ) with ( @@ -418,7 +418,7 @@ async def test_claim_token_sets_accepted_at_after_password_written(): data = InvitationClaim( invitation_link="invite-abc", user_id="user-123", - password="NewP@ssw0rd", + password="NewP@ssw0rd123", ) mock_token_response = {"token": "sk-generated-key", "user_id": "user-123"} @@ -477,7 +477,7 @@ async def test_claim_token_rolls_back_invite_when_session_key_mint_fails(): data = InvitationClaim( invitation_link="invite-abc", user_id="user-123", - password="NewP@ssw0rd", + password="NewP@ssw0rd123", ) with ( diff --git a/tests/test_litellm/proxy/auth/test_password_policy.py b/tests/test_litellm/proxy/auth/test_password_policy.py new file mode 100644 index 00000000000..f6e7d443907 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_password_policy.py @@ -0,0 +1,136 @@ +""" +Tests for the configurable password-strength policy in +`litellm.proxy.auth.password_policy`, enforced on every path that persists a +new or changed password for a locally-managed user. +""" + +import pytest + +from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.auth.password_policy import ( + DEFAULT_MIN_LENGTH, + MIN_ALLOWED_LENGTH, + PasswordPolicy, + get_password_policy, + validate_password_policy, +) + +STRONG_PASSWORD = "Str0ng!Passw0rd" + + +def test_get_password_policy_defaults_to_pif_baseline(): + policy = get_password_policy({}) + assert policy == PasswordPolicy( + min_length=DEFAULT_MIN_LENGTH, + require_uppercase=True, + require_lowercase=True, + require_numbers=True, + require_special_characters=True, + ) + + +def test_get_password_policy_reads_overrides_from_general_settings(): + policy = get_password_policy( + { + "password_policy_min_length": 20, + "password_policy_require_uppercase": False, + "password_policy_require_lowercase": False, + "password_policy_require_numbers": False, + "password_policy_require_special_characters": False, + } + ) + assert policy == PasswordPolicy( + min_length=20, + require_uppercase=False, + require_lowercase=False, + require_numbers=False, + require_special_characters=False, + ) + + +def test_validate_password_policy_accepts_strong_password(): + assert validate_password_policy(STRONG_PASSWORD, {}) is None + + +@pytest.mark.parametrize( + "password,expected_fragment", + [ + ("Sh0rt!Pw", "12 characters"), + ("weakpassword123!", "uppercase"), + ("WEAKPASSWORD123!", "lowercase"), + ("WeakPassword!!!!", "number"), + ("WeakPassword12345", "special character"), + ], +) +def test_validate_password_policy_rejects_each_missing_class(password, expected_fragment): + with pytest.raises(ProxyException) as exc_info: + validate_password_policy(password, {}) + assert exc_info.value.code == "400" + assert exc_info.value.type == ProxyErrorTypes.validation_error + assert exc_info.value.param == "password" + assert expected_fragment in exc_info.value.message + + +def test_validate_password_policy_reports_every_violation_at_once(): + with pytest.raises(ProxyException) as exc_info: + validate_password_policy("weak", {}) + assert "12 characters" in exc_info.value.message + assert "uppercase" in exc_info.value.message + assert "number" in exc_info.value.message + assert "special character" in exc_info.value.message + + +def test_validate_password_policy_honors_relaxed_config(): + general_settings = { + "password_policy_min_length": MIN_ALLOWED_LENGTH, + "password_policy_require_special_characters": False, + } + # 8 chars, has upper/lower/number, no special char: fails default policy, + # passes the relaxed one above. + validate_password_policy("Abcd1234", general_settings) + with pytest.raises(ProxyException): + validate_password_policy("Abcd1234", {}) + + +def test_validate_password_policy_honors_stricter_min_length(): + general_settings = {"password_policy_min_length": 20} + with pytest.raises(ProxyException) as exc_info: + validate_password_policy(STRONG_PASSWORD, general_settings) + assert "20 characters" in exc_info.value.message + + +@pytest.mark.parametrize("configured_min_length", [0, -1, -100, 1, 7]) +def test_get_password_policy_floors_nonpositive_or_too_low_min_length(configured_min_length): + """A misconfigured min_length must never disable the length check + entirely: it floors at MIN_ALLOWED_LENGTH instead of passing through.""" + policy = get_password_policy({"password_policy_min_length": configured_min_length}) + assert policy.min_length == MIN_ALLOWED_LENGTH + + +def test_validate_password_policy_rejects_short_password_even_with_zero_min_length_configured(): + general_settings = {"password_policy_min_length": 0} + with pytest.raises(ProxyException) as exc_info: + validate_password_policy("a", general_settings) + assert f"{MIN_ALLOWED_LENGTH} characters" in exc_info.value.message + + +def test_get_password_policy_ignores_boolean_min_length(): + """`bool` is a subclass of `int` in Python; a stray `true`/`false` value + must not silently coerce into a min_length of 1 or 0.""" + policy = get_password_policy({"password_policy_min_length": False}) + assert policy.min_length == DEFAULT_MIN_LENGTH + + +def test_validate_password_policy_rejects_unicode_letter_as_special_character(): + """Regression: an ASCII-only `[^A-Za-z0-9]` check would miscount an + accented letter as the required special character, so a letters-and- + digits-only password like this one (no real symbol) must still be + rejected.""" + with pytest.raises(ProxyException) as exc_info: + validate_password_policy("Passwörd1234", {}) + assert "special character" in exc_info.value.message + + +def test_validate_password_policy_accepts_real_special_character_with_unicode_letters(): + """Same base password as the rejection test above, plus an actual symbol.""" + assert validate_password_policy("Passwörd1234!", {}) is None diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index d80e3acb4b8..c685d778c0e 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -1,12 +1,14 @@ import asyncio import json import sys +from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest from fastapi import HTTPException, Request from prisma import errors as prisma_errors +from prisma.engine.errors import BinaryNotFoundError, EngineConnectionError from prisma.errors import ( ClientNotConnectedError, DataError, @@ -317,6 +319,43 @@ def test_is_database_service_unavailable_error_in_chain_sees_through_wrapping(): assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(ValueError("nope")) is False +def test_find_database_service_unavailable_error_in_chain_returns_the_wrapped_outage_itself(): + """Wording a 503 by the kind of outage needs the wrapped database error, not the ValueError + get_user_object wrapped it in, so the finder must hand back the inner exception.""" + outage = _wrapped_like_get_user_object(ConnectionError("can't reach database server")) + found = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(outage) + assert isinstance(found, ConnectionError) + assert found is outage.__context__ + missing_user = _wrapped_like_get_user_object(Exception()) + assert PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(missing_user) is None + + +def _raised_while_handling(inner, outer): + try: + raise inner + except BaseException: + try: + raise outer + except BaseException as surfaced: + return surfaced + + +def test_permanent_fault_outranks_the_transient_error_that_surfaced_it(): + """A reconnect that dies on a missing engine binary raises the transport error last, with the + BinaryNotFoundError left as __context__. The binary is what keeps the database down, so both the + finder and the 503 wording must pick it over the outer transient error, whichever way they nest.""" + permanent = BinaryNotFoundError("query engine binary not found") + transient_over_permanent = _raised_while_handling(permanent, httpx.ConnectError("connection refused")) + permanent_over_transient = _raised_while_handling(httpx.ConnectError("connection refused"), permanent) + + for chain in (transient_over_permanent, permanent_over_transient): + assert PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(chain) is permanent + message = PrismaDBExceptionHandler.database_unavailable_message(chain) + assert "BinaryNotFoundError" in message + assert "will not clear by retrying" in message + assert "temporarily unreachable" not in message + + def test_is_database_service_unavailable_error_in_chain_terminates_on_a_cause_cycle(): """The walk must terminate on a pathological __cause__ cycle rather than hang. Neither link is an outage, so the bounded walk returns False instead of looping forever.""" @@ -508,6 +547,51 @@ def test_permanent_prisma_faults_are_still_reported_as_service_problems(prisma_e assert PrismaDBExceptionHandler.is_database_service_unavailable_error(prisma_error) is True +RECONNECTABLE_CLIENT_STATE_FAULTS = (prisma_errors.ClientNotConnectedError, prisma_errors.HTTPClientClosedError) + + +@pytest.mark.parametrize("prisma_error", PERMANENT_PRISMA_FAULTS) +def test_permanent_prisma_faults_are_worded_as_not_retryable(prisma_error): + """A 503 for a fault that never heals must not tell the operator to wait. + + The status stays 503 (the service is at fault), but the message has to say + the outage is not transient and name the engine fault, or an operator + watching a version-skewed engine keeps retrying a request that can never + succeed. The two client-state faults a reconnect can repair keep the retry + wording.""" + reconnectable = isinstance(prisma_error, RECONNECTABLE_CLIENT_STATE_FAULTS) + message = PrismaDBExceptionHandler.database_unavailable_message(prisma_error) + + assert PrismaDBExceptionHandler.is_permanent_database_fault(prisma_error) is (not reconnectable) + assert message.startswith("Service Unavailable") + assert ("temporarily unreachable" in message) is reconnectable + assert ("Please retry shortly" in message) is reconnectable + assert ("will not clear by retrying" in message) is (not reconnectable) + assert (type(prisma_error).__name__ in message) is (not reconnectable) + + +@pytest.mark.parametrize( + "transient_error", + [ + pytest.param(httpx.ConnectError("All connection attempts failed"), id="ConnectError"), + pytest.param(ConnectionError("connection refused"), id="ConnectionError"), + pytest.param(EngineConnectionError(), id="EngineConnectionError"), + pytest.param(prisma_errors.PrismaError("can't reach database server"), id="P1001_text"), + pytest.param( + ProxyException(message="no db", type=ProxyErrorTypes.no_db_connection, param=None, code=503), + id="ProxyException", + ), + ], +) +def test_transient_outages_keep_the_retry_wording(transient_error): + """A genuine outage is expected to come back, so the retry guidance is the + right message and must not be replaced by the permanent-fault text.""" + assert PrismaDBExceptionHandler.is_permanent_database_fault(transient_error) is False + assert PrismaDBExceptionHandler.database_unavailable_message(transient_error) == ( + "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly." + ) + + @pytest.mark.parametrize( "transient_error", [ @@ -579,3 +663,40 @@ def test_is_deadlock_error_matches_postgres_deadlock(error): def test_is_deadlock_error_excludes_non_deadlocks(error): """Non-deadlock prisma errors, connectivity failures, and non-prisma exceptions are not treated as deadlocks.""" assert PrismaDBExceptionHandler.is_deadlock_error(error) is False + + +MOCKED_PRISMA_PREDICATES: Final = ( + PrismaDBExceptionHandler.is_database_infrastructure_error, + PrismaDBExceptionHandler.is_database_transport_error, + PrismaDBExceptionHandler.is_deadlock_error, + PrismaDBExceptionHandler.is_prisma_engine_internal_error, + PrismaDBExceptionHandler.is_database_service_unavailable_error, +) + + +@pytest.mark.parametrize("predicate", MOCKED_PRISMA_PREDICATES, ids=lambda p: p.__name__) +def test_predicates_answer_false_for_a_plain_exception_when_prisma_is_mocked(predicate): + """Suites that swap ``sys.modules["prisma"]`` for a ``MagicMock`` hand the + predicates mocks in place of prisma's error classes. ``isinstance`` against + a mock raises ``TypeError``; the predicate must instead answer for the + non-prisma checks it still has.""" + with patch.dict(sys.modules, {"prisma": MagicMock()}): + assert predicate(Exception("db connection dropped")) is False + + +def test_infrastructure_error_still_recognizes_transport_errors_when_prisma_is_mocked(): + """Skipping the prisma classes must not skip the checks that do not need them.""" + with patch.dict(sys.modules, {"prisma": MagicMock()}): + no_db: Final = ProxyException(message="no db", type=ProxyErrorTypes.no_db_connection, param=None, code=503) + assert PrismaDBExceptionHandler.is_database_infrastructure_error(httpx.ConnectError("refused")) is True + assert PrismaDBExceptionHandler.is_database_infrastructure_error(no_db) is True + + +def test_connection_error_answers_when_prisma_is_mocked_after_import(): + """``prisma.engine`` is already loaded in a real process, so a mock parent + still resolves ``prisma.engine.errors``; its classes are then mocks too.""" + import prisma.engine.errors # noqa: F401 + + with patch.dict(sys.modules, {"prisma": MagicMock()}): + assert PrismaDBExceptionHandler.is_database_connection_error(Exception("x")) is False + assert PrismaDBExceptionHandler.is_database_connection_error(httpx.ConnectError("refused")) is True diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py index f4da8c941a4..c3f7b0100d8 100644 --- a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py +++ b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -18,7 +18,7 @@ def test_ui_discovery_endpoints_with_defaults(): with ( patch("litellm.proxy.utils.get_server_root_path", return_value="/"), patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): @@ -41,7 +41,7 @@ def test_ui_discovery_endpoints_with_custom_server_root_path(): with ( patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"), patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): @@ -66,7 +66,7 @@ def test_ui_discovery_endpoints_with_proxy_base_url_when_set(): "litellm.proxy.utils.get_proxy_base_url", return_value="https://proxy.example.com", ), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): @@ -91,7 +91,7 @@ def test_ui_discovery_endpoints_with_sso_configured_and_auto_redirect_enabled(): "litellm.proxy.utils.get_proxy_base_url", return_value="https://proxy.example.com", ), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=True), patch.dict( os.environ, {"AUTO_REDIRECT_UI_LOGIN_TO_SSO": "true", "DISABLE_ADMIN_UI": "false"}, @@ -121,7 +121,7 @@ def test_ui_discovery_endpoints_with_sso_configured_and_auto_redirect_not_set_de "litellm.proxy.utils.get_proxy_base_url", return_value="https://proxy.example.com", ), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=True), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): # Ensure AUTO_REDIRECT_UI_LOGIN_TO_SSO is not set (simulate default) @@ -148,7 +148,7 @@ def test_ui_discovery_endpoints_with_sso_configured_but_auto_redirect_disabled() "litellm.proxy.utils.get_proxy_base_url", return_value="https://proxy.example.com", ), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=True), patch.dict( os.environ, {"AUTO_REDIRECT_UI_LOGIN_TO_SSO": "false", "DISABLE_ADMIN_UI": "false"}, @@ -174,7 +174,7 @@ def test_ui_discovery_endpoints_with_sso_not_configured_but_auto_redirect_enable with ( patch("litellm.proxy.utils.get_server_root_path", return_value="/"), patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch.dict( os.environ, {"AUTO_REDIRECT_UI_LOGIN_TO_SSO": "true", "DISABLE_ADMIN_UI": "false"}, @@ -203,7 +203,7 @@ def test_ui_discovery_endpoints_both_routes_return_same_data(): "litellm.proxy.utils.get_proxy_base_url", return_value="https://proxy.example.com", ), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=True), patch.dict( os.environ, {"AUTO_REDIRECT_UI_LOGIN_TO_SSO": "true", "DISABLE_ADMIN_UI": "false"}, @@ -228,7 +228,7 @@ def test_ui_discovery_endpoints_with_auto_redirect_via_general_settings(): with ( patch("litellm.proxy.utils.get_server_root_path", return_value="/"), patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=True), patch( "litellm.proxy.proxy_server.general_settings", {"auto_redirect_ui_login_to_sso": True}, @@ -254,7 +254,7 @@ def test_ui_discovery_endpoints_with_auto_redirect_env_var_overrides_general_set with ( patch("litellm.proxy.utils.get_server_root_path", return_value="/"), patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=True), patch( "litellm.proxy.proxy_server.general_settings", {"auto_redirect_ui_login_to_sso": False}, @@ -281,7 +281,7 @@ def test_ui_discovery_endpoints_with_admin_ui_disabled(): with ( patch("litellm.proxy.utils.get_server_root_path", return_value="/"), patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "true"}, clear=False), ): @@ -311,7 +311,7 @@ def test_ui_discovery_endpoints_is_control_plane_true_when_workers_configured(): with ( patch("litellm.proxy.utils.get_server_root_path", return_value="/"), patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch("litellm.proxy.proxy_server.proxy_config", mock_config), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): @@ -336,7 +336,7 @@ def test_ui_discovery_endpoints_hide_default_credentials_hint_default_false(): with ( patch("litellm.proxy.utils.get_server_root_path", return_value="/"), patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): os.environ.pop("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", None) @@ -357,7 +357,7 @@ def test_ui_discovery_endpoints_hide_default_credentials_hint_via_env_var(): with ( patch("litellm.proxy.utils.get_server_root_path", return_value="/"), patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch.dict( os.environ, { @@ -384,7 +384,7 @@ def test_ui_discovery_endpoints_hide_default_credentials_hint_via_general_settin with ( patch("litellm.proxy.utils.get_server_root_path", return_value="/"), patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch( "litellm.proxy.proxy_server.general_settings", {"hide_default_credentials_hint": True}, @@ -411,7 +411,7 @@ def test_ui_discovery_endpoints_is_control_plane_false_when_no_workers(): with ( patch("litellm.proxy.utils.get_server_root_path", return_value="/"), patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), - patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), + patch("litellm.proxy.auth.auth_utils.has_user_setup_sso", return_value=False), patch("litellm.proxy.proxy_server.proxy_config", mock_config), patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False), ): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 4ee6741ee02..84f7611c0c0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -842,24 +842,37 @@ async def test_presidio_filter_scope_initializer(monkeypatch): params_input = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="input") guardrail_dict = {"guardrail_name": "g1"} - cb = initialize_presidio(params_input, guardrail_dict) - assert cb is created[0] + callbacks = initialize_presidio(params_input, guardrail_dict) + assert callbacks == (created[0],) assert created[0].apply_to_output is False # output-only created.clear() params_output = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="output") - cb = initialize_presidio(params_output, guardrail_dict) + callbacks = initialize_presidio(params_output, guardrail_dict) assert len(created) == 1 + assert callbacks == (created[0],) assert created[0].apply_to_output is True - # both -> expect two callbacks (input + output) + # both -> expect two callbacks (input + output), both returned, input first created.clear() params_both = LitellmParams(guardrail="presidio", mode="pre_call", presidio_filter_scope="both") - cb = initialize_presidio(params_both, guardrail_dict) + callbacks = initialize_presidio(params_both, guardrail_dict) assert len(created) == 2 - assert any(not c.apply_to_output for c in created) - assert any(c.apply_to_output for c in created) + assert callbacks == tuple(created) + assert callbacks[0].apply_to_output is False + assert callbacks[1].apply_to_output is True + + # both + output_parse_pii -> three callbacks, all returned, input first + created.clear() + params_all = LitellmParams( + guardrail="presidio", mode="pre_call", presidio_filter_scope="both", output_parse_pii=True + ) + callbacks = initialize_presidio(params_all, guardrail_dict) + assert len(created) == 3 + assert callbacks == tuple(created) + assert callbacks[0].apply_to_output is False + assert mgr.added[-3:] == list(created) @pytest.mark.asyncio @@ -3116,6 +3129,18 @@ def test_update_in_memory_applies_analyze_chunk_size(): assert guardrail.presidio_analyze_chunk_size_bytes == 99_000 +def test_update_in_memory_keeps_output_masker_from_unmasking(): + masker = _OPTIONAL_PresidioPIIMasking(mock_testing=True, apply_to_output=True, output_parse_pii=False) + unmasker = _OPTIONAL_PresidioPIIMasking(mock_testing=True, output_parse_pii=True) + params = LitellmParams(guardrail="presidio", mode="pre_call", output_parse_pii=True) + + masker.update_in_memory_litellm_params(params) + unmasker.update_in_memory_litellm_params(params) + + assert (masker.apply_to_output, masker.output_parse_pii) == (True, False) + assert (unmasker.apply_to_output, unmasker.output_parse_pii) == (False, True) + + def test_merge_drops_truncated_same_type_fragment_from_overlap(): """A boundary entity seen truncated by chunk 1 and whole by chunk 2 must merge to the single full span; keeping both overlapping spans corrupts the diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 2c0735970d3..24742e1bac2 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -1,3 +1,4 @@ +from collections.abc import Iterable from unittest.mock import AsyncMock, MagicMock import pytest @@ -491,6 +492,144 @@ def test_repeated_db_sync_does_not_accumulate_runner_instances(): cb_list[:] = snapshot +PRESIDIO_SIBLINGS_GID = "55555555-5555-5555-5555-555555555555" +PRESIDIO_SIBLINGS_NAME = "presidio-siblings" + + +def _presidio_db_guardrail(pii_entities_config: dict[str, str]) -> Guardrail: + return Guardrail( + guardrail_id=PRESIDIO_SIBLINGS_GID, + guardrail_name=PRESIDIO_SIBLINGS_NAME, + litellm_params={ + "guardrail": "presidio", + "mode": "pre_call", + "default_on": True, + "output_parse_pii": True, + "presidio_filter_scope": "both", + "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze", + "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize", + "pii_entities_config": pii_entities_config, + }, + ) + + +def _presidio_callbacks_in(cb_list: Iterable[object]) -> list[CustomGuardrail]: + return [ + callback + for callback in cb_list + if isinstance(callback, CustomGuardrail) and getattr(callback, "guardrail_name", None) == PRESIDIO_SIBLINGS_NAME + ] + + +def test_presidio_siblings_are_tracked_and_deleted_together(): + """ + A presidio guardrail scoped to both stages registers the pre_call primary plus + the post_call unmask and mask-output siblings. Deleting the guardrail must remove + all three from every callback list, not just the primary. + """ + import litellm + + handler = InMemoryGuardrailHandler() + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler.initialize_guardrail(_presidio_db_guardrail({"EMAIL_ADDRESS": "MASK"})) + + registered = _presidio_callbacks_in(litellm.callbacks) + assert len(registered) == 3 + primary = handler.guardrail_id_to_custom_guardrail[PRESIDIO_SIBLINGS_GID] + siblings = handler.guardrail_id_to_sibling_callbacks[PRESIDIO_SIBLINGS_GID] + assert primary is registered[0] + assert siblings == tuple(registered[1:]) + assert [sibling.event_hook for sibling in siblings] == [GuardrailEventHooks.post_call] * 2 + + for cb_list in lists[1:]: + cb_list.extend(registered) + + handler.delete_in_memory_guardrail(PRESIDIO_SIBLINGS_GID) + + for cb_list in lists: + assert _presidio_callbacks_in(cb_list) == [] + assert PRESIDIO_SIBLINGS_GID not in handler.guardrail_id_to_custom_guardrail + assert PRESIDIO_SIBLINGS_GID not in handler.guardrail_id_to_sibling_callbacks + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + +def test_update_in_memory_guardrail_reaches_presidio_siblings_and_keeps_their_stage(): + import litellm + + handler = InMemoryGuardrailHandler() + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + handler.initialize_guardrail(_presidio_db_guardrail({"EMAIL_ADDRESS": "MASK", "IP_ADDRESS": "MASK"})) + tracked = _presidio_callbacks_in(litellm.callbacks) + roles_before = [ + (callback.apply_to_output, callback.output_parse_pii, callback.event_hook) for callback in tracked + ] + assert roles_before == [ + (False, True, [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call]), + (False, True, GuardrailEventHooks.post_call), + (True, False, GuardrailEventHooks.post_call), + ] + + updated = Guardrail( + guardrail_id=PRESIDIO_SIBLINGS_GID, + guardrail_name=PRESIDIO_SIBLINGS_NAME, + litellm_params=LitellmParams( + guardrail="presidio", + mode="pre_call", + default_on=True, + output_parse_pii=True, + presidio_filter_scope="both", + presidio_analyzer_api_base="https://fakelink.com/v1/presidio/analyze", + presidio_anonymizer_api_base="https://fakelink.com/v1/presidio/anonymize", + pii_entities_config={"EMAIL_ADDRESS": "MASK"}, + ), + ) + handler.update_in_memory_guardrail(guardrail_id=PRESIDIO_SIBLINGS_GID, guardrail=updated) + + assert [callback.pii_entities_config for callback in tracked] == [{"EMAIL_ADDRESS": "MASK"}] * 3 + assert [ + (callback.apply_to_output, callback.output_parse_pii, callback.event_hook) for callback in tracked + ] == roles_before + assert _presidio_callbacks_in(litellm.callbacks) == tracked + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + +def test_repeated_db_sync_replaces_presidio_siblings_instead_of_leaking_stale_ones(): + """ + The callback manager dedupes custom loggers by their scalar attributes, so a + leaked post_call sibling blocks the re-initialized sibling from registering and + keeps serving the previous entity config. After every DB re-sync, each callback + list must hold exactly the three current instances, all on the latest config. + """ + import litellm + + handler = InMemoryGuardrailHandler() + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + entity_configs = [{"EMAIL_ADDRESS": "MASK"}, {"EMAIL_ADDRESS": "MASK", "IP_ADDRESS": "MASK"}] + for cycle in range(4): + latest = entity_configs[cycle % 2] + handler.sync_guardrail_from_db(_presidio_db_guardrail(latest)) + for cb_list in lists[1:]: + cb_list.extend(_presidio_callbacks_in(litellm.callbacks)) + + for cb_list in lists: + current = _presidio_callbacks_in(cb_list) + assert len({id(callback) for callback in current}) == 3 + assert all(callback.pii_entities_config == latest for callback in current) + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + def _judge_guardrail(guardrail_id: str) -> Guardrail: return Guardrail( guardrail_id=guardrail_id, diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 8bce967b316..f231eb66a50 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -4220,3 +4220,88 @@ async def test_user_new_persists_model_max_budget( ) assert captured["user_data"].get("model_max_budget") == expected_written + + +@pytest.fixture +def _admin_prisma(mocker): + """A mocked prisma_client wired in as proxy_server's module globals, for + the password-policy tests below (mirrors the pattern every other test in + this file repeats per-test; consolidated here since these three share it + verbatim).""" + mock_prisma_client = mocker.MagicMock() + mocker.patch( # test-quality-ok: same module-global mocking every test in this file already uses + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) + mocker.patch( # test-quality-ok: same module-global mocking every test in this file already uses + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ) + return mock_prisma_client + + +@pytest.mark.asyncio +async def test_user_update_rejects_weak_password(_admin_prisma): + """/user/update must reject a password that fails the configured + policy before it ever reaches the DB write.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + user_request = UpdateUserRequest(user_id="target-user", password="short1!") + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(ProxyException) as exc_info: + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) + + assert exc_info.value.code == "400" + _admin_prisma.db.litellm_usertable.find_first.assert_not_called() + + +@pytest.mark.asyncio +async def test_user_update_rejects_weak_password_against_configured_policy(_admin_prisma, mocker): + """A password that meets the default policy but not a stricter + admin-configured one must still be rejected.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mocker.patch( # test-quality-ok: same module-global mocking every test in this file already uses + "litellm.proxy.proxy_server.general_settings", + {"password_policy_min_length": 24}, + ) + + user_request = UpdateUserRequest(user_id="target-user", password="Str0ng!Passw0rd") + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(ProxyException) as exc_info: + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) + + assert "24 characters" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_user_update_hashes_and_persists_strong_password(_admin_prisma, mocker): + """A password meeting the policy is hashed (never stored in plaintext) + and reaches the DB write.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = _admin_prisma + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "target-user"} + existing_user.user_id = "target-user" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user"}) + mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + + strong_password = "Str0ng!Passw0rd" + user_request = UpdateUserRequest(user_id="target-user", password=strong_password) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) + + written_data = mock_prisma_client.update_data.call_args.kwargs["data"] + assert written_data.get("password") is not None + assert written_data["password"] != strong_password diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 87e0319f6a1..bca97915347 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -4555,3 +4555,114 @@ def test_create_file_async_pre_call_hook_rejection_blocks_upload(monkeypatch, ll assert response.status_code == 400, response.text assert "file upload not allowed" in response.text assert provider_route.call_count == 0 + + +def test_create_file_non_batch_over_max_file_size_mb_rejected_before_forwarding(monkeypatch, llm_router: Router): + """max_file_size_mb applies to every purpose, unlike the batch-only max_batch_file_size_mb.""" + import litellm.proxy.proxy_server as ps + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + monkeypatch.setitem(ps.general_settings, "max_file_size_mb", 1) + + oversized = b"x" * (2 * 1024 * 1024) + try: + response = client.post( + "/v1/files", + files={"file": ("labels.jsonl", oversized, "application/octet-stream")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 413, response.text + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert error["param"] == "file" + assert "max_file_size_mb" in error["message"] + assert "1 MB" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_non_batch_under_max_file_size_mb_forwards(monkeypatch, llm_router: Router): + import litellm.proxy.proxy_server as ps + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + monkeypatch.setitem(ps.general_settings, "max_file_size_mb", 1) + + try: + response = client.post( + "/v1/files", + files={"file": ("labels.jsonl", b"small content", "application/octet-stream")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 200, response.text + assert len(forwarded_calls) == 1 + + +def test_create_file_blocked_extension_rejected_before_forwarding(monkeypatch, llm_router: Router): + import litellm.proxy.proxy_server as ps + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + monkeypatch.setitem(ps.general_settings, "blocked_file_extensions", [".exe", ".sh"]) + + try: + response = client.post( + "/v1/files", + files={"file": ("payload.exe", b"MZ\x90\x00", "application/octet-stream")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert error["param"] == "file" + assert ".exe" in error["message"] + assert "blocked_file_extensions" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_blocked_extension_unset_allows_everything(monkeypatch, llm_router: Router): + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + + try: + response = client.post( + "/v1/files", + files={"file": ("payload.exe", b"MZ\x90\x00", "application/octet-stream")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 200, response.text + assert len(forwarded_calls) == 1 + + +def test_create_file_path_traversal_filename_rejected_before_forwarding(monkeypatch, llm_router: Router): + """A filename carrying a directory-traversal component must never reach storage or the provider.""" + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + + try: + response = client.post( + "/v1/files", + files={"file": ("../../etc/passwd", b"malicious content", "text/plain")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert error["param"] == "file" + assert "traversal" in error["message"].lower() + assert forwarded_calls == [] diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py b/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py new file mode 100644 index 00000000000..9f7dd914e4a --- /dev/null +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_general_upload_validation.py @@ -0,0 +1,134 @@ +import io + +import pytest + +from litellm.proxy._types import ProxyException +from litellm.proxy.openai_files_endpoints.general_upload_validation import ( + MB, + UploadedFileBlockedExtension, + UploadedFileTooLarge, + UploadedFileUnsafeFilename, + check_blocked_extension, + check_unsafe_filename, + check_upload_file_size, + raise_upload_validation_failure, +) + + +def test_size_under_cap_allowed(): + assert check_upload_file_size(b"x" * 100, 1) is None + + +def test_size_over_cap_rejected_for_bytes(): + content = b"x" * (2 * MB) + assert check_upload_file_size(content, 1) == UploadedFileTooLarge(size_bytes=len(content), limit_mb=1) + + +def test_size_over_cap_rejected_for_binaryio_and_restores_caller_position(): + """The handle is caller-owned; inspecting its size must not discard where the caller had it.""" + content = b"x" * (2 * MB) + handle = io.BytesIO(content) + handle.seek(17) + assert check_upload_file_size(handle, 1) == UploadedFileTooLarge(size_bytes=len(content), limit_mb=1) + assert handle.tell() == 17 + + +def test_size_under_cap_allowed_for_binaryio_restores_caller_position(): + handle = io.BytesIO(b"x" * 100) + handle.seek(42) + assert check_upload_file_size(handle, 1) is None + assert handle.tell() == 42 + + +def test_size_exactly_at_cap_allowed(): + content = b"x" * MB + assert check_upload_file_size(content, 1) is None + + +def test_no_cap_skips_size_check(): + assert check_upload_file_size(b"x" * (10 * MB), None) is None + + +@pytest.mark.parametrize("cap", [0, -3]) +def test_nonpositive_cap_disables_size_check(cap): + assert check_upload_file_size(b"x" * (10 * MB), cap) is None + + +def test_blocked_extension_rejected(): + assert check_blocked_extension("payload.exe", (".exe", ".sh")) == UploadedFileBlockedExtension(extension=".exe") + + +def test_blocked_extension_match_is_case_insensitive(): + assert check_blocked_extension("payload.EXE", (".exe",)) == UploadedFileBlockedExtension(extension=".exe") + + +def test_blocked_extension_match_is_case_insensitive_for_configured_value(): + """A config entry like blocked_file_extensions: ['.EXE'] must still catch a lowercase upload.""" + assert check_blocked_extension("payload.exe", (".EXE",)) == UploadedFileBlockedExtension(extension=".exe") + + +def test_extension_not_in_blocklist_allowed(): + assert check_blocked_extension("report.pdf", (".exe", ".sh")) is None + + +def test_empty_blocklist_allows_everything(): + assert check_blocked_extension("payload.exe", ()) is None + + +def test_no_filename_skips_extension_check(): + assert check_blocked_extension(None, (".exe",)) is None + + +def test_path_traversal_filename_rejected(): + assert check_unsafe_filename("../../etc/passwd") == UploadedFileUnsafeFilename(filename="../../etc/passwd") + + +def test_windows_style_path_traversal_filename_rejected(): + assert check_unsafe_filename("..\\..\\windows\\system32\\config") == UploadedFileUnsafeFilename( + filename="..\\..\\windows\\system32\\config" + ) + + +def test_traversal_embedded_after_extension_rejected(): + assert check_unsafe_filename("report.jsonl/../../etc/cron.d/evil") == UploadedFileUnsafeFilename( + filename="report.jsonl/../../etc/cron.d/evil" + ) + + +def test_null_byte_filename_rejected(): + assert check_unsafe_filename("report.pdf\x00.exe") == UploadedFileUnsafeFilename(filename="report.pdf\x00.exe") + + +@pytest.mark.parametrize("filename", ["report.pdf", ".env", "a.b.c.jsonl", "my file (1).csv", None]) +def test_ordinary_filenames_allowed(filename): + assert check_unsafe_filename(filename) is None + + +@pytest.mark.parametrize( + "failure, expected_code, expected_fragments", + [ + ( + UploadedFileTooLarge(size_bytes=15728640, limit_mb=10), + "413", + ("15.0 MB", "max_file_size_mb", "10 MB", "not forwarded"), + ), + ( + UploadedFileBlockedExtension(extension=".exe"), + "400", + (".exe", "blocked_file_extensions", "not forwarded"), + ), + ( + UploadedFileUnsafeFilename(filename="../../etc/passwd"), + "400", + ("../../etc/passwd", "traversal", "not forwarded"), + ), + ], +) +def test_failures_map_to_openai_shaped_proxy_exceptions(failure, expected_code, expected_fragments): + with pytest.raises(ProxyException) as exc_info: + raise_upload_validation_failure(failure) + assert exc_info.value.code == expected_code + assert exc_info.value.type == "invalid_request_error" + assert exc_info.value.param == "file" + for fragment in expected_fragments: + assert fragment in exc_info.value.message diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py index af37dbe85fe..45460dcecf1 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py @@ -29,7 +29,7 @@ def _install_login_mocks(monkeypatch, raise_on_auth: bool = False) -> None: """ from litellm.proxy import proxy_server as ps - async def _fake_auth(username, password, master_key, prisma_client): + async def _fake_auth(username, password, master_key, prisma_client, general_settings=None): if raise_on_auth: raise Exception("boom-auth-failure") fake = MagicMock() diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py b/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py index 5cc22cca7a0..778acc1baab 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_onboarding.py @@ -234,7 +234,7 @@ def test_claim_onboarding_link_happy(client, monkeypatch, mock_prisma): json={ "invitation_link": "inv-123", "user_id": "user-abc", - "password": "hunter2", + "password": "Hunter2Strong!", }, headers={"Authorization": f"Bearer {onboarding_jwt}"}, ) @@ -260,7 +260,7 @@ def test_claim_onboarding_link_invalid_invite_401(client, monkeypatch, mock_pris json={ "invitation_link": "missing", "user_id": "user-abc", - "password": "hunter2", + "password": "Hunter2Strong!", }, headers={"Authorization": "Bearer irrelevant"}, ) @@ -287,7 +287,7 @@ def test_claim_onboarding_link_user_id_mismatch_401( json={ "invitation_link": "inv-123", "user_id": "user-attacker", - "password": "hunter2", + "password": "Hunter2Strong!", }, headers={"Authorization": "Bearer irrelevant"}, ) @@ -339,7 +339,7 @@ def test_claim_onboarding_link_bad_onboarding_jwt_401( json={ "invitation_link": "inv-123", "user_id": "user-abc", - "password": "hunter2", + "password": "Hunter2Strong!", }, headers={"Authorization": f"Bearer {bogus_jwt}"}, ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 91fca8f1e27..5f692f8f109 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -130,6 +130,7 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): password="secret", master_key="test-master-key", prisma_client=mock_prisma_client, + general_settings={}, ) mock_create_ui_token_object.assert_called_once_with( login_result=mock_login_result, @@ -10782,11 +10783,54 @@ def test_update_config_redacts_all_environment_variable_values(_update_config_se class _EnvBuiltRedisCache(RedisCache): """RedisCache stand-in that records its constructor kwargs and never opens a network connection, so tests can assert which connection params - the proxy used to build its coordination Redis.""" + the proxy used to build its coordination Redis. `ping()` reports reachable + by default, matching a real Redis the env fallback should adopt.""" def __init__(self, **kwargs): self.init_kwargs = kwargs + async def ping(self) -> bool: + return True + + +class _UnreachableRedisCache(_EnvBuiltRedisCache): + """Same stand-in, but `ping()` fails like a REDIS_* env var naming a Redis + that is not actually reachable (wrong host, no service running, ...).""" + + async def ping(self) -> bool: + raise ConnectionError("connection refused") + + +@contextlib.contextmanager +def _patched_coordination_redis_module_state( + *, + spend_cache: DualCache, + config_cache: types.SimpleNamespace, + redis_cache_class: type = _EnvBuiltRedisCache, +): + """Stub every `litellm.proxy.proxy_server` global that + `_attach_redis_usage_cache` (and its callers) can write to, shared by the + whole coordination-Redis test family below. + + Centralizing this is not just DRY: `_attach_redis_usage_cache` always sets + `cli_sso_session_cache.redis_cache` unconditionally, and a call site that + forgets to patch that one real (persistent) global leaks a throwaway + Redis stand-in into it for the rest of the pytest session, breaking + unrelated tests that run later. One patched-state helper means a new call + site cannot forget a global this family already knows to isolate. + """ + with ( + patch.object(proxy_server_module, "redis_usage_cache", None), + patch.object(proxy_server_module, "spend_counter_cache", spend_cache), + patch.object(proxy_server_module, "user_api_key_cache", DualCache()), + patch.object(proxy_server_module, "cli_sso_session_cache", DualCache()), + patch.object(proxy_server_module, "llm_router", None), + patch.object(proxy_server_module, "litellm_config_cache", config_cache), + patch.object(proxy_server_module, "RedisCache", redis_cache_class), + patch.object(proxy_server_module, "RedisClusterCache", _EnvBuiltClusterCache), + ): + yield + def _run_init_cache_with_backend(cache_backend, redis_env_kwargs): """Run ProxyConfig._init_cache with a stubbed response-cache backend and a @@ -10798,12 +10842,7 @@ def _run_init_cache_with_backend(cache_backend, redis_env_kwargs): fresh_config_cache = types.SimpleNamespace(redis_cache=None) with ( - patch.object(proxy_server_module, "redis_usage_cache", None), - patch.object(proxy_server_module, "spend_counter_cache", fresh_spend_cache), - patch.object(proxy_server_module, "user_api_key_cache", DualCache()), - patch.object(proxy_server_module, "llm_router", None), - patch.object(proxy_server_module, "litellm_config_cache", fresh_config_cache), - patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache), + _patched_coordination_redis_module_state(spend_cache=fresh_spend_cache, config_cache=fresh_config_cache), patch( "litellm._redis._redis_kwargs_from_environment", return_value=redis_env_kwargs, @@ -10878,12 +10917,7 @@ def _run_init_coordination_redis(config, env=None): fresh_config_cache = types.SimpleNamespace(redis_cache=None) with ( - patch.object(proxy_server_module, "redis_usage_cache", None), - patch.object(proxy_server_module, "spend_counter_cache", fresh_spend_cache), - patch.object(proxy_server_module, "user_api_key_cache", DualCache()), - patch.object(proxy_server_module, "litellm_config_cache", fresh_config_cache), - patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache), - patch.object(proxy_server_module, "RedisClusterCache", _EnvBuiltClusterCache), + _patched_coordination_redis_module_state(spend_cache=fresh_spend_cache, config_cache=fresh_config_cache), mock.patch.dict(os.environ, env or {}, clear=False), ): built = proxy_server_module.ProxyConfig()._init_coordination_redis(config=config) @@ -10971,13 +11005,7 @@ def test_explicit_coordination_redis_takes_precedence_over_cache_backend(): mock_litellm_cache.cache = cache_backend with ( - patch.object(proxy_server_module, "redis_usage_cache", None), - patch.object(proxy_server_module, "spend_counter_cache", fresh_spend_cache), - patch.object(proxy_server_module, "user_api_key_cache", DualCache()), - patch.object(proxy_server_module, "llm_router", None), - patch.object(proxy_server_module, "litellm_config_cache", fresh_config_cache), - patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache), - patch.object(proxy_server_module, "RedisClusterCache", _EnvBuiltClusterCache), + _patched_coordination_redis_module_state(spend_cache=fresh_spend_cache, config_cache=fresh_config_cache), patch("litellm.Cache", return_value=mock_litellm_cache), ): litellm.cache = None @@ -10995,6 +11023,95 @@ def test_explicit_coordination_redis_takes_precedence_over_cache_backend(): assert fresh_spend_cache.redis_cache is usage_cache +async def _run_init_coordination_redis_env_fallback( + litellm_settings, redis_env_kwargs, redis_cache_class=_EnvBuiltRedisCache +): + """Run ProxyConfig._init_coordination_redis_env_fallback against a + stubbed module state and a controlled REDIS_* environment, returning + (built, spend_counter redis).""" + fresh_spend_cache = DualCache() + fresh_config_cache = types.SimpleNamespace(redis_cache=None) + + with ( + _patched_coordination_redis_module_state( + spend_cache=fresh_spend_cache, config_cache=fresh_config_cache, redis_cache_class=redis_cache_class + ), + patch( + "litellm._redis._redis_kwargs_from_environment", + return_value=redis_env_kwargs, + ), + ): + built = await proxy_server_module.ProxyConfig._init_coordination_redis_env_fallback( + litellm_settings=litellm_settings + ) + return built, fresh_spend_cache.redis_cache + + +@pytest.mark.asyncio +async def test_init_coordination_redis_env_fallback_builds_from_environment(): + """A deployment with no coordination_redis block and no litellm_settings.cache + but with bare REDIS_HOST/REDIS_PORT env vars must still get a coordination + Redis: otherwise spend counters, budget-window enforcement, and the + reset_spend cache-eviction broadcast stay per-pod local and a reset issued + on one pod never clears another pod's stale enforcement.""" + built, spend_redis = await _run_init_coordination_redis_env_fallback( + litellm_settings={}, + redis_env_kwargs={"host": "env-fallback-host", "port": "6390"}, + ) + + assert isinstance(built, _EnvBuiltRedisCache) + assert built.init_kwargs["host"] == "env-fallback-host" + assert spend_redis is built + + +@pytest.mark.asyncio +async def test_init_coordination_redis_env_fallback_without_redis_env_returns_none(): + """With no REDIS_* connection info at all, the fallback must leave the + coordination Redis unset rather than building a broken client.""" + built, spend_redis = await _run_init_coordination_redis_env_fallback( + litellm_settings={}, + redis_env_kwargs={}, + ) + + assert built is None + assert spend_redis is None + + +@pytest.mark.asyncio +async def test_init_coordination_redis_env_fallback_unreachable_stays_in_memory(): + """REDIS_* env vars can name a Redis that is not actually reachable (wrong + host, leftover from an unrelated job/service). Guessing "coordination + available" from bare env vars must not turn a previously harmless + in-memory-only proxy into one that raises on its next cache write, so an + unreachable ping must leave everything exactly as if no REDIS_* vars were + set at all.""" + built, spend_redis = await _run_init_coordination_redis_env_fallback( + litellm_settings={}, + redis_env_kwargs={"host": "unreachable-host", "port": "6390"}, + redis_cache_class=_UnreachableRedisCache, + ) + + assert built is None + assert spend_redis is None + + +@pytest.mark.asyncio +async def test_init_coordination_redis_env_fallback_malformed_cluster_nodes_stays_in_memory(): + """REDIS_CLUSTER_NODES can be set to a malformed value nothing here ever + asked to be parsed. Unlike the explicit coordination_redis block (a + deliberate opt-in, so a bad value there should fail loudly), this + inferred fallback must not abort proxy startup over it -- it has to + decline the same way it does for an absent or unreachable Redis.""" + with mock.patch.dict(os.environ, {"REDIS_CLUSTER_NODES": "not-valid-json"}, clear=False): + built, spend_redis = await _run_init_coordination_redis_env_fallback( + litellm_settings={}, + redis_env_kwargs={}, + ) + + assert built is None + assert spend_redis is None + + def test_env_fallback_builds_cluster_client_from_cluster_nodes_env(): """A deployment whose only Redis env is REDIS_CLUSTER_NODES must still get a coordination Redis from the env fallback, and it must be a cluster @@ -11036,11 +11153,7 @@ async def test_startup_applies_coordination_redis_saved_in_database(): fresh_config_cache = types.SimpleNamespace(redis_cache=None) with ( - patch.object(proxy_server_module, "spend_counter_cache", fresh_spend_cache), - patch.object(proxy_server_module, "user_api_key_cache", DualCache()), - patch.object(proxy_server_module, "litellm_config_cache", fresh_config_cache), - patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache), - patch.object(proxy_server_module, "RedisClusterCache", _EnvBuiltClusterCache), + _patched_coordination_redis_module_state(spend_cache=fresh_spend_cache, config_cache=fresh_config_cache), patch.object( proxy_server_module, "get_persisted_coordination_redis_settings", @@ -12189,3 +12302,97 @@ async def test_load_config_router_authorizes_fallback_targets_against_the_callin router, _, _ = await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) assert router.fallback_access_check is router_fallback_access_check + + +def test_docs_redoc_openapi_are_reachable_by_default(): + """ + LIT-6745: the interactive/machine-readable docs surfaces are on by + default (the customer-facing production toggle is opt-in, not opt-out). + """ + client = TestClient(app) + + assert client.get("/redoc").status_code == 200 + openapi_response = client.get("/openapi.json") + assert openapi_response.status_code == 200 + assert "paths" in openapi_response.json() + + +def test_production_app_docs_urls_are_wired_to_the_real_env_helpers(): + """ + LIT-6745: pins the actual `FastAPI(docs_url=..., redoc_url=..., openapi_url=...)` + construction in proxy_server.py to _get_docs_url/_get_redoc_url/_get_openapi_url, + so a hardcoded or drifted value at that call site fails this test even though + the helpers themselves are covered separately. + """ + from litellm.proxy import utils as proxy_utils + + assert app.docs_url == proxy_utils._get_docs_url() + assert app.redoc_url == proxy_utils._get_redoc_url() + assert app.openapi_url == proxy_utils._get_openapi_url() + + +def _build_app_with_docs_env(monkeypatch, *, disabled: bool) -> FastAPI: + from litellm.proxy import utils as proxy_utils + from litellm.proxy.health_endpoints._health_endpoints import router as health_router + + for flag in ("DOCS_URL", "REDOC_URL", "OPENAPI_URL"): + monkeypatch.delenv(flag, raising=False) + for flag in ("NO_DOCS", "NO_REDOC", "NO_OPENAPI"): + if disabled: + monkeypatch.setenv(flag, "True") + else: + monkeypatch.delenv(flag, raising=False) + + # Mirrors the exact FastAPI() construction in proxy_server.py, so this + # exercises the real gating mechanism rather than a reimplementation of it. + app_under_test = FastAPI( + docs_url=proxy_utils._get_docs_url(), + redoc_url=proxy_utils._get_redoc_url(), + openapi_url=proxy_utils._get_openapi_url(), + ) + app_under_test.include_router(health_router) + return app_under_test + + +def test_docs_endpoints_enabled_when_env_unset(monkeypatch): + app_under_test = _build_app_with_docs_env(monkeypatch, disabled=False) + assert app_under_test.docs_url == "/" + assert app_under_test.redoc_url == "/redoc" + assert app_under_test.openapi_url == "/openapi.json" + + client = TestClient(app_under_test) + assert client.get(app_under_test.docs_url).status_code == 200 + assert client.get(app_under_test.redoc_url).status_code == 200 + assert client.get(app_under_test.openapi_url).status_code == 200 + + +def test_no_docs_no_redoc_no_openapi_disable_every_documentation_surface(monkeypatch): + """ + LIT-6745: NO_DOCS, NO_REDOC and NO_OPENAPI must each 404 their surface + with no schema in the body, so a production/air-gapped deployment can + restrict every doc route consistently. + """ + app_under_test = _build_app_with_docs_env(monkeypatch, disabled=True) + assert app_under_test.docs_url is None + assert app_under_test.redoc_url is None + assert app_under_test.openapi_url is None + + client = TestClient(app_under_test) + for route in ("/", "/redoc", "/openapi.json"): + response = client.get(route) + assert response.status_code == 404 + assert "openapi" not in response.text.lower() + assert "paths" not in response.text.lower() + + +def test_disabling_docs_does_not_disable_other_routes(monkeypatch): + """ + LIT-6745: disabling the doc surfaces must not affect inference/management + routes, since NO_DOCS/NO_REDOC/NO_OPENAPI only remove the routes FastAPI + itself auto-registers for docs_url/redoc_url/openapi_url. + """ + app_under_test = _build_app_with_docs_env(monkeypatch, disabled=True) + client = TestClient(app_under_test) + + assert client.get("/redoc").status_code == 404 + assert client.get("/health/liveliness").status_code == 200 diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index dc256ccf718..709447d23c0 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3006,6 +3006,128 @@ def test_update_mcp_semantic_filter_settings_requires_proxy_admin(monkeypatch): app.dependency_overrides.pop(user_api_key_auth, None) +class TestMcpToolSearchSettingsEndpoints: + """`litellm_settings.mcp_tool_search` drives the native `mcp_tool_search` virtual tool, so the UI must round-trip it.""" + + @staticmethod + def _override_auth(role: LitellmUserRoles): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="u", api_key="hashed", user_role=role + ) + + def test_get_returns_stored_values_and_field_schema(self, mock_proxy_config, mock_auth, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", object()) + mock_proxy_config["config"]["litellm_settings"]["mcp_tool_search"] = { + "embedding_model": "text-embedding-3-small", + "core_tools": ["treasury-get_rates"], + } + + resp = client.get("/get/mcp_tool_search_settings") + + assert resp.status_code == 200, resp.text + assert resp.json()["values"] == { + "embedding_model": "text-embedding-3-small", + "top_k": 5, + "similarity_threshold": 0.0, + "core_tools": ["treasury-get_rates"], + } + assert resp.json()["field_schema"]["properties"]["core_tools"]["type"] == "array" + + def test_update_requires_proxy_admin(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + self._override_auth(LitellmUserRoles.INTERNAL_USER) + try: + resp = client.patch("/update/mcp_tool_search_settings", json={"top_k": 3}) + finally: + app.dependency_overrides.clear() + assert resp.status_code == 403 + + def test_update_persists_and_applies_in_memory(self, mock_proxy_config, monkeypatch): + import litellm + + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + monkeypatch.setattr(litellm, "mcp_tool_search", None) + self._override_auth(LitellmUserRoles.PROXY_ADMIN) + payload = { + "embedding_model": "text-embedding-3-small", + "top_k": 3, + "similarity_threshold": 0.25, + "core_tools": ["treasury-get_rates"], + } + try: + resp = client.patch("/update/mcp_tool_search_settings", json=payload) + finally: + app.dependency_overrides.clear() + + assert resp.status_code == 200, resp.text + assert mock_proxy_config["save_call_count"]() == 1 + assert litellm.mcp_tool_search == payload + assert mock_proxy_config["config"]["litellm_settings"]["mcp_tool_search"] == payload + + def test_update_rejects_out_of_range_top_k(self, mock_proxy_config, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + self._override_auth(LitellmUserRoles.PROXY_ADMIN) + try: + resp = client.patch("/update/mcp_tool_search_settings", json={"top_k": 0}) + finally: + app.dependency_overrides.clear() + assert resp.status_code == 422 + assert mock_proxy_config["save_call_count"]() == 0 + + +def test_upload_logo_requires_proxy_admin(monkeypatch): + """Any authenticated key could previously write a file to the server's disk here.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + async def _internal_user_auth(): + return UserAPIKeyAuth( + user_id="internal-user-1", + api_key="hashed-internal-key", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + app.dependency_overrides[user_api_key_auth] = _internal_user_auth + try: + resp = client.post( + "/upload/logo", + files={"file": ("logo.png", b"\x89PNG\r\n\x1a\n" + b"x" * 32, "image/png")}, + ) + assert resp.status_code == 403 + assert "proxy admin" in resp.json()["detail"].lower() + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + +def test_upload_logo_allows_proxy_admin(monkeypatch, tmp_path): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + async def _admin_auth(): + return UserAPIKeyAuth( + user_id="admin-1", + api_key="hashed-admin-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + app.dependency_overrides[user_api_key_auth] = _admin_auth + try: + resp = client.post( + "/upload/logo", + files={"file": ("logo.png", b"\x89PNG\r\n\x1a\n" + b"x" * 32, "image/png")}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["status"] == "success" + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + uploaded_path = resp.json().get("file_path") + if uploaded_path and os.path.exists(uploaded_path): + os.remove(uploaded_path) + + class TestPtuCostAttributionUISetting: """``enable_ptu_cost_attribution`` is derived from the environment on every GET. diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index 087a1c8b3ad..bf3757e6886 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -1,26 +1,30 @@ import ast import asyncio import json +import logging import re import sys +import time +from io import StringIO from pathlib import Path from typing import List import pytest -import logging - import litellm from litellm._logging import ( _COLOR_LOG_FORMAT, + _MAX_SCRUBBED_ACCESS_ARG, _PLAIN_LOG_FORMAT, ALL_LOGGERS, + AccessLogRedactionFilter, CorrelationContextFilter, CorrelationPlainFormatter, JsonFormatter, LevelRoutingStreamHandler, SecretRedactionFilter, StdoutLogTruncationFilter, + _get_uvicorn_json_log_config, _initialize_loggers_with_handler, _parse_json_logs_env, _plain_log_format, @@ -968,3 +972,209 @@ def test_plain_log_format_survives_none_streams(): """sys.stdout/sys.stderr can be None in embedded interpreters; import must not crash.""" assert _plain_log_format(None, None) == _PLAIN_LOG_FORMAT assert _plain_log_format(_FakeStream(True), None) == _PLAIN_LOG_FORMAT + + +# --------------------------------------------------------------------------- +# Access-log redaction (LIT-5909) +# --------------------------------------------------------------------------- + +_LEAKED_KEY = "sk-mx5ous1o9Iezz5fj3pkLuA" + + +def _access_record(full_path: str) -> logging.LogRecord: + """A record shaped exactly like the one uvicorn.access emits per request.""" + return logging.LogRecord( + name="uvicorn.access", + level=logging.INFO, + pathname="", + lineno=0, + msg='%s - "%s %s HTTP/%s" %d', + args=("127.0.0.1:1", "GET", full_path, "1.1", 200), + exc_info=None, + ) + + +@pytest.mark.parametrize( + "full_path", + [ + f"/key/info?key={_LEAKED_KEY}", + f"/global/spend/report?api_key={_LEAKED_KEY}&start_date=2026-08-01", + f"/key/spend/report?api_key={_LEAKED_KEY}", + f"/spend/logs?api_key={_LEAKED_KEY}", + f"/user/daily/activity?api_key={_LEAKED_KEY}", + f"/gemini/v1beta/models/gemini-2.0-flash:generateContent?key={_LEAKED_KEY}", + ], +) +def test_access_log_filter_redacts_a_credential_query_parameter(full_path): + record = _access_record(full_path) + assert AccessLogRedactionFilter().filter(record) is True + assert _LEAKED_KEY not in record.getMessage() + assert "REDACTED" in record.getMessage() + + +def test_access_log_filter_keeps_the_record_formattable_by_uvicorn(): + """uvicorn's AccessFormatter unpacks record.args, so the filter must scrub the + args in place rather than collapse them the way SecretRedactionFilter does.""" + from uvicorn.logging import AccessFormatter + + record = _access_record(f"/key/info?key={_LEAKED_KEY}") + AccessLogRedactionFilter().filter(record) + assert isinstance(record.args, tuple) + assert len(record.args) == 5 + + formatted = AccessFormatter('%(client_addr)s - "%(request_line)s" %(status_code)s', use_colors=False).format(record) + assert _LEAKED_KEY not in formatted + assert "GET" in formatted + assert "200 OK" in formatted + + +@pytest.mark.parametrize( + "full_path, want", + [ + # The delimiter must survive so the logged request line stays well formed. + (f"/key/info?key={_LEAKED_KEY}&page=2", "/key/info?REDACTED&page=2"), + ("/download?sig=AbCd1234%2Fxy&page=2", "/download?REDACTED&page=2"), + ( + f"/global/spend/report?api_key={_LEAKED_KEY}&start_date=2026-01-01", + "/global/spend/report?REDACTED&start_date=2026-01-01", + ), + ("/sso/callback?client_secret=abcdefgh12345&state=xyz", "/sso/callback?REDACTED&state=xyz"), + (f"/v1/models?token={_LEAKED_KEY}&page=2", "/v1/models?REDACTED&page=2"), + ], +) +def test_access_log_filter_keeps_the_query_delimiter(full_path, want): + record = _access_record(full_path) + AccessLogRedactionFilter().filter(record) + assert record.args[2] == want + + +@pytest.mark.parametrize( + "full_path, want", + [ + # Both the param name and the value are encoded, so neither is literal text + # the patterns can see, yet the request parser decodes it into a working key. + (f"/key/info?k%65y=sk%2D{_LEAKED_KEY[3:]}", "/key/info?REDACTED"), + (f"/key/info?k%65y=sk%2D{_LEAKED_KEY[3:]}&page=2", "/key/info?REDACTED"), + (f"/v1/models/sk%2D{_LEAKED_KEY[3:]}", "REDACTED"), + # A decoded credential must never be echoed back: it can carry a newline and + # forge a following log line. + (f"/v1/models?k%65y=sk%2D{_LEAKED_KEY[3:]}%0AINFO:%20forged", "/v1/models?REDACTED"), + ], +) +def test_access_log_filter_redacts_a_percent_encoded_credential(full_path, want): + record = _access_record(full_path) + AccessLogRedactionFilter().filter(record) + assert record.args[2] == want + + +@pytest.mark.parametrize( + "full_path", + [ + "/v1/models?filter=gpt%2D4o&page=2", + "/gemini/v1beta/models/gemini-2.0-flash%3AgenerateContent", + ], +) +def test_access_log_filter_leaves_harmless_percent_encoding_alone(full_path): + """Decoding is a detector, not a rewrite, so a request line with no credential + in it survives encoded exactly as the client sent it.""" + record = _access_record(full_path) + AccessLogRedactionFilter().filter(record) + assert record.args[2] == full_path + + +def test_access_log_filter_caps_how_much_of_a_request_target_it_scans(): + """The request target is the only input to the secret regex an unauthenticated + caller controls end to end, so it is bounded before it is scanned, and the + dropped tail must not reach the log either.""" + record = _access_record("/v1/models?u=" + "a://" * 8192 + f"&key={_LEAKED_KEY}") + + started = time.perf_counter() + AccessLogRedactionFilter().filter(record) + elapsed = time.perf_counter() - started + + scrubbed = record.args[2] + assert _LEAKED_KEY not in scrubbed + assert len(scrubbed) < 1024 + assert elapsed < 1.0, f"scrubbing one access line took {elapsed:.2f}s" + + +@pytest.mark.parametrize("chars_before_the_cut", range(1, 12)) +def test_access_log_filter_never_logs_a_half_scanned_credential(chars_before_the_cut): + """Cutting mid-value would leave a prefix too short for the key= pattern to match, + and that prefix would then be logged raw, so the cut lands on a param boundary.""" + prefix = "/v1/models?u=" + padding = _MAX_SCRUBBED_ACCESS_ARG - len(prefix) - len("&key=") - chars_before_the_cut + record = _access_record(f"{prefix}{'a' * padding}&key={_LEAKED_KEY}") + + AccessLogRedactionFilter().filter(record) + + assert f"key={_LEAKED_KEY[:chars_before_the_cut]}" not in record.args[2] + + +def test_access_log_filter_leaves_a_credential_free_request_line_intact(): + record = _access_record("/v1/chat/completions") + AccessLogRedactionFilter().filter(record) + assert record.getMessage() == '127.0.0.1:1 - "GET /v1/chat/completions HTTP/1.1" 200' + + +def test_access_log_filter_redacts_a_record_that_carries_no_positional_args(): + record = logging.LogRecord( + name="uvicorn.access", + level=logging.INFO, + pathname="", + lineno=0, + msg=f'127.0.0.1:1 - "GET /key/info?key={_LEAKED_KEY} HTTP/1.1" 200', + args=None, + exc_info=None, + ) + assert AccessLogRedactionFilter().filter(record) is True + assert _LEAKED_KEY not in record.getMessage() + + +def _emit_access_line(full_path: str) -> str: + """Hand one real record to uvicorn.access and return what a handler wrote out.""" + from uvicorn.logging import AccessFormatter + + logger = logging.getLogger("uvicorn.access") + stream = StringIO() + handler = logging.StreamHandler(stream) + handler.setFormatter(AccessFormatter('%(client_addr)s - "%(request_line)s" %(status_code)s', use_colors=False)) + saved_level, saved_propagate = logger.level, logger.propagate + logger.addHandler(handler) + logger.setLevel(logging.INFO) + logger.propagate = False + try: + logger.handle(_access_record(full_path)) + finally: + logger.removeHandler(handler) + logger.setLevel(saved_level) + logger.propagate = saved_propagate + return stream.getvalue() + + +def test_uvicorn_access_logger_redacts_a_credential_it_is_handed(): + """Registration happens at litellm import; without it the filter never runs.""" + emitted = _emit_access_line(f"/key/info?key={_LEAKED_KEY}") + + assert _LEAKED_KEY not in emitted + assert "REDACTED" in emitted + + +def test_access_redaction_survives_the_uvicorn_json_log_config(): + """litellm hands uvicorn a dictConfig when json_logs is on. dictConfig clears a + logger's handlers but not its filters, so redaction has to still be attached.""" + import logging.config + + names = ("uvicorn", "uvicorn.error", "uvicorn.access") + saved = tuple((logging.getLogger(n), logging.getLogger(n).handlers[:], logging.getLogger(n).level) for n in names) + try: + logging.config.dictConfig(_get_uvicorn_json_log_config()) + emitted = _emit_access_line(f"/key/info?key={_LEAKED_KEY}") + + assert _LEAKED_KEY not in emitted + assert "REDACTED" in emitted + finally: + for lg, handlers, level in saved: + lg.handlers[:] = handlers + lg.setLevel(level) + lg.propagate = True diff --git a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py new file mode 100644 index 00000000000..1ecd9490f78 --- /dev/null +++ b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py @@ -0,0 +1,107 @@ +import json +from pathlib import Path + +import pytest + +import litellm +from litellm.cost_calculator import cost_per_token +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import StandardBuiltInToolCostTracking + +MUSE_SPARK_STANDARD = "meta/muse-spark-1.3" +MUSE_SPARK_CONTRIBUTOR = "meta/muse-spark-1.3-contributor" +WEB_SEARCH_COST_PER_QUERY = 0.0025 + +PRICING = ( + (MUSE_SPARK_STANDARD, 1.25e-06, 1.5e-07, 4.25e-06), + (MUSE_SPARK_CONTRIBUTOR, 1e-07, 2e-09, 2e-07), +) + + +def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> dict: + with open(Path(__file__).parents[2] / filename) as f: + return json.load(f) + + + +@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) +def test_muse_spark_1_3_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float): + info = _load_cost_map().get(model) + assert info is not None, f"{model} not found in model_prices_and_context_window.json" + + assert info["litellm_provider"] == "meta" + assert info["mode"] == "chat" + + assert info["input_cost_per_token"] == input_cost + assert info["output_cost_per_token"] == output_cost + assert info["cache_read_input_token_cost"] == cached_cost + + assert info["max_input_tokens"] == 1048576 + assert info["max_output_tokens"] == 131072 + assert info["max_tokens"] == 131072 + + assert info["supports_function_calling"] is True + assert info["supports_parallel_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_response_schema"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + assert info["supports_pdf_input"] is True + assert info["supports_web_search"] is True + assert info["supports_minimal_reasoning_effort"] is True + assert info["supports_xhigh_reasoning_effort"] is True + + assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"] + assert info["supported_modalities"] == ["text", "image", "video"] + assert info["supported_output_modalities"] == ["text"] + + assert info["search_context_cost_per_query"] == { + "search_context_size_high": WEB_SEARCH_COST_PER_QUERY, + "search_context_size_low": WEB_SEARCH_COST_PER_QUERY, + "search_context_size_medium": WEB_SEARCH_COST_PER_QUERY, + } + + +@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) +def test_muse_spark_1_3_cost_per_token( + local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float +): + prompt_cost, completion_cost = cost_per_token(model=model, prompt_tokens=1000, completion_tokens=500) + + assert prompt_cost == pytest.approx(1000 * input_cost) + assert completion_cost == pytest.approx(500 * output_cost) + + +@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) +def test_muse_spark_1_3_routes_to_meta_model_api(model: str): + routed_model, provider, _, api_base = get_llm_provider(model=model, api_key="sk-test") + + assert routed_model == model.split("/", 1)[1] + assert provider == "meta" + assert api_base == "https://api.meta.ai/v1" + + +@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) +def test_muse_spark_1_3_web_search_cost_per_query(local_model_cost_map, model: str): + info = litellm.get_model_info(model=model) + + assert StandardBuiltInToolCostTracking.get_cost_for_web_search(model_info=info) == WEB_SEARCH_COST_PER_QUERY + + +@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) +def test_muse_spark_1_3_backup_matches_main(model: str): + """Ensure the bundled model cost map stays in sync with the canonical file.""" + main_cost = _load_cost_map() + backup_cost = _load_cost_map("litellm/model_prices_and_context_window_backup.json") + + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" + + +def test_muse_spark_contributor_tier_is_cheaper_than_standard(): + cost_map = _load_cost_map() + standard = cost_map[MUSE_SPARK_STANDARD] + contributor = cost_map[MUSE_SPARK_CONTRIBUTOR] + + for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): + assert contributor[field] < standard[field], f"contributor {field} should undercut the standard tier" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index ff5db18a3ea..dcaa6cfd602 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -6,6 +6,7 @@ import logging import os import threading from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -1537,6 +1538,91 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" +@pytest.mark.asyncio +async def test_ageneric_api_call_resolves_realtime_session_model(): + """ + Regression for #36742: realtime client secret requests carry the model inside `session` too, and the proxy + fills it with the pre-routing model group name. The underlying litellm function reads session.model first, + so it must see the resolved deployment, while a caller's nested transcription model stays untouched. + """ + routed: Final = AsyncMock(return_value={"result": "ok"}) + + router = litellm.Router( + model_list=[ + { + "model_name": "my-realtime-group", + "litellm_params": { + "model": "openai/gpt-realtime-2.1-mini", + "api_key": "fake-key", + }, + "model_info": {"mode": "realtime"}, + } + ] + ) + + await router._ageneric_api_call_with_fallbacks( + model="my-realtime-group", + original_function=routed, + session={ + "type": "realtime", + "model": "my-realtime-group", + "audio": {"input": {"transcription": {"model": "gpt-4o-transcribe"}}}, + }, + ) + + sent: Final = routed.call_args.kwargs + assert sent["model"] == "openai/gpt-realtime-2.1-mini" + assert sent["session"]["model"] == "openai/gpt-realtime-2.1-mini" + assert sent["session"]["audio"]["input"]["transcription"]["model"] == "gpt-4o-transcribe" + + +@pytest.mark.asyncio +async def test_ageneric_api_call_does_not_add_session_model(): + """ + A session that never carried a model must not gain one from routing: the underlying function then falls back + to the resolved `model` kwarg itself, and the outgoing session body keeps the caller's shape. + """ + routed: Final = AsyncMock(return_value={"result": "ok"}) + + router = litellm.Router( + model_list=[ + { + "model_name": "my-realtime-group", + "litellm_params": { + "model": "openai/gpt-realtime-2.1-mini", + "api_key": "fake-key", + }, + "model_info": {"mode": "realtime"}, + } + ] + ) + + await router._ageneric_api_call_with_fallbacks( + model="my-realtime-group", + original_function=routed, + session={"type": "realtime"}, + ) + + sent: Final = routed.call_args.kwargs + assert sent["model"] == "openai/gpt-realtime-2.1-mini" + assert sent["session"] == {"type": "realtime"} + + +@pytest.mark.parametrize( + "session, expected", + [ + ({"type": "realtime", "model": "my-realtime-group"}, {"session": {"type": "realtime", "model": "resolved"}}), + ({"type": "realtime"}, {}), + (None, {}), + ("not-a-session", {}), + ], +) +def test_with_router_resolved_session_model(session, expected): + from litellm.router import _with_router_resolved_session_model + + assert dict(_with_router_resolved_session_model(session, "resolved")) == expected + + def test_router_get_model_access_groups_team_only_models(): """ Test that Router.get_model_access_groups returns the correct response for team-only models @@ -11907,3 +11993,55 @@ class TestPreRoutingTierDrivesFallbacks: response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-b" + + +@pytest.mark.asyncio +async def test_prompt_management_factory_marks_injection_for_every_deployment(monkeypatch): + """The factory stamps a provisional deployment's model_info into kwargs before the + prompt pass runs, then routes on the returned model, so any deployment can end up + billed. An injection recorded there must carry the every-deployment sentinel, never + the provisional deployment's id, or a differently-billed deployment loses the credit.""" + import time + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + router = litellm.Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + "model_info": {"id": "provisional-dep"}, + } + ] + ) + captured: dict = {} + + async def _capture_acompletion(**kwargs): + captured.update(kwargs) + return litellm.ModelResponse() + + monkeypatch.setattr(litellm, "acompletion", _capture_acompletion) + logging_obj = LiteLLMLogging( + model="cached-claude", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="lit-6445", + function_id="f", + ) + await router.acompletion( + model="cached-claude", + messages=[ + {"role": "system", "content": "a static system prompt"}, + {"role": "user", "content": "hi"}, + ], + cache_control_injection_points=[{"location": "message", "role": "system"}], + litellm_logging_obj=logging_obj, + ) + bucket = captured.get("litellm_metadata") or captured["metadata"] + assert captured["model_info"]["id"] == "provisional-dep" + assert bucket["litellm_gateway_injected_cache"] == "" diff --git a/tests/test_litellm/test_secret_redaction.py b/tests/test_litellm/test_secret_redaction.py index 7d50694c805..303efcab9c7 100644 --- a/tests/test_litellm/test_secret_redaction.py +++ b/tests/test_litellm/test_secret_redaction.py @@ -1,8 +1,10 @@ import logging import logging.config import sys +import time from collections.abc import Callable from io import StringIO +from typing import Final from unittest.mock import patch import pytest @@ -67,6 +69,50 @@ def test_redact_string_catches_secret_patterns(): assert redact_string(normal) == normal +@pytest.mark.parametrize( + "connection_string", + [ + "postgres://admin:pass3cret@db.example.com:5432/mydb", + "redis://:pass3cret@cache.example.com:6379", + "postgres://admin:pass/s3cret@db.example.com:5432/mydb", + "amqp://admin:pass:s3cret@rabbit:5672", + "https://ad@min:pass3cret@host", + # An unencoded "@" inside the password, with a ":" after it + "postgresql://admin:p@ss3cret:2026@db.example.com:5432/mydb", + "amqp://guest:gu@st3cret:1@rabbit:5672/", + # An AWS RDS IAM auth token is a presigned query string used as the + # password, so the userinfo runs to several hundred characters. + "postgresql://litellm:host%3A5432%2F%3FAction%3Dconnect%26X-Amz-Signature%3D" + + "f" * 540 + + "s3cret@db.host:5432/litellm", + ], +) +def test_redact_string_still_catches_connection_string_credentials(connection_string): + """The bounded userinfo pattern must keep matching real connection strings.""" + assert "s3cret" not in redact_string(connection_string) + + +def _redaction_cost(url_bytes: int) -> float: + url: Final = "/x?u=" + "a://" * (url_bytes // 4) + + def once() -> float: + started = time.perf_counter() + redact_string(url) + return time.perf_counter() - started + + return min(once() for _ in range(3)) + + +def test_redact_string_stays_sub_quadratic_on_a_long_adversarial_url(): + """Access-log redaction runs on attacker-controlled request lines, so quadrupling + a URL of scheme separators must not multiply the cost by sixteen. Comparing two + sizes rather than asserting a wall-clock ceiling keeps this honest on a slow box: + the unbounded pattern this replaced cost 5s at 4 KB and 314s at 16 KB.""" + growth: Final = _redaction_cost(16 * 1024) / _redaction_cost(4 * 1024) + + assert growth < 11.0, f"cost grew {growth:.1f}x for 4x the URL length" + + def test_redact_string_catches_minimum_length_virtual_key(): """Regression test for LIT-4355: keys at the enforced 16-char minimum (MINIMUM_CUSTOM_KEY_LENGTH) must be treated as key-shaped by the scrubber.""" diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 0790b41c349..200cfd02197 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -41,7 +41,6 @@ from litellm.utils import ( _snapshot_exception_for_hook, async_post_call_failure_deployment_hook, client, - get_api_key, get_llm_provider, get_non_default_completion_params, get_optional_params_image_gen, @@ -4917,17 +4916,6 @@ def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkey _invalidate_model_cost_lowercase_map() -def test_ai21_api_key_is_resolved_from_the_documented_env_var(monkeypatch: pytest.MonkeyPatch) -> None: - """The ai21 branch resolved a misspelled env var, so the name every other ai21 code path - reads, and the only name documented, was ignored.""" - monkeypatch.setattr(litellm, "api_key", None) - monkeypatch.setattr(litellm, "ai21_key", None) - monkeypatch.delenv("AI211_API_KEY", raising=False) - monkeypatch.setenv("AI21_API_KEY", "sk-ai21-resolved-from-env") - - assert get_api_key(llm_provider="ai21", dynamic_api_key=None) == "sk-ai21-resolved-from-env" - - class _JsonCapture(logging.Handler): def __init__(self): super().__init__() diff --git a/tests/test_rust_python_harness.py b/tests/test_rust_python_harness.py new file mode 100644 index 00000000000..90dc38663dd --- /dev/null +++ b/tests/test_rust_python_harness.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +import importlib +import json +from pathlib import Path + +import pytest + +catalog = importlib.import_module("tests.rust-python-harness.catalog") +cli = importlib.import_module("tests.rust-python-harness.cli") +models = importlib.import_module("tests.rust-python-harness.models") +runner = importlib.import_module("tests.rust-python-harness.runner") +ui = importlib.import_module("tests.rust-python-harness.ui") + +load_catalog = catalog.load_catalog +_pick_values = cli._pick_values +_coverage_pytest_args = cli._coverage_pytest_args +_select = cli._select +CaseResult = models.CaseResult +Coverage = models.Coverage +HarnessCase = models.HarnessCase +HarnessRun = models.HarnessRun +RunStatus = models.RunStatus +SDK_FUNCTIONS = models.SDK_FUNCTIONS +section_confidence = models.section_confidence +run_pytest = runner.run_pytest +runnable_selectors = runner.runnable_selectors +selector_matches_node = runner.selector_matches_node +_format_duration = ui._format_duration +_rerun_command = ui._rerun_command +_summary = ui._summary + + +def _case( + *, selectors: tuple[str, ...] = (), coverage: Coverage = Coverage.COMPLETE +) -> HarnessCase: + return HarnessCase( + strategy_id="example", + strategy_label="Example", + sdk_function="messages", + coverage=coverage, + selectors=selectors, + ) + + +def _manifest() -> dict[str, object]: + return { + "order": 1, + "id": "example", + "label": "Example strategy", + "description": "Example description", + "functions": { + function: {"coverage": "planned", "selectors": []} + for function in SDK_FUNCTIONS + }, + } + + +def test_should_load_the_three_harness_strategies_in_order() -> None: + strategies = load_catalog() + + assert [strategy.id for strategy in strategies] == [ + "e2e_fuzz_tests", + "unit_tests_rust", + "validate_sub_methods", + ] + assert all( + tuple(case.sdk_function for case in strategy.cases) == SDK_FUNCTIONS + for strategy in strategies + ) + + +def test_should_reject_a_manifest_missing_an_sdk_function(tmp_path: Path) -> None: + strategy_directory = tmp_path / "example" + strategy_directory.mkdir() + manifest = _manifest() + del manifest["functions"]["count_tokens"] # type: ignore[index] + (strategy_directory / "strategy.json").write_text( + json.dumps(manifest), encoding="utf-8" + ) + + with pytest.raises(ValueError, match="functions must exactly match"): + load_catalog(tmp_path) + + +@pytest.mark.parametrize( + ("selector", "nodeid", "matches"), + [ + ("tests/test_parity.py", "tests/test_parity.py::test_one", True), + ("tests/test_parity.py::test_one", "tests/test_parity.py::test_one", True), + ( + "tests/test_parity.py::test_one", + "tests/test_parity.py::test_one[value]", + True, + ), + ("tests/test_parity.py::test_one", "tests/test_parity.py::test_two", False), + ], +) +def test_should_match_pytest_file_and_node_selectors( + selector: str, nodeid: str, matches: bool +) -> None: + assert selector_matches_node(selector, nodeid) is matches + + +def test_should_only_return_selectors_whose_files_exist(tmp_path: Path) -> None: + existing = tmp_path / "tests" / "test_parity.py" + existing.parent.mkdir() + existing.write_text("", encoding="utf-8") + case = _case( + selectors=("tests/test_parity.py", "tests/test_missing.py::test_missing") + ) + + assert runnable_selectors((case,), tmp_path) == ("tests/test_parity.py",) + + +def test_should_mark_planned_and_not_applicable_cases_without_running() -> None: + planned = CaseResult(case=_case(coverage=Coverage.PLANNED)) + not_applicable = CaseResult(case=_case(coverage=Coverage.NOT_APPLICABLE)) + + planned.set_initial_status() + not_applicable.set_initial_status() + + assert planned.status is RunStatus.PLANNED + assert not_applicable.status is RunStatus.NOT_APPLICABLE + + +def test_should_treat_an_all_planned_filtered_run_as_success(tmp_path: Path) -> None: + exit_code, run = run_pytest( + cases=(_case(coverage=Coverage.PLANNED),), + repo_root=tmp_path, + on_update=lambda _: None, + ) + + assert exit_code == 0 + assert next(iter(run.results.values())).status is RunStatus.PLANNED + + +def test_should_finalize_a_fully_passing_case() -> None: + result = CaseResult(case=_case(selectors=("tests/test_parity.py",))) + result.set_initial_status() + result.collected.update({"one", "two"}) + result.completed.update({"one", "two"}) + result.passed = 2 + + result.finalize() + + assert result.status is RunStatus.PASSED + + +def test_should_replace_a_pass_with_a_teardown_error() -> None: + result = CaseResult(case=_case(selectors=("tests/test_parity.py",))) + result.set_initial_status() + result.collected.add("one") + + result.record("one", RunStatus.PASSED, 0.1) + result.record("one", RunStatus.ERROR, 0.2) + + assert result.status is RunStatus.ERROR + assert result.passed == 0 + assert result.errors == 1 + assert result.duration == pytest.approx(0.3) + + +def test_should_filter_the_catalog_by_strategy_and_sdk_function() -> None: + strategies = load_catalog() + + cases = _select(strategies, {"e2e_fuzz_tests"}, {"messages"}) + + assert len(cases) == 1 + assert cases[0].key == "e2e_fuzz_tests:messages" + + +def test_should_reject_an_unknown_strategy() -> None: + with pytest.raises(ValueError, match="Unknown strategy"): + _select(load_catalog(), {"not-real"}, set()) + + +def test_should_pick_multiple_interactive_filters() -> None: + answers = iter(["nope", "1, 3"]) + + selected = _pick_values( + "Examples", + (("one", "One"), ("two", "Two"), ("three", "Three")), + input_fn=lambda _: next(answers), + ) + + assert selected == {"one", "three"} + + +def test_should_format_developer_facing_run_context() -> None: + run = HarnessRun.from_cases((_case(selectors=("tests/test_parity.py",)),)) + result = next(iter(run.results.values())) + result.collected.add("tests/test_parity.py::test_one") + result.record("tests/test_parity.py::test_one", RunStatus.PASSED, 1.25) + + assert _summary(run) == (1, 0, 0, 0) + assert _format_duration(1.25) == "1.2s" + assert _rerun_command("tests/test_parity.py::test_one") == ( + "poetry run pytest tests/test_parity.py::test_one -q" + ) + assert _rerun_command("tests/test_parity.py::test_one[value with spaces]") == ( + "poetry run pytest 'tests/test_parity.py::test_one[value with spaces]' -q" + ) + + +def test_should_build_python_coverage_reports_below_the_target_directory( + tmp_path: Path, +) -> None: + args = _coverage_pytest_args(tmp_path) + + assert tmp_path.is_dir() + assert "--cov=litellm" in args + assert "--cov-context=test" in args + assert f"--cov-report=json:{tmp_path / 'python.json'}" in args + assert f"--cov-report=xml:{tmp_path / 'python.xml'}" in args + assert f"--cov-report=html:{tmp_path / 'python-html'}" in args + + +def test_should_report_confidence_for_each_sdk_section() -> None: + strategies = load_catalog() + cases = tuple(case for strategy in strategies for case in strategy.cases) + run = HarnessRun.from_cases(cases) + passing = run.results["e2e_fuzz_tests:responses"] + passing.collected.add("tests/test_parity.py::test_one") + passing.record("tests/test_parity.py::test_one", RunStatus.PASSED) + + scores = { + score.sdk_function: score for score in section_confidence(run, strategies) + } + + assert scores["responses"].verified_strategies == 1 + assert scores["responses"].required_strategies == 3 + assert scores["responses"].percentage == 33 + assert scores["responses"].level.value == "MEDIUM" + assert scores["count_tokens"].percentage == 0 + assert scores["count_tokens"].level.value == "LOW" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 2836b22c3d0..78595d5c0b6 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22364 + "limit": 22358 }, "LIT002": { - "limit": 26776 + "limit": 26773 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16506 + "limit": 16493 }, "LIT011": { "limit": 5535 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpToolSearchSettings/useMCPToolSearchSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpToolSearchSettings/useMCPToolSearchSettings.ts new file mode 100644 index 00000000000..8645e867c16 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpToolSearchSettings/useMCPToolSearchSettings.ts @@ -0,0 +1,47 @@ +import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; +import { apiClient } from "@/components/networking"; +import type { components } from "@/lib/http/schema"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +export type MCPToolSearchSettings = components["schemas"]["MCPToolSearchSettings"]; +export type MCPToolSearchSettingsResponse = components["schemas"]["MCPToolSearchSettingsResponse"]; + +const GET_PATH = "/get/mcp_tool_search_settings"; +const UPDATE_PATH = "/update/mcp_tool_search_settings"; + +const mcpToolSearchSettingsKeys = createQueryKeys("mcpToolSearchSettings"); + +export const getMCPToolSearchSettings = (accessToken: string): Promise => + apiClient.get(GET_PATH, { accessToken }); + +export const updateMCPToolSearchSettings = ( + accessToken: string, + settings: MCPToolSearchSettings, +): Promise => + apiClient.patch(UPDATE_PATH, { accessToken, body: settings }); + +export const useMCPToolSearchSettings = () => { + const { accessToken } = useAuthorized(); + return useQuery({ + queryKey: mcpToolSearchSettingsKeys.list({}), + queryFn: () => getMCPToolSearchSettings(accessToken), + enabled: !!accessToken, + }); +}; + +export const useUpdateMCPToolSearchSettings = () => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + return useMutation({ + mutationFn: (settings) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return updateMCPToolSearchSettings(accessToken, settings); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: mcpToolSearchSettingsKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx index 6e70b1ac59a..e6148d5d997 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx @@ -36,6 +36,7 @@ import type { Team, } from "@/components/mcp_tools/types"; import MCPSemanticFilterSettings from "@/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings"; +import MCPToolSearchSettings from "@/components/Settings/AdminSettings/MCPToolSearchSettings/MCPToolSearchSettings"; import MCPNetworkSettings from "./MCPNetworkSettings"; import MCPDiscovery from "./mcp_discovery"; import { ByokCredentialModal } from "@/components/mcp_tools/ByokCredentialModal"; @@ -544,6 +545,11 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) Semantic Filter )} + {isAdminRole(userRole) && ( + + Tool Search + + )} {isAdminRole(userRole) && ( Network Settings @@ -726,6 +732,11 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) )} + {isAdminRole(userRole) && ( + + + + )} {isAdminRole(userRole) && ( diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/MCPToolSearchSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/MCPToolSearchSettings.test.tsx new file mode 100644 index 00000000000..6e312bea651 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/MCPToolSearchSettings.test.tsx @@ -0,0 +1,96 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { render, screen, act, fireEvent } from "@testing-library/react"; +import MCPToolSearchSettings from "./MCPToolSearchSettings"; +import { + useMCPToolSearchSettings, + useUpdateMCPToolSearchSettings, +} from "@/app/(dashboard)/hooks/mcpToolSearchSettings/useMCPToolSearchSettings"; + +vi.mock("@/app/(dashboard)/hooks/mcpToolSearchSettings/useMCPToolSearchSettings", () => ({ + useMCPToolSearchSettings: vi.fn(), + useUpdateMCPToolSearchSettings: vi.fn(), +})); + +vi.mock("@/components/llm_calls/fetch_models", () => ({ + fetchAvailableModels: vi.fn().mockResolvedValue([{ model_group: "text-embedding-3-small", mode: "embedding" }]), +})); + +vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), fromError: vi.fn() } })); + +const mockMutate = vi.fn(); + +const EDITED_PAYLOAD = { + embedding_model: "text-embedding-3-small", + top_k: 8, + similarity_threshold: 0.25, + core_tools: ["treasury-get_rates", "weather-forecast"], +}; + +const STORED = { + field_schema: {}, + values: { + embedding_model: "text-embedding-3-small", + top_k: 3, + similarity_threshold: 0.25, + core_tools: ["treasury-get_rates"], + }, +}; + +type SettingsQuery = ReturnType; +type SettingsMutation = ReturnType; + +const settled = (data: typeof STORED | undefined, overrides: Partial = {}) => + ({ data, isLoading: false, isError: false, error: null, ...overrides }) as SettingsQuery; + +async function renderSettings(accessToken: string | null = "token") { + const result = render(); + await act(async () => {}); + return result; +} + +describe("MCPToolSearchSettings", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(useMCPToolSearchSettings).mockReturnValue(settled(STORED)); + vi.mocked(useUpdateMCPToolSearchSettings).mockReturnValue({ + mutate: mockMutate, + isPending: false, + } as unknown as SettingsMutation); + }); + + it("shows the stored settings and keeps Save disabled until something changes", async () => { + await renderSettings(); + + expect(screen.getByLabelText(/top k results/i)).toHaveValue(3); + expect(screen.getByLabelText(/always returned first/i)).toHaveValue("treasury-get_rates"); + expect(screen.getByRole("slider", { hidden: true })).toHaveAttribute("aria-valuenow", "0.25"); + expect(screen.getByRole("button", { name: /save settings/i })).toBeDisabled(); + }); + + it("sends the edited settings as the proxy's PATCH payload", async () => { + await renderSettings(); + + fireEvent.change(screen.getByLabelText(/top k results/i), { target: { value: "8" } }); + fireEvent.change(screen.getByLabelText(/always returned first/i), { + target: { value: "treasury-get_rates\nweather-forecast" }, + }); + await act(async () => { + fireEvent.click(screen.getByRole("button", { name: /save settings/i })); + }); + + expect(mockMutate).toHaveBeenCalledTimes(1); + expect(mockMutate.mock.calls[0][0]).toEqual(EDITED_PAYLOAD); + }); + + it("asks the user to log in without a token and surfaces load errors", async () => { + await renderSettings(null); + expect(screen.getByText(/please log in/i)).toBeInTheDocument(); + + vi.mocked(useMCPToolSearchSettings).mockReturnValue( + settled(undefined, { isError: true, error: new Error("Database not connected") }), + ); + await renderSettings(); + expect(screen.getByText("Database not connected")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/MCPToolSearchSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/MCPToolSearchSettings.tsx new file mode 100644 index 00000000000..2ad9a56d0f3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPToolSearchSettings/MCPToolSearchSettings.tsx @@ -0,0 +1,251 @@ +"use client"; + +import { + useMCPToolSearchSettings, + useUpdateMCPToolSearchSettings, +} from "@/app/(dashboard)/hooks/mcpToolSearchSettings/useMCPToolSearchSettings"; +import { toast } from "@/lib/toast"; +import { Skeleton } from "@/components/ui/skeleton"; +import { CircleHelp, Info, Save } from "lucide-react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { useEffect, useState } from "react"; +import { useForm } from "react-hook-form"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; +import { FieldGroup } from "@/components/ui/field"; +import { FormField } from "@/components/shared/form/FormField"; +import { SearchSelect } from "@/components/shared/SearchSelect"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Input } from "@/components/ui/input"; +import { Slider } from "@/components/ui/slider"; +import { Textarea } from "@/components/ui/textarea"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { + DEFAULT_FORM_VALUES, + TOP_K_MAX, + TOP_K_MIN, + clampTopK, + formToPayload, + storedValuesToForm, + ToolSearchFormValues, +} from "./toolSearchForm"; + +interface MCPToolSearchSettingsProps { + accessToken: string | null; +} + +const SIMILARITY_THRESHOLD_MARKS = [0, 0.3, 0.5, 0.7, 1]; + +const labelWithHint = (label: string, hint: string): React.ReactNode => ( + <> + {label} + + } /> + {hint} + + +); + +export default function MCPToolSearchSettings({ accessToken }: MCPToolSearchSettingsProps) { + const { data, isLoading, isError, error } = useMCPToolSearchSettings(); + const { mutate: updateSettings, isPending: isUpdating } = useUpdateMCPToolSearchSettings(); + const form = useForm({ defaultValues: DEFAULT_FORM_VALUES }); + const isDirty = form.formState.isDirty; + const [embeddingModels, setEmbeddingModels] = useState([]); + const [loadingModels, setLoadingModels] = useState(true); + const storedValues = data?.values; + + useEffect(() => { + if (!accessToken) return; + fetchAvailableModels(accessToken) + .then((models) => setEmbeddingModels(models.filter((model) => model.mode === "embedding"))) + .catch((fetchError: unknown) => console.error("Error fetching embedding models:", fetchError)) + .finally(() => setLoadingModels(false)); + }, [accessToken]); + + useEffect(() => { + if (!storedValues) return; + form.reset(storedValuesToForm(storedValues)); + }, [storedValues, form]); + + const handleSave = (formValues: ToolSearchFormValues) => { + updateSettings(formToPayload(formValues), { + onSuccess: () => { + form.reset(formValues); + toast.success("Settings updated successfully. Changes will be applied across all pods within 10 seconds."); + }, + onError: (saveError) => toast.fromError(saveError), + }); + }; + + if (!accessToken) { + return
Please log in to configure tool search.
; + } + + if (isLoading) { + return ( +
+ + + +
+ ); + } + + if (isError) { + return ( + + Could not load MCP tool search settings + {error instanceof Error && {error.message}} + + ); + } + + return ( +
+ + + Native MCP Tool Search + + Controls the mcp_tool_search virtual tool that native MCP clients call to discover tools. With an + embedding model set, tools are ranked by the meaning of their name and description, so a query like + "FX" finds a "foreign exchange rates" tool. Without one, keyword matching is used. Callers + only ever see tools their key, team and server permissions already allow. + + + + +
event.preventDefault()} noValidate> + + + Ranking + + + + + {({ value, onChange, id }) => ( + ({ label: model.model_group, value: model.model_group }))} + value={value} + onValueChange={onChange} + allowClear + placeholder={loadingModels ? "Loading models..." : "Keyword matching (no embedding model)"} + emptyText={loadingModels ? "Loading..." : "No embedding models available"} + disabled={isUpdating || loadingModels} + /> + )} + + + + {({ ref, value, onChange, onBlur, id }) => ( + onChange(event.target.valueAsNumber)} + onBlur={() => { + onChange(Number.isNaN(value) ? DEFAULT_FORM_VALUES.top_k : clampTopK(value)); + onBlur(); + }} + disabled={isUpdating} + /> + )} + + + + {({ value, onChange, id }) => ( +
+ onChange(Array.isArray(next) ? next[0] : next)} + disabled={isUpdating} + /> +
+ {SIMILARITY_THRESHOLD_MARKS.map((mark) => ( + + {mark.toFixed(1)} + + ))} +
+
+ )} + + + + + + + + Core Tools + + + + + {({ ref, value, onChange, onBlur, id }) => ( +