scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

merge: origin/litellm_internal_staging into litellm_lite_login_prefill_code

Browse source
This commit is contained in:
mateo-berri 2026-09-02 16:38:30 -07:00
commit 4ebac321a1
119 changed files with 7984 additions and 718 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
basedpyright-code-budget.json
View file

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

4
enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py
View file

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

46
litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs
View file

@ -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<Map<String, Value>>,
) -> Result<BTreeMap<String, String>, 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<String, String>, 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)")
}
}

84
litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs
View file

@ -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<Value, Error> {
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<String, Value>,
) -> Result<ProviderAudioTranscriptionRequest, Error> {
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,
&region,
&credentials,
SystemTime::now(),
)?);
}
}
Ok(ProviderAudioTranscriptionRequest {
upstream_headers: headers.into_iter().collect(),
..request.clone()
})
}
pub(super) fn environment_lookup(key: &str) -> Option<String> {
std::env::var(key).ok()
}

82
litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs
View file

@ -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<ProviderAudioTranscriptionRequest, Error> {
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::<Vec<_>>();
if matches!(auth, AudioTranscriptionAuth::Bearer)
&& !has_header(
&upstream_headers
.iter()
.cloned()
.collect::<std::collections::BTreeMap<_, _>>(),
"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(

4
litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs
View file

@ -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<Value, Error> {

11
litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs
View file

@ -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<Duration>,
}

14
litellm-rust/crates/core/src/audio_transcription/client.rs Normal file
View file

@ -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<reqwest::Client> = OnceLock::new();
CLIENT.get_or_init(|| {
reqwest::Client::builder()
.timeout(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS))
.build()
.unwrap_or_else(|_| reqwest::Client::new())
})
}

91
litellm-rust/crates/core/src/audio_transcription/handler.rs Normal file
View file

@ -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<Value, Error> {
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<Vec<(String, String)>, 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<String, String> = 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<Vec<(String, String)>, 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()),
}
}

18
litellm-rust/crates/core/src/audio_transcription/mod.rs
View file

@ -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<Value, Error> {
execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?)
.await
}
#[cfg(test)]
mod tests;

72
litellm-rust/crates/core/src/audio_transcription/prepare.rs Normal file
View file

@ -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<ProviderAudioTranscriptionRequest, Error> {
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,
})
}

50
litellm-rust/crates/core/src/audio_transcription/tests.rs Normal file
View file

@ -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");
}

53
litellm-rust/crates/core/src/audio_transcription/types.rs
View file

@ -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<Map<String, Value>>,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
}
#[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<String, Value>,
pub(super) timeout: Option<Duration>,
}
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 {

2
litellm-rust/crates/core/src/constants.rs
View file

@ -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";

14
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -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)?;

3
litellm/__init__.py
View file

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

77
litellm/_logging.py
View file

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

1
litellm/constants.py
View file

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

32
litellm/integrations/anthropic_cache_control_hook.py
View file

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

60
litellm/integrations/otel/langfuse_logger.py Normal file
View file

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

27
litellm/integrations/otel/logger.py
View file

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

8
litellm/integrations/otel/mappers/langfuse.py
View file

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

90
litellm/integrations/otel/model/request_io.py Normal file
View file

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

18
litellm/litellm_core_utils/litellm_logging.py
View file

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

17
litellm/litellm_core_utils/secret_redaction.py
View file

@ -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.<provider>_key=<value>
@ -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"[^{}]*(?:\{[^{}]*\}[^{}]*)*\}',
]

27
litellm/llms/base_llm/files/azure_blob_storage_backend.py
View file

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

1
litellm/main.py
View file

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

82
litellm/model_prices_and_context_window_backup.json
View file

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

5
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

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

32
litellm/proxy/_experimental/mcp_server/bridge_token_flow.py
View file

@ -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 _:

27
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
View file

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

150
litellm/proxy/_experimental/mcp_server/tool_search.py
View file

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

42
litellm/proxy/_types.py
View file

@ -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=(

226
litellm/proxy/agent_endpoints/agent_registry.py
View file

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

132
litellm/proxy/agent_endpoints/agent_search.py
View file

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

57
litellm/proxy/agent_endpoints/endpoints.py
View file

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

18
litellm/proxy/auth/auth_checks.py
View file

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

4
litellm/proxy/auth/auth_exception_handler.py
View file

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

60
litellm/proxy/auth/auth_utils.py
View file

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

33
litellm/proxy/auth/login_utils.py
View file

@ -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 <master_key>), 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.",

92
litellm/proxy/auth/password_policy.py Normal file
View file

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

142
litellm/proxy/common_utils/semantic_text_index.py Normal file
View file

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

105
litellm/proxy/db/exception_handler.py
View file

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

4
litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py
View file

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

2
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

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

39
litellm/proxy/guardrails/guardrail_initializers.py
View file

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

159
litellm/proxy/guardrails/guardrail_registry.py
View file

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

14
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

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

15
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -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"
```

4
litellm/proxy/management_endpoints/ui_sso.py
View file

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

30
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

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

150
litellm/proxy/openai_files_endpoints/general_upload_validation.py Normal file
View file

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

77
litellm/proxy/proxy_server.py
View file

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

13
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -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"
```

82
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -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.<key> 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

1
litellm/proxy/utils.py
View file

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

26
litellm/router.py
View file

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

29
litellm/types/mcp.py
View file

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

43
litellm/utils.py
View file

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

82
model_prices_and_context_window.json
View file

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

6
ruff-strict-budget.json
View file

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

2
tests/code_coverage_tests/recursive_detector.py
View file

@ -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.
]

1
tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
View file

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

145
tests/rust-python-harness/README.md Normal file
View file

@ -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.

5
tests/rust-python-harness/__init__.py Normal file
View file

@ -0,0 +1,5 @@
"""Interactive Rust/Python SDK parity test harness."""
from .catalog import load_catalog
__all__ = ["load_catalog"]

4
tests/rust-python-harness/__main__.py Normal file
View file

@ -0,0 +1,4 @@
from .cli import main
if __name__ == "__main__":
raise SystemExit(main())

93
tests/rust-python-harness/catalog.py Normal file
View file

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

180
tests/rust-python-harness/cli.py Normal file
View file

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

3
tests/rust-python-harness/e2e_fuzz_tests/README.md Normal file
View file

@ -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.

12
tests/rust-python-harness/e2e_fuzz_tests/strategy.json Normal file
View file

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

219
tests/rust-python-harness/models.py Normal file
View file

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

160
tests/rust-python-harness/runner.py Normal file
View file

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

295
tests/rust-python-harness/ui.py Normal file
View file

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

3
tests/rust-python-harness/unit_tests_rust/README.md Normal file
View file

@ -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.

12
tests/rust-python-harness/unit_tests_rust/strategy.json Normal file
View file

@ -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": []}
}
}

3
tests/rust-python-harness/validate_sub_methods/README.md Normal file
View file

@ -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.

12
tests/rust-python-harness/validate_sub_methods/strategy.json Normal file
View file

@ -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": []}
}
}

359
tests/test_litellm/integrations/otel/test_langfuse_logger.py Normal file
View file

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

21
tests/test_litellm/integrations/test_anthropic_cache_control_hook.py
View file

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

26
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

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

44
tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py
View file

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

38
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

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

87
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

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

29
tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
View file

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

236
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
View file

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

507
tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py
View file

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

3
tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
View file

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

124
tests/test_litellm/proxy/agent_endpoints/test_endpoints.py
View file

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

61
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

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

91
tests/test_litellm/proxy/auth/test_auth_exception_handler.py
View file

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

156
tests/test_litellm/proxy/auth/test_auth_utils.py
View file

@ -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", "<EntityDescriptor/>")
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:

201
tests/test_litellm/proxy/auth/test_login_utils.py
View file

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

18
tests/test_litellm/proxy/auth/test_onboarding.py
View file

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

136
tests/test_litellm/proxy/auth/test_password_policy.py Normal file
View file

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

121
tests/test_litellm/proxy/db/test_exception_handler.py
View file

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

32
tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py
View file

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

39
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py
View file

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

139
tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
View file

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

85
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

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

111
tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py
View file

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

Some files were not shown because too many files have changed in this diff Show more