Merge remote-tracking branch 'origin/main' into litellm_fix_tpm_window_reset_sibling_counters

This commit is contained in:
Devin AI 2026-09-18 21:38:26 +00:00
commit 16d63eaf88
179 changed files with 17846 additions and 1033 deletions

View file

@ -41,4 +41,5 @@ jobs:
"$RUNNER_TEMP/osv-scanner" scan source \
--config osv-scanner.toml \
-L uv.lock \
-L ui/litellm-dashboard/package-lock.json
-L ui/litellm-dashboard/package-lock.json \
-L vscode-extension/package-lock.json

View file

@ -0,0 +1,65 @@
name: VS Code Extension
permissions:
contents: read
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "vscode-extension/**"
- ".github/workflows/test-vscode-extension.yml"
push:
branches:
- main
paths:
- "vscode-extension/**"
- ".github/workflows/test-vscode-extension.yml"
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
vscode-extension:
runs-on: ubuntu-latest
timeout-minutes: 10
defaults:
run:
working-directory: vscode-extension
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 1
persist-credentials: false
- name: Set up Node.js
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
with:
node-version: "24"
cache: npm
cache-dependency-path: vscode-extension/package-lock.json
- name: Install dependencies
run: npm ci
- name: Typecheck
run: npm run typecheck
- name: Unit tests
run: npm test
- name: Package extension
run: npm run package
- name: Upload VSIX
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: litellm-vscode
path: vscode-extension/*.vsix
if-no-files-found: error

View file

@ -82,6 +82,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/anthropic/",
"/azure/",
"/azure_ai/",
"/azure_speech/",
"/aws/",
"/bedrock/",
"/comprehendmedical",

View file

@ -66,7 +66,7 @@
"/v1/video" "/v1/videos" "/video" "/videos" "/v1/search" "/search"
"/v1/containers" "/containers" "/v1/evals" "/v1/memory" "/queue/chat"
"/v1beta" "/interactions"
"/anthropic" "/azure" "/azure_ai" "/aws" "/bedrock" "/comprehendmedical" "/transcribe" "/cohere" "/gemini" "/google"
"/anthropic" "/azure" "/azure_ai" "/azure_speech" "/aws" "/bedrock" "/comprehendmedical" "/transcribe" "/cohere" "/gemini" "/google"
"/vertex_ai" "/vertex-ai" "/assemblyai" "/eu.assemblyai" "/langfuse" "/vllm"
"/mistral" "/groq" "/voyage" "/cursor" "/milvus" "/openai_passthrough"
"/toolset"

View file

@ -559,6 +559,21 @@ dependencies = [
"vsimd",
]
[[package]]
name = "bit-set"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3"
dependencies = [
"bit-vec",
]
[[package]]
name = "bit-vec"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
[[package]]
name = "bitflags"
version = "2.13.1"
@ -1166,6 +1181,17 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "fancy-regex"
version = "0.19.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d301f5bf187b3c295fce6468d3875037a0bccc5f6b151c63cac2f85babf21912"
dependencies = [
"bit-set",
"regex-automata",
"regex-syntax",
]
[[package]]
name = "fastrand"
version = "2.5.0"
@ -2059,7 +2085,9 @@ dependencies = [
name = "litellm-core-utils"
version = "0.1.0"
dependencies = [
"fancy-regex",
"litellm-types",
"rstest",
"serde",
"serde_json",
"serde_path_to_error",

View file

@ -50,6 +50,7 @@ strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"
fancy-regex = "0.19.2"
veil = "0.3.0"
[profile.release]

View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
fancy-regex.workspace = true
litellm-types.workspace = true
serde.workspace = true
serde_json.workspace = true
@ -13,3 +14,6 @@ serde_path_to_error = "0.1"
serde_with.workspace = true
thiserror.workspace = true
url.workspace = true
[dev-dependencies]
rstest.workspace = true

View file

@ -0,0 +1,115 @@
use super::public::PublicError;
use super::rules::{Rule, contains_any};
/// The text branches of `_map_cohere_exception`, in its order.
pub(super) const RULES: &[Rule] = &[
Rule::new(
|mapping| {
contains_any(
&mapping.error_str,
&["invalid api token", "No API key provided."],
)
},
PublicError::Authentication,
),
Rule::new(
|mapping| mapping.error_str.contains("invalid type: parameter"),
PublicError::BadRequest,
),
Rule::new(
|mapping| mapping.error_str.contains("too many tokens"),
PublicError::ContextWindowExceeded,
),
Rule::new(
|mapping| {
mapping
.error_str
.to_lowercase()
.contains("internal server error")
},
PublicError::InternalServer,
),
Rule::new(
|mapping| mapping.status.is_none() && mapping.error_str.contains("invalid type:"),
PublicError::BadRequest,
),
Rule::new(
|mapping| mapping.status.is_none() && mapping.error_str.contains("Unexpected server error"),
PublicError::InternalServer,
),
];
#[cfg(test)]
mod tests {
use super::super::rules::first_match;
use super::super::testing::mapping;
use super::*;
fn classified(text: &str) -> Option<PublicError> {
classified_with(Some(400), text)
}
fn classified_with(status: Option<u16>, text: &str) -> Option<PublicError> {
first_match(RULES, &mapping(status, text)).map(|rule| rule.error)
}
#[rstest::rstest]
#[case::invalid_token("invalid api token", PublicError::Authentication)]
#[case::no_api_key("No API key provided.", PublicError::Authentication)]
#[case::invalid_parameter("invalid type: parameter x", PublicError::BadRequest)]
#[case::too_many_tokens("too many tokens", PublicError::ContextWindowExceeded)]
#[case::internal_server_text("Internal Server Error", PublicError::InternalServer)]
#[case::internal_server_any_case("INTERNAL server ERROR", PublicError::InternalServer)]
fn each_text_rule_claims_its_marker(#[case] text: &str, #[case] expected: PublicError) {
assert_eq!(classified(text), Some(expected));
}
#[rstest::rstest]
#[case::token_before_parameter(
"invalid api token invalid type: parameter",
PublicError::Authentication
)]
#[case::parameter_before_tokens(
"invalid type: parameter too many tokens",
PublicError::BadRequest
)]
#[case::tokens_before_internal(
"too many tokens Internal Server Error",
PublicError::ContextWindowExceeded
)]
fn the_earlier_rule_wins_when_two_apply(#[case] text: &str, #[case] expected: PublicError) {
assert_eq!(classified(text), Some(expected));
}
#[rstest::rstest]
#[case::invalid_type(None, "invalid type: x", Some(PublicError::BadRequest))]
#[case::unexpected_server_error(
None,
"Unexpected server error",
Some(PublicError::InternalServer)
)]
#[case::invalid_type_before_unexpected(
None,
"invalid type: x Unexpected server error",
Some(PublicError::BadRequest)
)]
#[case::internal_before_invalid_type(
None,
"internal server error invalid type: x",
Some(PublicError::InternalServer)
)]
#[case::invalid_type_with_a_status(Some(500), "invalid type: x", None)]
#[case::unexpected_with_a_status(Some(400), "Unexpected server error", None)]
fn the_trailing_rules_only_claim_failures_without_a_status(
#[case] status: Option<u16>,
#[case] text: &str,
#[case] expected: Option<PublicError>,
) {
assert_eq!(classified_with(status, text), expected);
}
#[test]
fn text_without_a_marker_is_left_to_the_status_table() {
assert_eq!(classified("rejected"), None);
}
}

View file

@ -0,0 +1,542 @@
//! A port of Python's `exception_type` for the routes that run in Rust. Rust decides the
//! public class, the message and the debug text; Python only builds the class.
//!
//! DIVERGENCES: where the Python mapper is inconsistent, the port follows one rule instead.
//! - The message is always `{Provider}Exception - {redacted text}`. Python's per-branch
//! labels (`RateLimitError: `, `litellm.RateLimitError: `, `Vertex_aiException BadRequestError`)
//! are dropped because every public class already prefixes `litellm.{Class}: `.
//! - The upstream response is always the real one. Python swaps in made-up `httpx.Response`
//! stubs on some Vertex branches, losing the body and `retry-after`.
//! - The debug text is always attached; Python passes it on some branches only.
//! - No family rule turns a status into a class; the shared status table owns that. So a
//! Vertex 502 is a `BadGatewayError` and an OpenAI-family 403 is a `PermissionDeniedError`.
//! Three rules read the status only to gate a text match, as Python does: the standalone
//! `429`, Vertex's wrapped 429 behind a 5xx, and Cohere's rules for failures with no status.
//! - A timeout text marker on an HTTP failure keeps the upstream response. Python's `Timeout`
//! carries none.
//! - Every family matches and reports the redacted text. Python's OpenAI mapper builds the
//! message from the unredacted text.
//! - A refused connection is an `APIConnectionError`, not the 500 Python's HTTP handler
//! synthesizes.
//! - Dropped Python rules: Vertex's bare `403` substring (it matches `4031 tokens`), Vertex's
//! `IndexError` quota marker (a Python client crash), the OpenAI SDK's missing-`api_key`
//! text and its `OPENAI` renaming, Cohere's `llm_provider="cohere"` override, and Cohere's
//! `CohereConnectionError` check (a Python SDK class name).
//!
//! KNOWN_GAPS: differences from the Python mapper that no Rust route can reach today. Each
//! one stops being acceptable at its trigger.
//! - The Vertex partner-model API base for "claude" models is not built into the debug text.
//! Trigger: a Vertex route whose models include Anthropic partner models.
//! - The debug text's `API Base` line is only the non-streaming Vertex URL. Python prefers an
//! explicit or provider-resolved `api_base`, uses `:streamGenerateContent` when streaming,
//! and has Gemini and OpenAI defaults. Trigger: the first route wired to this mapper, since
//! every route knows its `api_base`.
//! - The debug text has no `Messages:` line, which Python adds when
//! `redact_messages_in_exceptions` is off. Trigger: a wired route that carries messages.
//! - Python reports the provider `get_llm_provider` resolves for a stripped model name when
//! that name happens to be in the model cost map. Trigger: a route whose model names
//! overlap the cost map; that needs the provider resolution port, not a classifier change.
//! - `litellm_proxy` errors are not unwrapped into the proxied exception. Trigger: a Rust
//! route that calls a LiteLLM proxy.
//! - Only the OpenAI-compatible, Vertex AI and Cohere mappers are ported; every other
//! provider goes straight to the status table. Trigger: a Rust route for such a provider.
use super::secret_redaction::SecretRedactor;
mod cohere;
mod openai;
mod original;
mod public;
mod rules;
mod status;
mod vertex_ai;
pub use original::{ExceptionFamily, OriginalException};
pub use public::{MappedFailure, PublicError, UpstreamResponse};
use rules::{Rule, contains_any, first_match};
const TIMEOUT_MARKERS: &[&str] = &[
"Request Timeout Error",
"Request timed out",
"Timed out generating response",
"The read operation timed out",
];
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ExceptionContext {
pub model: String,
pub custom_llm_provider: String,
pub asynchronous: bool,
pub vertex_project: Option<String>,
pub vertex_location: Option<String>,
pub model_group: Option<String>,
pub deployment: Option<String>,
pub user_api_key_alias: Option<String>,
pub user_api_key_team_alias: Option<String>,
}
/// What the rules read: the status of a provider response, if any, and the redacted text.
struct Mapping {
status: Option<u16>,
error_str: String,
}
pub fn exception_type(
context: &ExceptionContext,
redactor: Option<&SecretRedactor>,
original: &OriginalException,
) -> MappedFailure {
let (status, text, upstream) = match original {
OriginalException::Http {
status,
body,
headers,
} => (
Some(*status),
body.clone(),
Some(UpstreamResponse {
status: *status,
body: body.clone(),
headers: headers.clone(),
}),
),
OriginalException::Connection { message } | OriginalException::Plain { message } => {
(None, message.clone(), None)
}
OriginalException::Timeout {
timeout_seconds,
elapsed_seconds,
} => (
None,
timeout_message(context.asynchronous, *timeout_seconds, *elapsed_seconds),
None,
),
};
let mapping = Mapping {
status,
error_str: match redactor {
Some(redactor) => redactor.redact(&text),
None => text,
},
};
let family = ExceptionFamily::for_provider(&context.custom_llm_provider);
let (error, hint) = classify(family, original, &mapping);
MappedFailure {
error,
message: format!(
"{} - {}{hint}",
exception_provider(&context.custom_llm_provider),
mapping.error_str
),
upstream,
debug_info: extra_information(context, api_base(context).as_deref()),
}
}
fn classify(
family: ExceptionFamily,
original: &OriginalException,
mapping: &Mapping,
) -> (PublicError, &'static str) {
const TIMEOUT: PublicError = PublicError::Timeout { status: 408 };
if matches!(original, OriginalException::Timeout { .. })
|| contains_any(&mapping.error_str, TIMEOUT_MARKERS)
{
return (TIMEOUT, "");
}
if let Some(rule) = first_match(family_rules(family), mapping) {
return (rule.error, rule.hint);
}
let by_status = mapping.status.and_then(status::classify);
(by_status.unwrap_or(PublicError::ApiConnection), "")
}
fn family_rules(family: ExceptionFamily) -> &'static [Rule] {
match family {
ExceptionFamily::OpenAiCompatible => openai::RULES,
ExceptionFamily::VertexAi => vertex_ai::RULES,
ExceptionFamily::Cohere => cohere::RULES,
ExceptionFamily::Other => &[],
}
}
/// The text the Python HTTP handler's timeout carries: the sync and async handlers word it
/// differently.
fn timeout_message(
asynchronous: bool,
timeout_seconds: Option<f64>,
elapsed_seconds: Option<f64>,
) -> String {
let timeout = python_float(timeout_seconds);
if asynchronous {
let elapsed =
python_float(elapsed_seconds.map(|seconds| (seconds * 1000.0).round() / 1000.0));
format!("Connection timed out. Timeout passed={timeout}, time taken={elapsed} seconds")
} else {
format!("Connection timed out after {timeout} seconds.")
}
}
fn python_float(value: Option<f64>) -> String {
match value {
None => "None".to_string(),
Some(value) if value.fract() == 0.0 => format!("{value:.1}"),
Some(value) => value.to_string(),
}
}
fn exception_provider(provider: &str) -> String {
if provider == "openai" {
return "OpenAIException".to_string();
}
let mut characters = provider.chars();
match characters.next() {
Some(first) => format!("{}{}Exception", first.to_uppercase(), characters.as_str()),
None => String::new(),
}
}
fn api_base(context: &ExceptionContext) -> Option<String> {
match (&context.vertex_location, &context.vertex_project) {
(Some(location), Some(project)) => Some(format!(
"{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/publishers/google/models/{}:generateContent",
context.model
)),
_ => None,
}
}
fn extra_information(context: &ExceptionContext, api_base: Option<&str>) -> String {
let lines = [
Some(format!("\nModel: {}", context.model)),
api_base.map(|api_base| format!("\nAPI Base: `{api_base}`")),
context
.model_group
.as_ref()
.map(|value| format!("\nmodel_group: `{value}`\n")),
context
.deployment
.as_ref()
.map(|value| format!("\ndeployment: `{value}`\n")),
context
.vertex_project
.as_ref()
.map(|value| format!("\nvertex_project: `{value}`\n")),
context
.vertex_location
.as_ref()
.map(|value| format!("\nvertex_location: `{value}`\n")),
];
let information: String = lines.into_iter().flatten().collect();
match &context.user_api_key_alias {
Some(alias) => format!(
"\n\nKey Name: `{alias}`\nTeam: `{}`{information}",
context.user_api_key_team_alias.as_deref().unwrap_or("None")
),
None => information,
}
}
#[cfg(test)]
mod testing {
use super::Mapping;
pub(super) fn mapping(status: Option<u16>, text: &str) -> Mapping {
Mapping {
status,
error_str: text.into(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
const DEBUG: &str = "\nModel: ocr-model";
fn context(provider: &str) -> ExceptionContext {
ExceptionContext {
model: "ocr-model".into(),
custom_llm_provider: provider.into(),
..ExceptionContext::default()
}
}
fn redactor() -> SecretRedactor {
SecretRedactor::new(16)
}
fn headers() -> Vec<(String, String)> {
vec![("retry-after".into(), "7".into())]
}
fn http(status: u16, body: &str) -> OriginalException {
OriginalException::Http {
status,
body: body.into(),
headers: headers(),
}
}
fn upstream(status: u16, body: &str) -> Option<UpstreamResponse> {
Some(UpstreamResponse {
status,
body: body.into(),
headers: headers(),
})
}
fn mapped(provider: &str, original: &OriginalException) -> MappedFailure {
exception_type(&context(provider), Some(&redactor()), original)
}
#[rstest::rstest]
#[case::openai_family("mistral", "rate limit reached", PublicError::RateLimit)]
#[case::vertex_family("vertex_ai", "Resource exhausted", PublicError::RateLimit)]
#[case::cohere_family("cohere", "too many tokens", PublicError::ContextWindowExceeded)]
fn a_family_text_rule_beats_the_status_and_keeps_the_real_response(
#[case] provider: &str,
#[case] body: &str,
#[case] expected: PublicError,
) {
let failure = mapped(provider, &http(401, body));
assert_eq!(failure.error, expected);
assert_eq!(failure.upstream, upstream(401, body));
}
#[test]
fn the_other_family_has_no_text_rules() {
assert_eq!(
mapped("reducto", &http(401, "rate limit reached")).error,
PublicError::Authentication
);
}
#[rstest::rstest]
#[case::openai_403_is_permission_denied("mistral", 403, PublicError::PermissionDenied)]
#[case::openai_409_is_bad_request("mistral", 409, PublicError::BadRequest)]
#[case::vertex_502_is_bad_gateway("vertex_ai", 502, PublicError::BadGateway)]
#[case::vertex_504_is_a_timeout("vertex_ai", 504, PublicError::Timeout { status: 504 })]
#[case::cohere_498_is_bad_request("cohere", 498, PublicError::BadRequest)]
#[case::other_503("reducto", 503, PublicError::ServiceUnavailable)]
fn without_a_text_rule_every_family_uses_the_status_table(
#[case] provider: &str,
#[case] status: u16,
#[case] expected: PublicError,
) {
assert_eq!(
mapped(provider, &http(status, "rejected")),
MappedFailure {
error: expected,
message: format!("{} - rejected", exception_provider(provider)),
upstream: upstream(status, "rejected"),
debug_info: DEBUG.into(),
}
);
}
#[rstest::rstest]
#[case::request_timeout_error("Request Timeout Error")]
#[case::request_timed_out("Request timed out")]
#[case::timed_out_generating("Timed out generating response")]
#[case::read_operation("The read operation timed out")]
fn timeout_markers_win_over_every_family(#[case] marker: &str) {
let body = format!("rate limit invalid api token {marker}");
for provider in ["mistral", "vertex_ai", "cohere", "reducto"] {
assert_eq!(
mapped(provider, &http(429, &body)).error,
PublicError::Timeout { status: 408 },
"{provider}"
);
}
}
#[test]
fn a_handler_timeout_is_a_408_without_a_response() {
let original = OriginalException::Timeout {
timeout_seconds: Some(0.5),
elapsed_seconds: Some(0.5031),
};
assert_eq!(
mapped("mistral", &original),
MappedFailure {
error: PublicError::Timeout { status: 408 },
message: "MistralException - Connection timed out after 0.5 seconds.".into(),
upstream: None,
debug_info: DEBUG.into(),
}
);
}
#[rstest::rstest]
#[case::refused_connection(OriginalException::Connection { message: "refused".into() })]
#[case::unparseable_response(OriginalException::Plain { message: "refused".into() })]
#[case::informational_status(OriginalException::Http { status: 399, body: "refused".into(), headers: Vec::new() })]
fn a_failure_no_rule_or_status_claims_is_a_connection_error(
#[case] original: OriginalException,
) {
let failure = mapped("reducto", &original);
assert_eq!(failure.error, PublicError::ApiConnection);
assert_eq!(failure.message, "ReductoException - refused");
}
#[test]
fn a_timeout_marker_on_a_response_keeps_the_response() {
let failure = mapped("reducto", &http(429, "Request timed out"));
assert_eq!(failure.error, PublicError::Timeout { status: 408 });
assert_eq!(failure.upstream, upstream(429, "Request timed out"));
}
#[test]
fn family_text_rules_also_classify_failures_without_a_response() {
let original = OriginalException::Plain {
message: "Request too large".into(),
};
assert_eq!(mapped("mistral", &original).error, PublicError::RateLimit);
}
#[rstest::rstest]
#[case::openai_family("mistral", "MistralException - rejected REDACTED")]
#[case::vertex_family("vertex_ai", "Vertex_aiException - rejected REDACTED")]
#[case::other_family("reducto", "ReductoException - rejected REDACTED")]
fn every_family_reports_the_redacted_text(#[case] provider: &str, #[case] message: &str) {
let failure = mapped(provider, &http(400, "rejected Bearer abcdefghijklmnop"));
assert_eq!(failure.message, message);
}
#[test]
fn redaction_runs_before_the_rules_see_the_text() {
let body = "db_password=rate_limit";
assert_eq!(
mapped("mistral", &http(400, body)).error,
PublicError::BadRequest
);
assert_eq!(
exception_type(&context("mistral"), None, &http(400, body)).error,
PublicError::RateLimit
);
}
#[test]
fn without_a_redactor_the_text_is_kept() {
let body = "rejected Bearer abcdefghijklmnop";
assert_eq!(
exception_type(&context("reducto"), None, &http(400, body)).message,
format!("ReductoException - {body}")
);
}
#[test]
fn a_rule_hint_follows_the_message() {
let failure = mapped("mistral", &http(400, "invalid_encrypted_content"));
assert_eq!(failure.error, PublicError::BadRequest);
assert!(
failure
.message
.starts_with("MistralException - invalid_encrypted_content\n\n This error occurs")
);
}
#[rstest::rstest]
#[case::sync(
false,
Some(0.5),
Some(0.5031),
"Connection timed out after 0.5 seconds."
)]
#[case::async_rounds_the_elapsed_time(
true,
Some(0.5),
Some(0.5031),
"Connection timed out. Timeout passed=0.5, time taken=0.503 seconds"
)]
#[case::whole_seconds_keep_a_decimal(
true,
Some(600.0),
Some(2.0),
"Connection timed out. Timeout passed=600.0, time taken=2.0 seconds"
)]
#[case::unknown_values_render_as_none(
true,
None,
None,
"Connection timed out. Timeout passed=None, time taken=None seconds"
)]
fn timeout_text_follows_the_delivery_mode(
#[case] asynchronous: bool,
#[case] timeout_seconds: Option<f64>,
#[case] elapsed_seconds: Option<f64>,
#[case] expected: &str,
) {
assert_eq!(
timeout_message(asynchronous, timeout_seconds, elapsed_seconds),
expected
);
}
#[test]
fn debug_information_follows_the_python_layout() {
let context = ExceptionContext {
vertex_project: Some("project".into()),
vertex_location: Some("region".into()),
model_group: Some("ocr".into()),
deployment: Some("deployment".into()),
user_api_key_alias: Some("key".into()),
..context("vertex_ai")
};
assert_eq!(
exception_type(&context, None, &http(400, "rejected")).debug_info,
concat!(
"\n\nKey Name: `key`\nTeam: `None`",
"\nModel: ocr-model",
"\nAPI Base: `region-aiplatform.googleapis.com/v1/projects/project/locations/region/publishers/google/models/ocr-model:generateContent`",
"\nmodel_group: `ocr`\n",
"\ndeployment: `deployment`\n",
"\nvertex_project: `project`\n",
"\nvertex_location: `region`\n",
)
);
}
#[rstest::rstest]
#[case::bare(ExceptionContext::default(), "\nModel: ")]
#[case::team_alias(
ExceptionContext { model: "m".into(), user_api_key_alias: Some("key".into()), user_api_key_team_alias: Some("team".into()), ..ExceptionContext::default() },
"\n\nKey Name: `key`\nTeam: `team`\nModel: m"
)]
#[case::team_alias_without_key_is_ignored(
ExceptionContext { model: "m".into(), user_api_key_team_alias: Some("team".into()), ..ExceptionContext::default() },
"\nModel: m"
)]
#[case::project_without_location_has_no_api_base(
ExceptionContext { model: "m".into(), vertex_project: Some("p".into()), ..ExceptionContext::default() },
"\nModel: m\nvertex_project: `p`\n"
)]
#[case::location_without_project_has_no_api_base(
ExceptionContext { model: "m".into(), vertex_location: Some("l".into()), ..ExceptionContext::default() },
"\nModel: m\nvertex_location: `l`\n"
)]
fn each_optional_context_field_adds_its_own_line(
#[case] context: ExceptionContext,
#[case] expected: &str,
) {
assert_eq!(
extra_information(&context, api_base(&context).as_deref()),
expected
);
}
#[rstest::rstest]
#[case::openai_keeps_its_brand("openai", "OpenAIException")]
#[case::lowercase("mistral", "MistralException")]
#[case::keeps_the_rest("azure_ai", "Azure_aiException")]
#[case::empty("", "")]
fn exception_provider_capitalizes_only_the_first_letter(
#[case] provider: &str,
#[case] expected: &str,
) {
assert_eq!(exception_provider(provider), expected);
}
}

View file

@ -0,0 +1,192 @@
use super::public::PublicError;
use super::rules::{Rule, contains_any, is_context_window_exceeded, is_rate_limit};
const ENCRYPTED_CONTENT_HELP: &str = "\n\n This error occurs when load balancing Responses API across deployments with different API keys.\n Encrypted content is tied to the organization that created it and cannot be decrypted by other organizations.\n\n Solution: Enable 'encrypted_content_affinity' to route follow-up requests to the correct deployment:\n\n router_settings:\n enable_pre_call_checks: true\n optional_pre_call_checks:\n - encrypted_content_affinity\n\n Learn more: https://docs.litellm.ai/docs/response_api#encrypted-content-affinity-multi-region-load-balancing";
/// The text branches of `_map_openai_exception`, in its order.
pub(super) const RULES: &[Rule] = &[
Rule::new(
|mapping| is_rate_limit(&mapping.error_str, mapping.status),
PublicError::RateLimit,
),
Rule::new(
|mapping| is_context_window_exceeded(&mapping.error_str),
PublicError::ContextWindowExceeded,
),
Rule::new(
|mapping| {
mapping.error_str.contains("invalid_request_error")
&& mapping.error_str.contains("model_not_found")
},
PublicError::NotFound,
),
Rule::new(
|mapping| mapping.error_str.contains("A timeout occurred"),
PublicError::Timeout { status: 408 },
),
Rule::new(
|mapping| {
let error_str = &mapping.error_str;
(error_str.contains("invalid_request_error")
&& error_str.contains("content_policy_violation"))
|| (error_str.contains("Invalid prompt")
&& error_str.contains("violating our usage policy"))
|| error_str
.to_lowercase()
.contains("request was rejected as a result of the safety system")
},
PublicError::ContentPolicyViolation,
),
Rule {
hint: ENCRYPTED_CONTENT_HELP,
..Rule::new(
|mapping| {
contains_any(
&mapping.error_str,
&["invalid_encrypted_content", "could not be verified"],
)
},
PublicError::BadRequest,
)
},
Rule::new(
|mapping| {
mapping.error_str.contains("invalid_request_error")
&& !mapping.error_str.contains("Incorrect API key provided")
},
PublicError::BadRequest,
),
Rule::new(
|mapping| {
contains_any(
&mapping.error_str,
&[
"Web server is returning an unknown error",
"The server had an error processing your request.",
],
)
},
PublicError::InternalServer,
),
Rule::new(
|mapping| mapping.error_str.contains("Request too large"),
PublicError::RateLimit,
),
Rule::new(
|mapping| {
mapping
.error_str
.contains("Mistral API raised a streaming error")
},
PublicError::Api { status: 500 },
),
];
#[cfg(test)]
mod tests {
use super::super::rules::first_match;
use super::super::testing::mapping;
use super::*;
fn classified(status: Option<u16>, text: &str) -> Option<PublicError> {
first_match(RULES, &mapping(status, text)).map(|rule| rule.error)
}
#[rstest::rstest]
#[case::rate_limit_phrase("rate limit reached", PublicError::RateLimit)]
#[case::context_window(
"This model's maximum context length is 10",
PublicError::ContextWindowExceeded
)]
#[case::model_not_found("invalid_request_error model_not_found", PublicError::NotFound)]
#[case::timeout_occurred("A timeout occurred", PublicError::Timeout { status: 408 })]
#[case::content_policy_error_code(
"invalid_request_error content_policy_violation",
PublicError::ContentPolicyViolation
)]
#[case::content_policy_usage_policy(
"Invalid prompt violating our usage policy",
PublicError::ContentPolicyViolation
)]
#[case::content_policy_safety_system(
"Request was rejected as a result of the safety system",
PublicError::ContentPolicyViolation
)]
#[case::encrypted_content("invalid_encrypted_content", PublicError::BadRequest)]
#[case::unverifiable_content("could not be verified", PublicError::BadRequest)]
#[case::invalid_request("invalid_request_error bad field", PublicError::BadRequest)]
#[case::unknown_server_error(
"Web server is returning an unknown error",
PublicError::InternalServer
)]
#[case::server_had_an_error(
"The server had an error processing your request.",
PublicError::InternalServer
)]
#[case::request_too_large("Request too large", PublicError::RateLimit)]
#[case::mistral_streaming_error(
"Mistral API raised a streaming error",
PublicError::Api { status: 500 }
)]
fn each_text_rule_claims_its_marker(#[case] text: &str, #[case] expected: PublicError) {
assert_eq!(classified(Some(400), text), Some(expected));
}
#[rstest::rstest]
#[case::rate_limit_before_context_window(
"rate limit and This model's maximum context length is 10",
PublicError::RateLimit
)]
#[case::context_window_before_content_policy(
"This model's maximum context length is 10 invalid_request_error content_policy_violation",
PublicError::ContextWindowExceeded
)]
#[case::model_not_found_before_invalid_request(
"invalid_request_error model_not_found",
PublicError::NotFound
)]
#[case::timeout_before_invalid_request(
"A timeout occurred invalid_request_error",
PublicError::Timeout { status: 408 }
)]
#[case::content_policy_before_invalid_request(
"invalid_request_error content_policy_violation",
PublicError::ContentPolicyViolation
)]
#[case::encrypted_content_before_invalid_request(
"invalid_request_error invalid_encrypted_content",
PublicError::BadRequest
)]
fn the_earlier_rule_wins_when_two_apply(#[case] text: &str, #[case] expected: PublicError) {
assert_eq!(classified(Some(400), text), Some(expected));
}
#[rstest::rstest]
#[case::encrypted_content("invalid_encrypted_content", ENCRYPTED_CONTENT_HELP)]
#[case::plain_invalid_request("invalid_request_error bad field", "")]
fn only_encrypted_content_failures_carry_the_affinity_help(
#[case] text: &str,
#[case] hint: &str,
) {
assert_eq!(
first_match(RULES, &mapping(Some(400), text)).map(|rule| rule.hint),
Some(hint)
);
}
#[rstest::rstest]
#[case::bad_key_is_left_to_the_status("invalid_request_error Incorrect API key provided")]
#[case::echoed_429_is_not_a_rate_limit("token 429 in the prompt")]
#[case::unmarked("rejected")]
fn text_without_a_marker_is_left_to_the_status_table(#[case] text: &str) {
assert_eq!(classified(Some(400), text), None);
}
#[test]
fn a_standalone_429_counts_with_a_429_status() {
assert_eq!(
classified(Some(429), "got 429 back"),
Some(PublicError::RateLimit)
);
}
}

View file

@ -0,0 +1,144 @@
/// A failure a Rust route produced, before any public class is chosen.
#[derive(Clone, Debug, PartialEq)]
pub enum OriginalException {
Http {
status: u16,
body: String,
headers: Vec<(String, String)>,
},
Connection {
message: String,
},
Timeout {
timeout_seconds: Option<f64>,
elapsed_seconds: Option<f64>,
},
/// A failure with no HTTP response behind it, such as an unparseable body or a local
/// file error.
Plain {
message: String,
},
}
/// Which provider-specific text rules apply before the shared status table.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExceptionFamily {
OpenAiCompatible,
VertexAi,
Cohere,
Other,
}
/// `openai_compatible_providers` in `litellm/constants.py`.
const OPENAI_COMPATIBLE_PROVIDERS: &[&str] = &[
"anyscale",
"groq",
"nvidia_nim",
"cerebras",
"baseten",
"sambanova",
"ai21_chat",
"ai21",
"volcengine",
"codestral",
"deepseek",
"tencent",
"deepinfra",
"perplexity",
"xinference",
"xai",
"zai",
"together_ai",
"fireworks_ai",
"empower",
"friendliai",
"azure_ai",
"github",
"litellm_proxy",
"hosted_vllm",
"llamafile",
"lm_studio",
"galadriel",
"github_copilot",
"chatgpt",
"novita",
"meta_llama",
"publicai",
"synthetic",
"tensormesh",
"apertis",
"nano-gpt",
"poe",
"chutes",
"parasail",
"libertai",
"featherless_ai",
"nscale",
"nebius",
"dashscope",
"qwencloud",
"qwen_ai_platform",
"modelscope",
"moonshot",
"v0",
"helicone",
"morph",
"lambda_ai",
"inception",
"hyperbolic",
"vercel_ai_gateway",
"aiml",
"wandb",
"cometapi",
"clarifai",
"docker_model_runner",
"ragflow",
"pinstripes",
"darkbloom",
"meta",
"cognition",
"scx-ai",
];
impl ExceptionFamily {
/// The provider dispatch at the top of Python's `exception_type`, in its order.
pub fn for_provider(provider: &str) -> Self {
match provider {
"openai" | "text-completion-openai" | "custom_openai" | "mistral" | "runwayml" => {
Self::OpenAiCompatible
}
provider if OPENAI_COMPATIBLE_PROVIDERS.contains(&provider) => Self::OpenAiCompatible,
"vertex_ai" | "vertex_ai_beta" | "gemini" => Self::VertexAi,
"cohere" | "cohere_chat" => Self::Cohere,
_ => Self::Other,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[rstest::rstest]
#[case::openai("openai", ExceptionFamily::OpenAiCompatible)]
#[case::text_completion_openai("text-completion-openai", ExceptionFamily::OpenAiCompatible)]
#[case::custom_openai("custom_openai", ExceptionFamily::OpenAiCompatible)]
#[case::mistral("mistral", ExceptionFamily::OpenAiCompatible)]
#[case::runwayml("runwayml", ExceptionFamily::OpenAiCompatible)]
#[case::listed_compatible("azure_ai", ExceptionFamily::OpenAiCompatible)]
#[case::compatible_list_wins_over_its_own_mapper(
"together_ai",
ExceptionFamily::OpenAiCompatible
)]
#[case::vertex_ai("vertex_ai", ExceptionFamily::VertexAi)]
#[case::vertex_ai_beta("vertex_ai_beta", ExceptionFamily::VertexAi)]
#[case::gemini("gemini", ExceptionFamily::VertexAi)]
#[case::cohere("cohere", ExceptionFamily::Cohere)]
#[case::cohere_chat("cohere_chat", ExceptionFamily::Cohere)]
#[case::unported_mapper("anthropic", ExceptionFamily::Other)]
#[case::unknown("reducto", ExceptionFamily::Other)]
#[case::empty("", ExceptionFamily::Other)]
fn provider_selects_the_family(#[case] provider: &str, #[case] family: ExceptionFamily) {
assert_eq!(ExceptionFamily::for_provider(provider), family);
}
}

View file

@ -0,0 +1,77 @@
/// The public LiteLLM exception classes a Rust route failure can become. Python builds the
/// class; Rust decides which one.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PublicError {
BadRequest,
ContextWindowExceeded,
ContentPolicyViolation,
Authentication,
PermissionDenied,
NotFound,
Timeout { status: u16 },
RateLimit,
InternalServer,
BadGateway,
ServiceUnavailable,
ApiConnection,
Api { status: u16 },
}
impl PublicError {
/// The `status_code` the Python class carries.
pub const fn status_code(self) -> u16 {
match self {
Self::BadRequest | Self::ContextWindowExceeded | Self::ContentPolicyViolation => 400,
Self::Authentication => 401,
Self::PermissionDenied => 403,
Self::NotFound => 404,
Self::RateLimit => 429,
Self::InternalServer | Self::ApiConnection => 500,
Self::BadGateway => 502,
Self::ServiceUnavailable => 503,
Self::Timeout { status } | Self::Api { status } => status,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UpstreamResponse {
pub status: u16,
pub body: String,
pub headers: Vec<(String, String)>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MappedFailure {
pub error: PublicError,
pub message: String,
pub upstream: Option<UpstreamResponse>,
pub debug_info: String,
}
#[cfg(test)]
mod tests {
use super::*;
#[rstest::rstest]
#[case::bad_request(PublicError::BadRequest, 400)]
#[case::context_window(PublicError::ContextWindowExceeded, 400)]
#[case::content_policy(PublicError::ContentPolicyViolation, 400)]
#[case::authentication(PublicError::Authentication, 401)]
#[case::permission_denied(PublicError::PermissionDenied, 403)]
#[case::not_found(PublicError::NotFound, 404)]
#[case::request_timeout(PublicError::Timeout { status: 408 }, 408)]
#[case::gateway_timeout(PublicError::Timeout { status: 504 }, 504)]
#[case::rate_limit(PublicError::RateLimit, 429)]
#[case::internal_server(PublicError::InternalServer, 500)]
#[case::api_connection(PublicError::ApiConnection, 500)]
#[case::bad_gateway(PublicError::BadGateway, 502)]
#[case::service_unavailable(PublicError::ServiceUnavailable, 503)]
#[case::api(PublicError::Api { status: 501 }, 501)]
fn status_codes_are_the_ones_the_python_classes_set(
#[case] error: PublicError,
#[case] status: u16,
) {
assert_eq!(error.status_code(), status);
}
}

View file

@ -0,0 +1,176 @@
use std::sync::LazyLock;
use fancy_regex::Regex;
use serde_json::Value;
use super::Mapping;
use super::public::PublicError;
/// One text branch of a Python `_map_*_exception` function: when it applies, the class it
/// raises, and any help text appended to the message.
pub(super) struct Rule {
pub(super) when: fn(&Mapping) -> bool,
pub(super) error: PublicError,
pub(super) hint: &'static str,
}
impl Rule {
pub(super) const fn new(when: fn(&Mapping) -> bool, error: PublicError) -> Self {
Self {
when,
error,
hint: "",
}
}
}
/// The first rule that applies decides the class, as the `if`/`elif` chain does in Python.
pub(super) fn first_match<'r>(rules: &'r [Rule], mapping: &Mapping) -> Option<&'r Rule> {
rules.iter().find(|rule| (rule.when)(mapping))
}
pub(super) fn contains_any(text: &str, markers: &[&str]) -> bool {
markers.iter().any(|marker| text.contains(marker))
}
static STANDALONE_429: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"\b429\b").expect("valid regex"));
static RATE_LIMIT_PHRASE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"rate[\s_\-]*limit").expect("valid regex"));
/// `ExceptionCheckers.is_error_str_rate_limit`.
pub(super) fn is_rate_limit(error_str: &str, status: Option<u16>) -> bool {
if STANDALONE_429.is_match(error_str).unwrap_or(false) && matches!(status, None | Some(429)) {
return true;
}
let lower = error_str.to_lowercase();
RATE_LIMIT_PHRASE.is_match(&lower).unwrap_or(false)
|| lower.contains("service tier capacity exceeded")
}
/// `ExceptionCheckers.is_error_str_context_window_exceeded`.
pub(super) fn is_context_window_exceeded(error_str: &str) -> bool {
let lower = error_str.to_lowercase();
if lower.contains("string_above_max_length") {
return false;
}
if lower.contains("invalid 'user'") && lower.contains("string too long") {
return false;
}
contains_any(
&lower,
&[
"exceed context limit",
"this model's maximum context length is",
"string too long. expected a string with maximum length",
"model's maximum context limit",
"is longer than the model's context length",
"input tokens exceed the configured limit",
"`inputs` tokens + `max_new_tokens` must be",
"exceeds the available context size",
"exceeds the maximum number of tokens allowed",
],
) || (lower.contains("current length is") && lower.contains("while limit is"))
|| (lower.contains("maximum input length is") && lower.contains("tokens"))
}
/// The integer `error.code` of a JSON error body, read the way Python's `int()` would.
pub(super) fn body_error_code(error_str: &str) -> Option<i64> {
let body: Value = serde_json::from_str(error_str).ok()?;
let Some(Value::Object(error)) = body.as_object()?.get("error") else {
return None;
};
match error.get("code")? {
Value::Number(number) => number
.as_i64()
.or_else(|| number.as_f64().map(|value| value.trunc() as i64)),
Value::String(code) => code.trim().replace('_', "").parse().ok(),
Value::Bool(flag) => Some(i64::from(*flag)),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::super::testing::mapping;
use super::*;
const ORDERED: &[Rule] = &[
Rule::new(
|mapping| mapping.error_str.contains("first"),
PublicError::NotFound,
),
Rule::new(|_| true, PublicError::ApiConnection),
];
#[rstest::rstest]
#[case::earlier_rule_wins("first and second", PublicError::NotFound)]
#[case::later_rule_when_the_earlier_does_not_apply("second", PublicError::ApiConnection)]
fn the_first_applicable_rule_decides(#[case] text: &str, #[case] expected: PublicError) {
let rule = first_match(ORDERED, &mapping(Some(400), text));
assert_eq!(rule.map(|rule| rule.error), Some(expected));
}
#[test]
fn no_applicable_rule_leaves_the_failure_to_the_caller() {
assert!(first_match(&ORDERED[..1], &mapping(Some(400), "second")).is_none());
}
#[rstest::rstest]
#[case::standalone_429_with_429_status("got 429 back", Some(429), true)]
#[case::standalone_429_with_other_status("got 429 back", Some(400), false)]
#[case::standalone_429_with_unknown_status("got 429 back", None, true)]
#[case::embedded_429("token4290", Some(429), false)]
#[case::phrase_spaced("Rate Limit reached", None, true)]
#[case::phrase_underscored("rate_limit", None, true)]
#[case::phrase_hyphenated("rate-limit", None, true)]
#[case::service_tier("Service tier capacity exceeded", None, true)]
#[case::unrelated("rejected", Some(429), false)]
fn rate_limit_detection(
#[case] text: &str,
#[case] status: Option<u16>,
#[case] expected: bool,
) {
assert_eq!(is_rate_limit(text, status), expected);
}
#[rstest::rstest]
#[case::exceed_context_limit("Exceed context limit", true)]
#[case::maximum_context_length("This model's maximum context length is 10", true)]
#[case::string_too_long("string too long. Expected a string with maximum length 5", true)]
#[case::maximum_context_limit("the model's maximum context limit", true)]
#[case::longer_than_context("prompt is longer than the model's context length", true)]
#[case::configured_limit("input tokens exceed the configured limit", true)]
#[case::max_new_tokens("`inputs` tokens + `max_new_tokens` must be <= 10", true)]
#[case::available_context("exceeds the available context size", true)]
#[case::maximum_tokens("exceeds the maximum number of tokens allowed", true)]
#[case::current_and_limit("current length is 9 while limit is 8", true)]
#[case::current_without_limit("current length is 9", false)]
#[case::maximum_input_tokens("maximum input length is 8 tokens", true)]
#[case::maximum_input_without_tokens("maximum input length is 8", false)]
#[case::string_above_max_length_wins("string_above_max_length exceed context limit", false)]
#[case::user_field_is_not_context(
"invalid 'user': string too long. expected a string with maximum length",
false
)]
#[case::unrelated("rejected", false)]
fn context_window_detection(#[case] text: &str, #[case] expected: bool) {
assert_eq!(is_context_window_exceeded(text), expected);
}
#[rstest::rstest]
#[case::integer(r#"{"error": {"code": 429}}"#, Some(429))]
#[case::float(r#"{"error": {"code": 429.9}}"#, Some(429))]
#[case::string(r#"{"error": {"code": " 4_29 "}}"#, Some(429))]
#[case::boolean(r#"{"error": {"code": true}}"#, Some(1))]
#[case::unparseable_string(r#"{"error": {"code": "slow"}}"#, None)]
#[case::null(r#"{"error": {"code": null}}"#, None)]
#[case::no_code(r#"{"error": {}}"#, None)]
#[case::error_not_an_object(r#"{"error": "429"}"#, None)]
#[case::no_error(r#"{"code": 429}"#, None)]
#[case::not_an_object("[429]", None)]
#[case::not_json("429", None)]
fn body_error_code_reads_the_nested_code(#[case] body: &str, #[case] expected: Option<i64>) {
assert_eq!(body_error_code(body), expected);
}
}

View file

@ -0,0 +1,49 @@
use super::public::PublicError;
/// `_map_exception_by_status`, the one place a provider status picks a class. Statuses
/// below 400 are not failures the table claims.
pub(super) fn classify(status: u16) -> Option<PublicError> {
let error = match status {
..400 => return None,
401 => PublicError::Authentication,
403 => PublicError::PermissionDenied,
404 => PublicError::NotFound,
408 | 504 => PublicError::Timeout { status },
429 => PublicError::RateLimit,
500 => PublicError::InternalServer,
502 => PublicError::BadGateway,
503 => PublicError::ServiceUnavailable,
400..500 => PublicError::BadRequest,
_ => PublicError::Api { status },
};
Some(error)
}
#[cfg(test)]
mod tests {
use super::*;
#[rstest::rstest]
#[case::below_client_errors(399, None)]
#[case::lowest_client_error(400, Some(PublicError::BadRequest))]
#[case::authentication(401, Some(PublicError::Authentication))]
#[case::permission_denied(403, Some(PublicError::PermissionDenied))]
#[case::not_found(404, Some(PublicError::NotFound))]
#[case::request_timeout(408, Some(PublicError::Timeout { status: 408 }))]
#[case::other_client_error(409, Some(PublicError::BadRequest))]
#[case::unprocessable(422, Some(PublicError::BadRequest))]
#[case::rate_limited(429, Some(PublicError::RateLimit))]
#[case::highest_client_error(499, Some(PublicError::BadRequest))]
#[case::internal_server(500, Some(PublicError::InternalServer))]
#[case::other_server_error(501, Some(PublicError::Api { status: 501 }))]
#[case::bad_gateway(502, Some(PublicError::BadGateway))]
#[case::service_unavailable(503, Some(PublicError::ServiceUnavailable))]
#[case::gateway_timeout(504, Some(PublicError::Timeout { status: 504 }))]
#[case::highest_server_error(599, Some(PublicError::Api { status: 599 }))]
fn every_mapped_status_and_the_fallback(
#[case] status: u16,
#[case] expected: Option<PublicError>,
) {
assert_eq!(classify(status), expected);
}
}

View file

@ -0,0 +1,177 @@
use super::public::PublicError;
use super::rules::{Rule, body_error_code, contains_any, is_context_window_exceeded};
const QUOTA_MARKERS: &[&str] = &[
"429 Quota exceeded",
"Quota exceeded for",
"Resource exhausted",
"429 Unable to submit request because the service is temporarily out of capacity.",
];
/// The text branches of `_map_vertex_exception`, in its order.
pub(super) const RULES: &[Rule] = &[
Rule::new(
|mapping| {
contains_any(
&mapping.error_str,
&[
"Vertex AI API has not been used in project",
"Unable to find your project",
],
)
},
PublicError::BadRequest,
),
Rule::new(
|mapping| {
mapping
.error_str
.contains("400 Request payload size exceeds")
|| is_context_window_exceeded(&mapping.error_str)
},
PublicError::ContextWindowExceeded,
),
Rule::new(
|mapping| {
contains_any(
&mapping.error_str,
&["None Unknown Error.", "Content has no parts."],
)
},
PublicError::InternalServer,
),
Rule::new(
|mapping| mapping.error_str.contains("API key not valid."),
PublicError::Authentication,
),
Rule::new(
|mapping| {
contains_any(
&mapping.error_str,
&[
"The response was blocked.",
"Output blocked by content filtering policy",
],
)
},
PublicError::ContentPolicyViolation,
),
Rule::new(
|mapping| {
contains_any(&mapping.error_str, QUOTA_MARKERS)
|| (mapping
.status
.is_some_and(|status| (500..600).contains(&status))
&& body_error_code(&mapping.error_str) == Some(429))
},
PublicError::RateLimit,
),
Rule::new(
|mapping| {
contains_any(
&mapping.error_str,
&["500 Internal Server Error", "The model is overloaded."],
)
},
PublicError::InternalServer,
),
];
#[cfg(test)]
mod tests {
use super::super::rules::first_match;
use super::super::testing::mapping;
use super::*;
fn classified(status: Option<u16>, text: &str) -> Option<PublicError> {
first_match(RULES, &mapping(status, text)).map(|rule| rule.error)
}
#[rstest::rstest]
#[case::api_not_enabled(
"Vertex AI API has not been used in project x",
PublicError::BadRequest
)]
#[case::project_not_found("Unable to find your project", PublicError::BadRequest)]
#[case::payload_too_large(
"400 Request payload size exceeds the limit",
PublicError::ContextWindowExceeded
)]
#[case::context_window(
"This model's maximum context length is 10",
PublicError::ContextWindowExceeded
)]
#[case::unknown_error("None Unknown Error.", PublicError::InternalServer)]
#[case::no_parts("Content has no parts.", PublicError::InternalServer)]
#[case::api_key_not_valid("API key not valid.", PublicError::Authentication)]
#[case::response_blocked("The response was blocked.", PublicError::ContentPolicyViolation)]
#[case::output_blocked(
"Output blocked by content filtering policy",
PublicError::ContentPolicyViolation
)]
#[case::quota_exceeded_429("429 Quota exceeded", PublicError::RateLimit)]
#[case::quota_exceeded_for("Quota exceeded for aiplatform", PublicError::RateLimit)]
#[case::resource_exhausted("Resource exhausted", PublicError::RateLimit)]
#[case::out_of_capacity(
"429 Unable to submit request because the service is temporarily out of capacity.",
PublicError::RateLimit
)]
#[case::internal_server_text("500 Internal Server Error", PublicError::InternalServer)]
#[case::overloaded("The model is overloaded.", PublicError::InternalServer)]
fn each_text_rule_claims_its_marker(#[case] text: &str, #[case] expected: PublicError) {
assert_eq!(classified(Some(400), text), Some(expected));
}
#[rstest::rstest]
#[case::server_error_wrapping_a_429(Some(503), Some(PublicError::RateLimit))]
#[case::lowest_server_error(Some(500), Some(PublicError::RateLimit))]
#[case::highest_server_error(Some(599), Some(PublicError::RateLimit))]
#[case::client_error(Some(400), None)]
#[case::no_status(None, None)]
fn a_wrapped_429_is_a_rate_limit_only_behind_a_server_error(
#[case] status: Option<u16>,
#[case] expected: Option<PublicError>,
) {
assert_eq!(
classified(status, r#"{"error": {"code": "429"}}"#),
expected
);
}
#[rstest::rstest]
#[case::project_before_payload_size(
"Unable to find your project 400 Request payload size exceeds",
PublicError::BadRequest
)]
#[case::context_window_before_unknown_error(
"This model's maximum context length is 10 None Unknown Error.",
PublicError::ContextWindowExceeded
)]
#[case::unknown_error_before_api_key(
"Content has no parts. API key not valid.",
PublicError::InternalServer
)]
#[case::api_key_before_blocked(
"API key not valid. The response was blocked.",
PublicError::Authentication
)]
#[case::blocked_before_quota(
"The response was blocked. Resource exhausted",
PublicError::ContentPolicyViolation
)]
#[case::quota_before_overloaded(
"Resource exhausted The model is overloaded.",
PublicError::RateLimit
)]
fn the_earlier_rule_wins_when_two_apply(#[case] text: &str, #[case] expected: PublicError) {
assert_eq!(classified(Some(400), text), Some(expected));
}
#[rstest::rstest]
#[case::a_403_in_the_text("got a 403 from 4031 tokens")]
#[case::python_client_crash("IndexError: list index out of range")]
#[case::unmarked("rejected")]
fn text_without_a_marker_is_left_to_the_status_table(#[case] text: &str) {
assert_eq!(classified(Some(400), text), None);
}
}

View file

@ -1,7 +1,9 @@
pub mod call_arguments;
pub mod core_helpers;
pub mod exception_mapping_utils;
pub mod get_llm_provider_logic;
pub mod params;
pub mod prompt_templates;
pub mod secret_redaction;
pub mod serde_compat;
pub mod url_utils;

View file

@ -0,0 +1,109 @@
use fancy_regex::Regex;
pub const REDACTED: &str = "REDACTED";
const DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH: usize = 16;
fn minimum_custom_key_length() -> usize {
std::env::var("MINIMUM_CUSTOM_KEY_LENGTH")
.ok()
.and_then(|value| value.trim().parse().ok())
.unwrap_or(DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH)
}
fn secret_patterns(minimum_custom_key_length: usize) -> String {
let sk_suffix_length = minimum_custom_key_length.saturating_sub("sk-".len());
[
r"-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----",
r"\bya29\.[A-Za-z0-9_.~+/-]+",
r#"(?:client_secret|azure_password|azure_username)\s+[^\s,'"})\]{}>]+"#,
r"(?:AKIA|ASIA)[0-9A-Z]{16}",
r"Bearer\s+[A-Za-z0-9\-._~+/]{10,}=*",
r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}",
&format!(r"sk-[A-Za-z0-9\-_]{{{sk_suffix_length},}}"),
r#"(?<=[?&])(?:api[_-]?key|\w*(?:token|password|passwd|client_secret|secret_key|_secret))=[^\s&'"]+"#,
r#"(?:api[_-]?key)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]{8,}"#,
r#"(?:x-api-key|api-key)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#,
r"x-ak-[A-Za-z0-9\-_]{20,}",
r"AIza[0-9A-Za-z\-_]{35}",
r#"(?<=[?&])key=[^\s&'"]{8,}"#,
r#"(?:^|(?<=\W))\w*(?:password|passwd|client_secret|secret_key|_secret)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#,
r#"(?<=://)[^\s'":]{0,4096}:[^\s'"]{1,4096}(?=@)"#,
r"dapi[0-9a-f]{32}",
r#"litellm\.[A-Za-z0-9_]*_key['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#,
r#"private_key['"]?\s*[:=]\s*['"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'"})\]{}>]+)"#,
concat!(
r"(?:master_key|xai_key|database_url|db_url|connection_string|",
r"aws_secret_access_key|aws_session_token|aws_access_key_id|",
r"signing_key|encryption_key|",
r"auth_token|access_token|refresh_token|",
r"slack_webhook_url|webhook_url|",
r"database_connection_string|",
r"huggingface_token|jwt_secret)",
r#"['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#,
),
r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]*",
r"(?<=[?&])sig=[A-Za-z0-9%+/=]+",
r#"\{[^{}]*"type"\s*:\s*"service_account"[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"#,
]
.join("|")
}
/// Python's `_ENABLE_SECRET_REDACTION` pattern set, compiled once per configuration.
#[derive(Clone, Debug)]
pub struct SecretRedactor {
pattern: Regex,
}
impl SecretRedactor {
pub fn new(minimum_custom_key_length: usize) -> Self {
let pattern = Regex::new(&format!(
"(?i){}",
secret_patterns(minimum_custom_key_length)
))
.expect("secret redaction patterns compile");
Self { pattern }
}
/// `None` when `LITELLM_DISABLE_REDACT_SECRETS` turns redaction off.
pub fn from_env() -> Option<Self> {
let disabled = std::env::var("LITELLM_DISABLE_REDACT_SECRETS")
.is_ok_and(|value| value.eq_ignore_ascii_case("true"));
(!disabled).then(|| Self::new(minimum_custom_key_length()))
}
pub fn redact(&self, value: &str) -> String {
self.pattern.replace_all(value, REDACTED).into_owned()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[rstest::rstest]
#[case::bearer("auth failed: Bearer abcdefghijklmnop", "auth failed: REDACTED")]
#[case::sk_key("key sk-abcdefghijklmnopqrstuvwxyz rejected", "key REDACTED rejected")]
#[case::short_sk_key_is_kept("sk-abc", "sk-abc")]
#[case::query_param("GET /v1?api_key=secret123&x=1", "GET /v1?REDACTED&x=1")]
#[case::dict_repr("{'api_key': 'abcdefghij'}", "{'REDACTED'}")]
#[case::url_credentials("postgres://user:pass@host/db", "postgres://REDACTED@host/db")]
#[case::case_insensitive("BEARER ABCDEFGHIJKLMNOP", "REDACTED")]
#[case::aws_key("AKIAABCDEFGHIJKLMNOP", "REDACTED")]
#[case::sas_signature("https://x.blob/a?sv=1&sig=abc%2B=", "https://x.blob/a?sv=1&REDACTED")]
#[case::password_needs_word_boundary("db_password=hunter2", "REDACTED")]
#[case::plain_text_is_kept(r#"{"message": "rejected"}"#, r#"{"message": "rejected"}"#)]
fn redacts_the_same_spans_as_the_python_patterns(#[case] input: &str, #[case] expected: &str) {
assert_eq!(
SecretRedactor::new(DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH).redact(input),
expected
);
}
#[test]
fn sk_threshold_follows_the_minimum_custom_key_length() {
let redactor = SecretRedactor::new(8);
assert_eq!(redactor.redact("sk-abcde"), REDACTED);
assert_eq!(redactor.redact("sk-abcd"), "sk-abcd");
}
}

View file

@ -28,6 +28,7 @@
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
"web-fetch-2025-09-10": "web-fetch-2025-09-10",
"web-search-2025-03-05": "web-search-2025-03-05"
@ -59,6 +60,7 @@
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": null,
"token-efficient-tools-2025-02-19": null,
"web-fetch-2025-09-10": "web-fetch-2025-09-10",
"web-search-2025-03-05": "web-search-2025-03-05"
@ -90,6 +92,7 @@
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"token-efficient-tools-2025-02-19": null,
"tool-search-tool-2025-10-19": null,
"web-fetch-2025-09-10": null,
@ -122,6 +125,7 @@
"structured-outputs-2025-11-13": null,
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"token-efficient-tools-2025-02-19": null,
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
"web-fetch-2025-09-10": null,
@ -154,6 +158,7 @@
"structured-outputs-2025-11-13": null,
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"token-efficient-tools-2025-02-19": null,
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
"web-fetch-2025-09-10": null,
@ -187,6 +192,7 @@
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
"web-fetch-2025-09-10": "web-fetch-2025-09-10",
"web-search-2025-03-05": "web-search-2025-03-05"

View file

@ -183,6 +183,9 @@ MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIME
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH: Final = 8
MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE: Final = 4096
# Allowlist of commands permitted for MCP stdio transport.
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.
@ -1573,6 +1576,21 @@ ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS: Final = {
# Works for all LLM pass-through endpoints (Vertex AI, Anthropic, Bedrock, etc.)
PASS_THROUGH_HEADER_PREFIX: Final = "x-pass-"
AZURE_SPEECH_CUSTOM_LLM_PROVIDER: Final = "azure_speech"
AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX: Final = "/azure_speech"
AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX: Final = "/speech/"
AZURE_SPEECH_BATCH_PATH_PREFIX: Final = "/speechtotext/"
AZURE_SPEECH_FAST_TRANSCRIPTION_PATH: Final = "/speechtotext/transcriptions:transcribe"
AZURE_SPEECH_STT_DOMAIN: Final = "stt.speech.microsoft.com"
AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN: Final = "api.cognitive.microsoft.com"
AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER: Final = "Ocp-Apim-Subscription-Key"
AZURE_SPEECH_SHORT_AUDIO_MODEL: Final = "short-audio"
AZURE_SPEECH_BATCH_MODEL: Final = "batch-transcription"
AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL: Final = "fast-transcription"
AZURE_SPEECH_PRICING_MODEL: Final = "azure/speech/azure-stt"
AZURE_SPEECH_TICKS_PER_SECOND: Final = 10_000_000
AZURE_SPEECH_MILLISECONDS_PER_SECOND: Final = 1_000
BASE_MCP_ROUTE: Final = "/mcp"
TRANSCRIBE_JOB_POLLING_INTERVAL_SECONDS: Final = 10.0
@ -1657,6 +1675,11 @@ LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS: Final = int(
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE: Final = int(
os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000)
)
LOGIN_THROTTLE_CACHE_KEY_PREFIX: Final = "login_fail"
LOGIN_THROTTLE_UNKNOWN_SOURCE: Final = "unknown"
LOGIN_THROTTLE_MAX_TRACKED_COUNTERS: Final = 20_000
LOGIN_THROTTLE_MAX_TRACKED_BLOCKS: Final = 10_000
LOGIN_THROTTLE_NOT_BLOCKED: Final = (0, 0)
LITELLM_PROXY_ADMIN_NAME: Final = "default_user_id"
LITELLM_PROXY_BUDGET_NAME: Final = "litellm-proxy-budget"
GLOBAL_PROXY_SPEND_CACHE_KEY: Final = f"{LITELLM_PROXY_ADMIN_NAME}:spend"
@ -2049,6 +2072,7 @@ MCP_SPEND_LOG_MODEL_PREFIX: Final[str] = "MCP: "
PTU_SENTINEL_API_KEY: Final[str] = "__ptu_flat_cost__"
PTU_ROLLUP_JOB_ID: Final[str] = "ptu_flat_cost_rollup_job"
PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900
USAGE_TOP_API_KEYS_LIMIT: Final[int] = int(os.getenv("USAGE_TOP_API_KEYS_LIMIT", "100"))
# Furthest back the catch-up pass looks for unpriced PTU days when a deployment
# declares no ptu_effective_from, bounding the scan for an open-ended window.
PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90

View file

@ -846,7 +846,12 @@ def image_edit(
local_vars.update(kwargs)
# Get ImageEditOptionalRequestParams with only valid parameters
image_edit_optional_params: Final[ImageEditOptionalRequestParams] = (
_get_ImageEditRequestUtils().get_requested_image_edit_optional_param(local_vars)
_get_ImageEditRequestUtils().get_requested_image_edit_optional_param(
local_vars,
provider_supported_params=frozenset(
image_edit_provider_config.get_supported_openai_params(model)
).intersection(non_default_params),
)
)
# Get optional parameters for the responses API
image_edit_request_params: Final[dict] = _get_ImageEditRequestUtils().get_optional_params_image_edit(
@ -857,7 +862,7 @@ def image_edit(
additional_drop_params=kwargs.get("additional_drop_params"),
)
if (
if image_edit_provider_config.use_multipart_form_data() and (
custom_llm_provider == "openai"
or custom_llm_provider == "azure"
or custom_llm_provider in litellm.openai_compatible_providers

View file

@ -1,4 +1,4 @@
from collections.abc import Mapping
from collections.abc import Collection, Mapping
from io import BufferedReader, BytesIO
from typing import Any, Final, cast, get_type_hints
@ -63,6 +63,7 @@ class ImageEditRequestUtils:
@staticmethod
def get_requested_image_edit_optional_param(
params: Mapping[str, object],
provider_supported_params: Collection[str] = (),
) -> ImageEditOptionalRequestParams:
"""
Filter parameters to only include those defined in ImageEditOptionalRequestParams.
@ -73,7 +74,9 @@ class ImageEditRequestUtils:
Returns:
ImageEditOptionalRequestParams instance with only the valid parameters
"""
valid_keys: Final = get_type_hints(ImageEditOptionalRequestParams).keys()
valid_keys: Final = frozenset(get_type_hints(ImageEditOptionalRequestParams)) | frozenset(
provider_supported_params
)
filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None}
return cast(ImageEditOptionalRequestParams, filtered_params)

View file

@ -1848,6 +1848,9 @@ class CostCalculatorUtils:
return azure_ai_image_cost_calculator(
model=model,
image_response=completion_response,
size=resolved_size,
n=resolved_n,
optional_params=optional_params,
)
elif custom_llm_provider == litellm.LlmProviders.FAL_AI.value:
from litellm.llms.fal_ai.cost_calculator import (

View file

@ -1,5 +1,7 @@
import base64
from collections.abc import Mapping, Sequence
from io import BufferedReader
from types import MappingProxyType
from typing import Any, Final
from httpx._types import RequestFiles
@ -24,21 +26,12 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
Azure AI Foundry FLUX 2 image edit config
Supports FLUX 2 models (e.g., flux.2-pro) for image editing.
Uses the same /providers/blackforestlabs/v1/flux-2-pro endpoint as image generation,
Uses the model-specific /providers/blackforestlabs/v1/flux-2-* endpoint as image generation,
with the image passed as base64 in JSON body.
"""
def get_supported_openai_params(self, model: str) -> list:
"""
FLUX 2 supports a subset of OpenAI image edit params
"""
return [
"prompt",
"image",
"model",
"n",
"size",
]
return AzureFoundryFluxImageGenerationConfig().get_supported_openai_params(model)
def map_openai_params(
self,
@ -50,14 +43,14 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
Map OpenAI params to FLUX 2 params.
FLUX 2 uses the same param names as OpenAI for supported params.
"""
mapped_params: Final[dict[str, Any]] = {}
supported_params: Final = self.get_supported_openai_params(model)
for key, value in dict(image_edit_optional_params).items():
if key in supported_params and value is not None:
mapped_params[key] = value
return mapped_params
return AzureFoundryFluxImageGenerationConfig().map_openai_params(
non_default_params=MappingProxyType(
{key: value for key, value in image_edit_optional_params.items() if value is not None}
),
optional_params=MappingProxyType({}),
model=model,
drop_params=drop_params,
)
def use_multipart_form_data(self) -> bool:
"""FLUX 2 uses JSON requests, not multipart/form-data."""
@ -90,7 +83,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
self,
model: str,
prompt: str | None,
image: FileTypes | None,
image: FileTypes | Sequence[FileTypes] | None,
image_edit_optional_request_params: dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
@ -107,29 +100,29 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
if image is None:
raise ValueError("FLUX 2 image edit requires an image.")
image_b64: Final = self._convert_image_to_base64(image)
images: Final = tuple(image) if isinstance(image, list) else (image,)
if not images:
raise ValueError("FLUX 2 image edit requires at least one image.")
max_reference_images: Final = 10 if "flex" in model.lower() else 8
if len(images) > max_reference_images:
raise ValueError(f"{model} supports at most {max_reference_images} reference images.")
# Build request body with required params
reference_images: Final[Mapping[str, str]] = MappingProxyType(
{
"input_image" if index == 1 else f"input_image_{index}": self._convert_image_to_base64(reference_image)
for index, reference_image in enumerate(images, start=1)
}
)
request_body: Final[dict[str, Any]] = {
"prompt": prompt,
"image": image_b64,
"model": model,
**reference_images,
**image_edit_optional_request_params,
}
# Add mapped optional params (already filtered by map_openai_params)
request_body.update(image_edit_optional_request_params)
# Return JSON body and empty files list (FLUX 2 doesn't use multipart)
return request_body, []
def _convert_image_to_base64(self, image: Any) -> str:
"""Convert image file to base64 string"""
# Handle list of images (take first one)
if isinstance(image, list):
if len(image) == 0:
raise ValueError("Empty image list provided")
image = image[0]
if isinstance(image, BufferedReader):
image_bytes = image.read()
image.seek(0) # Reset file pointer for potential reuse
@ -151,7 +144,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
"""
Constructs a complete URL for Azure AI Foundry FLUX 2 image edits.
Uses the same /providers/blackforestlabs/v1/flux-2-pro endpoint as image generation.
Uses the same model-specific BFL provider endpoint as image generation.
"""
api_base = AzureFoundryModelInfo.get_api_base(api_base)

View file

@ -1,3 +1,4 @@
from collections.abc import Mapping
from typing import Any, Final
import litellm
@ -10,6 +11,9 @@ from litellm.types.utils import ImageResponse
def cost_calculator(
model: str,
image_response: Any,
size: str | None = None,
n: int | None = None,
optional_params: Mapping[str, object] | None = None,
) -> float:
"""
Azure AI image generation cost calculator
@ -28,10 +32,29 @@ def cost_calculator(
if token_based_cost is not None:
return token_based_cost
num_images: Final = n if n is not None else len(image_response.data or ())
output_cost_per_image: Final[float] = _model_info.get("output_cost_per_image") or 0.0
num_images: int = 0
if image_response.data:
num_images = len(image_response.data)
return output_cost_per_image * num_images
if output_cost_per_image:
return output_cost_per_image * num_images
model_cost: Final = litellm.model_cost[_model_info["key"]]
input_cost_per_pixel: Final[float] = model_cost.get("input_cost_per_pixel") or 0.0
if input_cost_per_pixel:
from litellm.cost_calculator import default_image_cost_calculator
width: Final = optional_params.get("width") if optional_params else None
height: Final = optional_params.get("height") if optional_params else None
pixel_size: Final = (
f"{width}x{height}"
if type(width) is int and type(height) is int and width > 0 and height > 0
else size or image_response.size
)
return default_image_cost_calculator(
model=_model_info["key"],
custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
size=pixel_size,
n=num_images,
)
return 0.0
raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}")

View file

@ -1,18 +1,22 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from litellm.exceptions import BadRequestError, UnsupportedParamsError
from litellm.llms.openai.image_generation import GPTImageGenerationConfig
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
FLUX2_DROPPED_OPENAI_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = (
"background",
"moderation",
"output_compression",
"quality",
"user",
)
class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
"""
Azure Foundry flux image generation config
From manual testing it follows the gpt-image-1 image generation config
(Azure Foundry does not have any docs on supported params at the time of writing)
From our test suite - following GPTImageGenerationConfig is working for this model
"""
"""Azure Foundry BFL API configuration for FLUX image generation."""
@staticmethod
def get_flux2_image_generation_url(
@ -25,11 +29,11 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
FLUX 2 models on Azure AI use a different URL pattern than standard Azure OpenAI:
- Standard: /openai/deployments/{model}/images/generations
- FLUX 2: /providers/blackforestlabs/v1/flux-2-pro
- FLUX 2: /providers/blackforestlabs/v1/{model-path}
Args:
api_base: Base URL (e.g., https://litellm-ci-cd-prod.services.ai.azure.com)
model: Model name (e.g., flux.2-pro)
model: Model name (e.g., FLUX.2-flex or FLUX.2-pro)
api_version: API version (e.g., preview)
Returns:
@ -47,9 +51,8 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
return api_base
return f"{api_base}?api-version={api_version}"
# Construct the FLUX 2 provider path
# Model name flux.2-pro maps to endpoint flux-2-pro
return f"{api_base}/providers/blackforestlabs/v1/flux-2-pro?api-version={api_version}"
provider_model_path: Final = AzureFoundryFluxImageGenerationConfig.get_flux2_provider_model_path(model)
return f"{api_base}/providers/blackforestlabs/v1/{provider_model_path}?api-version={api_version}"
@staticmethod
def is_flux2_model(model: str) -> bool:
@ -64,3 +67,90 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
"""
model_lower: Final = model.lower().replace(".", "-").replace("_", "-")
return "flux-2" in model_lower or "flux2" in model_lower
@staticmethod
def get_flux2_provider_model_path(model: str) -> str:
normalized_model: Final = model.lower().replace(".", "-").replace("_", "-")
return "flux-2-flex" if "flux-2-flex" in normalized_model else "flux-2-pro"
def get_supported_openai_params( # mutable-ok: inherited config contract returns a list
self, model: str
) -> list[OpenAIImageGenerationOptionalParams]:
if not self.is_flux2_model(model):
return super().get_supported_openai_params(model)
return [ # mutable-ok: BaseImageGenerationConfig requires a list
"n",
"size",
"output_format",
"seed",
"safety_tolerance",
"aspect_ratio",
"width",
"height",
"num_images",
"guidance",
"steps",
*FLUX2_DROPPED_OPENAI_PARAMS,
]
@staticmethod
def _map_parameter(name: str, value: object, model: str) -> tuple[tuple[str, object], ...]:
if name in FLUX2_DROPPED_OPENAI_PARAMS:
return ()
if isinstance(value, str):
if name in ("n", "num_images", "width", "height", "steps", "seed", "safety_tolerance"):
return (("num_images" if name == "n" else name, int(value)),)
if name == "guidance":
return ((name, float(value)),)
if name == "n":
return (("num_images", value),)
if name != "size":
return ((name, value),)
if str(value).lower() == "auto":
return ()
try:
width, height = (int(dimension) for dimension in str(value).lower().split("x"))
except (TypeError, ValueError):
raise BadRequestError(
message=f"Invalid size format '{value}'. Expected 'WxH', for example '1024x1024'.",
model=model,
llm_provider="azure_ai",
)
return (("width", width), ("height", height))
def map_openai_params(
self,
non_default_params: Mapping[str, object],
optional_params: Mapping[str, object],
model: str,
drop_params: bool,
) -> dict[str, object]: # mutable-ok: inherited config contract returns a dict
if not self.is_flux2_model(model):
return super().map_openai_params(
non_default_params=dict(non_default_params),
optional_params=dict(optional_params),
model=model,
drop_params=drop_params,
)
supported_params: Final = self.get_supported_openai_params(model)
unsupported_params: Final = tuple(name for name in non_default_params if name not in supported_params)
if unsupported_params and not drop_params:
raise UnsupportedParamsError(
message=(
f"Parameters {unsupported_params} are not supported for model {model}. "
f"Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters."
),
model=model,
llm_provider="azure_ai",
)
mapped_params: Final[Mapping[str, object]] = MappingProxyType(
{
mapped_name: mapped_value
for name, value in non_default_params.items()
if name in supported_params
for mapped_name, mapped_value in self._map_parameter(name, value, model)
}
)
return {**optional_params, **mapped_params} # mutable-ok: inherited config contract returns a dict

View file

@ -82,7 +82,14 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig):
)
# set optional params
image_response.size = image_response.size or optional_params.get("size", "1024x1024")
width: Final = optional_params.get("width")
height: Final = optional_params.get("height")
requested_size: Final = (
f"{width}x{height}"
if isinstance(width, int) and isinstance(height, int)
else optional_params.get("size", "1024x1024")
)
image_response.size = image_response.size or requested_size
image_response.quality = image_response.quality or optional_params.get("quality", "high")
image_response.output_format = image_response.output_format or optional_params.get("output_format", "png")

View file

@ -101,7 +101,8 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig):
extra_body: dict[str, object] | None = None,
) -> tuple[str, dict]:
encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id")
url: Final = f"{api_base}/{encoded_vector_store_id}/search"
base_url, query_separator, query_string = api_base.partition("?")
url: Final = f"{base_url}/{encoded_vector_store_id}/search{query_separator}{query_string}"
typed_request_body: Final = VectorStoreSearchRequest(
query=query,
filters=vector_store_search_optional_params.get("filters", None),

View file

@ -10820,6 +10820,25 @@
"/v1/images/generations"
]
},
"azure_ai/FLUX.2-flex": {
"input_cost_per_pixel": 5e-08,
"litellm_provider": "azure_ai",
"max_input_tokens": 32000,
"max_tokens": 32000,
"mode": "image_generation",
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/",
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"image"
]
},
"azure_ai/FW-DeepSeek-V3.2": {
"deprecation_date": "2027-07-01",
"cache_read_input_token_cost": 3.1e-07,
@ -16690,6 +16709,46 @@
"supports_tool_choice": true,
"supports_vision": true
},
"dashscope/qwen3.8-flash": {
"cache_creation_input_token_cost": 2e-07,
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "dashscope",
"max_input_tokens": 991808,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4.7e-07,
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true,
"supports_web_search": true
},
"dashscope/qwen3.8-omni-flash": {
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "dashscope",
"max_input_tokens": 991808,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4.7e-07,
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true,
"supports_web_search": true
},
"dashscope/qwq-plus": {
"input_cost_per_token": 8e-07,
"litellm_provider": "dashscope",
@ -18594,6 +18653,46 @@
"supports_tool_choice": true,
"supports_vision": true
},
"qwen_ai_platform/qwen3.8-flash": {
"cache_creation_input_token_cost": 2e-07,
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "qwen_ai_platform",
"max_input_tokens": 991808,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4.7e-07,
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true,
"supports_web_search": true
},
"qwen_ai_platform/qwen3.8-omni-flash": {
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "qwen_ai_platform",
"max_input_tokens": 991808,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4.7e-07,
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true,
"supports_web_search": true
},
"qwen_ai_platform/qwq-plus": {
"input_cost_per_token": 8e-07,
"litellm_provider": "qwen_ai_platform",
@ -22090,8 +22189,8 @@
"embed-english-light-v3.0": {
"input_cost_per_token": 1e-07,
"litellm_provider": "cohere",
"max_input_tokens": 1024,
"max_tokens": 1024,
"max_input_tokens": 512,
"max_tokens": 512,
"mode": "embedding",
"output_cost_per_token": 0.0
},
@ -22108,8 +22207,8 @@
"input_cost_per_image": 0.0001,
"input_cost_per_token": 1e-07,
"litellm_provider": "cohere",
"max_input_tokens": 1024,
"max_tokens": 1024,
"max_input_tokens": 512,
"max_tokens": 512,
"metadata": {
"notes": "'supports_image_input' is a deprecated field. Use 'supports_embedding_image_input' instead."
},
@ -22130,8 +22229,8 @@
"embed-multilingual-v3.0": {
"input_cost_per_token": 1e-07,
"litellm_provider": "cohere",
"max_input_tokens": 1024,
"max_tokens": 1024,
"max_input_tokens": 512,
"max_tokens": 512,
"mode": "embedding",
"output_cost_per_token": 0.0,
"supports_embedding_image_input": true
@ -22139,8 +22238,8 @@
"embed-multilingual-light-v3.0": {
"input_cost_per_token": 0.0001,
"litellm_provider": "cohere",
"max_input_tokens": 1024,
"max_tokens": 1024,
"max_input_tokens": 512,
"max_tokens": 512,
"mode": "embedding",
"output_cost_per_token": 0.0,
"supports_embedding_image_input": true
@ -57910,14 +58009,14 @@
"supports_tool_choice": true
},
"bedrock_mantle/openai.gpt-5.6-sol": {
"input_cost_per_token": 5.5e-06,
"input_cost_per_token_above_272k_tokens": 1.1e-05,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_272k_tokens": 1.375e-05,
"cache_read_input_token_cost": 5.5e-07,
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"input_cost_per_token": 4.4e-06,
"input_cost_per_token_above_272k_tokens": 8.8e-06,
"cache_creation_input_token_cost": 5.5e-06,
"cache_creation_input_token_cost_above_272k_tokens": 1.1e-05,
"cache_read_input_token_cost": 4.4e-07,
"cache_read_input_token_cost_above_272k_tokens": 8.8e-07,
"output_cost_per_token": 2.2e-05,
"output_cost_per_token_above_272k_tokens": 3.3e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
@ -65816,9 +65915,9 @@
"supports_web_search": false
},
"openrouter/z-ai/glm-5.3": {
"input_cost_per_token": 1.4e-06,
"output_cost_per_token": 4.4e-06,
"cache_read_input_token_cost": 2.6e-07,
"input_cost_per_token": 9.1e-07,
"output_cost_per_token": 2.86e-06,
"cache_read_input_token_cost": 1.69e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1310720,
"max_output_tokens": 943717,
@ -70539,14 +70638,14 @@
"supports_web_search": false
},
"openrouter/~deepseek/deepseek-flash-latest": {
"cache_read_input_token_cost": 1.5e-08,
"input_cost_per_token": 1.5e-07,
"cache_read_input_token_cost": 4.2e-09,
"input_cost_per_token": 1.4e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 393216,
"max_tokens": 393216,
"mode": "chat",
"output_cost_per_token": 6e-07,
"output_cost_per_token": 4.2e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -70831,14 +70930,14 @@
"supports_web_search": false
},
"openrouter/~z-ai/glm-latest": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost": 1.46625e-07,
"input_cost_per_token": 9e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1310720,
"max_output_tokens": 235929,
"max_tokens": 235929,
"mode": "chat",
"output_cost_per_token": 3e-06,
"output_cost_per_token": 2.805e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -73982,14 +74081,14 @@
"supports_web_search": false
},
"openrouter/tencent/hy3": {
"cache_read_input_token_cost": 3.3e-08,
"input_cost_per_token": 1.32e-07,
"cache_read_input_token_cost": 2.0625e-08,
"input_cost_per_token": 8.25e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 262144,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 5.28e-07,
"output_cost_per_token": 3.3e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,

View file

@ -1,8 +1,15 @@
import sys
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from pydantic import TypeAdapter
DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS: Final = 600.0
_SECONDS: Final = TypeAdapter(float)
_NO_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
def resolve_pass_through_request_timeout(
endpoint_timeout: float | None = None,
@ -31,26 +38,41 @@ def resolve_pass_through_request_timeout(
def resolve_llm_passthrough_timeout(
kwargs: dict | None = None,
litellm_params: dict | None = None,
router_timeout: float | None = None,
kwargs: Mapping[str, object] | None = None,
litellm_params: Mapping[str, object] | None = None,
router_timeout: float | str | None = None,
router_stream_timeout: float | str | None = None,
) -> float:
"""
Resolve upstream httpx timeout for SDK native passthrough (e.g. Bedrock /converse).
Resolve upstream httpx timeout for SDK native passthrough (e.g. Bedrock /converse,
Anthropic /v1/messages).
Precedence: kwargs timeout/request_timeout -> litellm_params timeout/request_timeout
-> router_timeout -> general_settings.pass_through_request_timeout -> 600s default.
Non-streaming precedence: kwargs timeout/request_timeout -> litellm_params
timeout/request_timeout -> router_timeout -> general_settings.pass_through_request_timeout
-> 600s default.
Streaming (``kwargs["stream"]`` truthy) resolves ``stream_timeout`` at every level before
any generic timeout, matching ``Router._get_stream_timeout`` on the completion route:
kwargs stream_timeout -> litellm_params stream_timeout -> router_stream_timeout, then the
non-streaming chain above.
Only the first set value is validated as seconds, so a value in a lower-precedence
field never fails the call.
"""
kwargs = kwargs or {}
litellm_params = litellm_params or {}
for source in (kwargs, litellm_params):
for key in ("timeout", "request_timeout"):
val = source.get(key)
if val is not None:
return float(val)
if router_timeout is not None:
return float(router_timeout)
return resolve_pass_through_request_timeout()
request: Final = kwargs if kwargs is not None else _NO_PARAMS
deployment: Final = litellm_params if litellm_params is not None else _NO_PARAMS
stream_candidates: Final = (
(request.get("stream_timeout"), deployment.get("stream_timeout"), router_stream_timeout)
if request.get("stream")
else ()
)
candidates: Final = (
*stream_candidates,
request.get("timeout"),
request.get("request_timeout"),
deployment.get("timeout"),
deployment.get("request_timeout"),
router_timeout,
)
winner: Final = next((val for val in candidates if val is not None), None)
return resolve_pass_through_request_timeout() if winner is None else _SECONDS.validate_python(winner)

View file

@ -17,6 +17,7 @@ _PASS_THROUGH_PROTECTED_HEADERS: Final[frozenset] = frozenset(
"api-key",
"x-api-key",
"x-goog-api-key",
"ocp-apim-subscription-key",
"host",
"content-length",
"accept-encoding",

View file

@ -0,0 +1,38 @@
"""Per-worker cache of stored BYOK credentials, keyed so peer workers can evict it over the auth cache pub/sub."""
from dataclasses import dataclass
from typing import Final
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE, MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS
_CACHE_KEY_PREFIX: Final = "mcp_byok_credential"
@dataclass(frozen=True, slots=True)
class CachedByokCredential:
credential: str | None
byok_credential_cache: Final = InMemoryCache(
max_size_in_memory=MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE,
default_ttl=MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS,
)
def byok_credential_cache_key(user_id: str, server_id: str) -> str:
return f"{_CACHE_KEY_PREFIX}:{user_id}:{server_id}"
def get_cached_byok_credential(user_id: str, server_id: str) -> CachedByokCredential | None:
cached: Final = byok_credential_cache.get_cache( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # InMemoryCache is untyped
byok_credential_cache_key(user_id, server_id)
)
return cached if isinstance(cached, CachedByokCredential) else None
def cache_byok_credential(user_id: str, server_id: str, credential: str | None) -> None:
byok_credential_cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
byok_credential_cache_key(user_id, server_id),
CachedByokCredential(credential=credential),
)

View file

@ -865,7 +865,7 @@ async def byok_token(
_invalidate_byok_cred_cache,
)
_invalidate_byok_cred_cache(user_id, server_id)
await _invalidate_byok_cred_cache(user_id, server_id)
except Exception as exc:
verbose_proxy_logger.error(
"byok_token: failed to store user credential for user=%s server=%s: %s",

View file

@ -24,6 +24,7 @@ from litellm.proxy._types import (
MCPApprovalStatus,
MCPEnvVar,
MCPEnvVarScope,
MCPServerUserCredentialListItem,
MCPSubmissionsSummary,
NewMCPServerRequest,
SpecialMCPServerName,
@ -1504,6 +1505,37 @@ async def get_user_oauth_credential(
return _parse_oauth_payload(decoded)
def _server_user_credential_item(
row: "prisma_db_models.LiteLLM_MCPUserCredentials",
) -> MCPServerUserCredentialListItem:
oauth_payload: Final = _decode_oauth_payload(row.credential_b64)
if oauth_payload is None:
return MCPServerUserCredentialListItem(
user_id=row.user_id,
credential_type="byok",
updated_at=row.updated_at.isoformat(),
)
return MCPServerUserCredentialListItem(
user_id=row.user_id,
credential_type="oauth2",
expires_at=oauth_payload.get("expires_at"),
connected_at=oauth_payload.get("connected_at"),
updated_at=row.updated_at.isoformat(),
)
async def list_server_user_credentials(
prisma_client: PrismaClient,
server_id: str,
) -> tuple[MCPServerUserCredentialListItem, ...]:
"""Every user's stored credential for one server, typed but without the secret, for admins."""
rows: Final = await _db_find_user_credential_rows(
prisma_client,
{"server_id": server_id}, # mutable-ok: prisma where-inputs must be plain dicts
)
return tuple(_server_user_credential_item(row) for row in rows)
async def list_user_oauth_credentials(
prisma_client: PrismaClient,
user_id: str,

View file

@ -295,12 +295,15 @@ class MCPPerUserTokenCache:
)
async def delete(self, user_id: str, server_id: str) -> None:
"""Invalidate the cached token (removes from both in-memory and Redis layers)."""
"""Invalidate the cached token in Redis, here, and in every peer worker's in-memory layer."""
try:
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( # noqa: PLC0415 # proxy import cycle
evict_and_broadcast,
)
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
key: Final = self._cache_key(user_id, server_id)
await user_api_key_cache.async_delete_cache(key)
await evict_and_broadcast((key,), user_api_key_cache)
except Exception as exc:
verbose_logger.debug(
"MCPPerUserTokenCache.delete failed for user=%s server=%s: %s",

View file

@ -28,7 +28,10 @@ from starlette.types import Message, Receive, Scope, Send
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
from litellm.constants import (
MAXIMUM_TRACEBACK_LINES_TO_LOG,
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -38,6 +41,12 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
_is_mcp_admitted_user_subject,
)
from litellm.proxy._experimental.mcp_server.byok_credential_cache import (
byok_credential_cache,
byok_credential_cache_key,
cache_byok_credential,
get_cached_byok_credential,
)
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
get_request_base_url,
)
@ -82,6 +91,9 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
publish_auth_cache_invalidation,
)
from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
get_chain_id_from_headers,
@ -91,6 +103,7 @@ from litellm.types.mcp import (
MCPGatewaySession,
MCPGatewaySessionGroupCount,
MCPGatewaySessionsResponse,
MCPGatewaySessionsTerminateResponse,
MCPSpecVersion,
)
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
@ -102,13 +115,6 @@ if TYPE_CHECKING:
from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload
# Short-lived in-memory cache for BYOK credentials.
# Keyed by (user_id, server_id); value is (credential_or_None, monotonic_timestamp).
# Storing the credential value (not just a bool) means _get_byok_credential and
# _check_byok_credential share a single DB round-trip per TTL window.
_byok_cred_cache: Final[dict[tuple[str, str], tuple[str | None, float]]] = {}
_BYOK_CRED_CACHE_TTL: Final = 60 # seconds
_BYOK_CRED_CACHE_MAX_SIZE: Final = 4096 # cap to prevent unbounded growth
_STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS: Final = 30 * 60
# Upper bound on concurrent stateful sessions a single caller may hold. Each
# `initialize` creates a session that survives until the idle timeout, so
@ -127,20 +133,11 @@ _MCP_TRANSPORT_SPAN_SCOPE_KEY: Final = "litellm_otel_transport_span"
_MCP_DESTINATIONS_SCOPE_KEY: Final = "litellm_otel_request_destinations"
def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
"""Remove a (user_id, server_id) entry from the BYOK credential cache.
Call this after storing or deleting a credential so subsequent calls
see the fresh value rather than a stale cached result.
"""
_byok_cred_cache.pop((user_id, server_id), None)
def _write_byok_cred_cache(user_id: str, server_id: str, credential: str | None) -> None:
"""Write a credential value to the cache, evicting all entries if at capacity."""
if len(_byok_cred_cache) >= _BYOK_CRED_CACHE_MAX_SIZE:
_byok_cred_cache.clear()
_byok_cred_cache[(user_id, server_id)] = (credential, time.monotonic())
async def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
"""Drop a stored-or-deleted BYOK credential from this worker's cache and from every peer worker's."""
cache_key: Final = byok_credential_cache_key(user_id, server_id)
byok_credential_cache.delete_cache(cache_key)
await publish_auth_cache_invalidation(cache_key=cache_key)
# Check if MCP is available
@ -618,6 +615,7 @@ if MCP_AVAILABLE:
_stateful_session_locks: Final[dict[str, asyncio.Lock]] = {}
_stateful_session_active_request_counts: Final[dict[str, int]] = {}
_stateful_session_client_info: Final[dict[str, Implementation]] = {} # mutable-ok: cleared on session teardown
_admin_terminated_session_ids: Final[dict[str, float]] = {} # mutable-ok: admin-closed id -> last replay
class _TerminableTransport(Protocol):
async def terminate(self) -> None: ...
@ -689,6 +687,7 @@ if MCP_AVAILABLE:
for session_id in list(_stateful_session_auth_context_last_seen):
if session_id not in _stateful_session_auth_contexts:
_remove_stateful_session_tracking(session_id)
_forget_expired_admin_terminated_session_ids(now)
async def _enforce_stateful_session_cap_for_owner(owner: str) -> bool:
"""
@ -2811,35 +2810,28 @@ if MCP_AVAILABLE:
mcp_server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
) -> str | None:
"""Retrieve the stored BYOK credential for a user+server pair.
Uses the shared _byok_cred_cache to avoid a DB round-trip on every
tool call within the TTL window.
"""
"""Retrieve the stored BYOK credential for a user+server pair, served from the worker cache within its TTL."""
if not mcp_server.is_byok:
return None
user_id: Final = (user_api_key_auth.user_id if user_api_key_auth else None) or ""
if not user_id:
return None
cache_key: Final = (user_id, mcp_server.server_id)
cached: Final = _byok_cred_cache.get(cache_key)
cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id)
if cached is not None:
credential, ts = cached
if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL:
return credential
return cached.credential
from litellm.proxy._experimental.mcp_server.db import get_user_credential
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return None
credential = await get_user_credential(
credential: Final = await get_user_credential(
prisma_client=prisma_client,
user_id=user_id,
server_id=mcp_server.server_id,
)
_write_byok_cred_cache(user_id, mcp_server.server_id, credential)
cache_byok_credential(user_id, mcp_server.server_id, credential)
return credential
async def _check_byok_credential(
@ -2868,27 +2860,23 @@ if MCP_AVAILABLE:
headers={"WWW-Authenticate": get_byok_www_authenticate()},
)
# Check shared credential cache before hitting the DB.
cache_key: Final = (user_id, mcp_server.server_id)
cached: Final = _byok_cred_cache.get(cache_key)
cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id)
if cached is not None:
cached_cred, ts = cached
if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL:
if cached_cred is None:
raise HTTPException(
status_code=401,
detail={
"error": "byok_auth_required",
"server_id": mcp_server.server_id,
"server_name": mcp_server.server_name or mcp_server.name,
"message": (
"No stored credential found for this BYOK server. "
"Complete the OAuth authorization flow to provide your API key."
),
},
headers={"WWW-Authenticate": get_byok_www_authenticate()},
)
return
if cached.credential is None:
raise HTTPException(
status_code=401,
detail={
"error": "byok_auth_required",
"server_id": mcp_server.server_id,
"server_name": mcp_server.server_name or mcp_server.name,
"message": (
"No stored credential found for this BYOK server. "
"Complete the OAuth authorization flow to provide your API key."
),
},
headers={"WWW-Authenticate": get_byok_www_authenticate()},
)
return
from litellm.proxy._experimental.mcp_server.db import get_user_credential
from litellm.proxy.proxy_server import prisma_client
@ -2912,7 +2900,7 @@ if MCP_AVAILABLE:
user_id=user_id,
server_id=mcp_server.server_id,
)
_write_byok_cred_cache(user_id, mcp_server.server_id, credential)
cache_byok_credential(user_id, mcp_server.server_id, credential)
if credential is None:
raise HTTPException(
status_code=401,
@ -3850,7 +3838,7 @@ if MCP_AVAILABLE:
client_info: Final = _stateful_session_client_info.get(session_id)
key_auth: Final = auth_user.user_api_key_auth
return MCPGatewaySession(
session_id_prefix=session_id[:8],
session_id_prefix=session_id[:MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH],
client_name=client_info.name if client_info is not None else None,
client_version=client_info.version if client_info is not None else None,
user_id=key_auth.user_id if key_auth is not None else None,
@ -3885,6 +3873,72 @@ if MCP_AVAILABLE:
sessions=sessions,
)
def _session_matches_admin_selector(
session_id: str,
auth_user: MCPAuthenticatedUser,
session_id_prefix: str | None,
user_id: str | None,
) -> bool:
if session_id_prefix is not None and not session_id.startswith(session_id_prefix):
return False
if user_id is None:
return True
key_auth: Final = auth_user.user_api_key_auth
return key_auth is not None and key_auth.user_id == user_id
def _forget_expired_admin_terminated_session_ids(now: float) -> None:
for session_id in [
session_id
for session_id, last_replayed in _admin_terminated_session_ids.items()
if now - last_replayed >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS
]:
del _admin_terminated_session_ids[session_id]
def _is_admin_terminated_session_id(session_id: str, now: float) -> bool:
last_replayed: Final = _admin_terminated_session_ids.get(session_id)
if last_replayed is None:
return False
if now - last_replayed >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS:
del _admin_terminated_session_ids[session_id]
return False
_admin_terminated_session_ids[session_id] = now
return True
async def terminate_mcp_gateway_sessions(
*,
session_id_prefix: str | None = None,
user_id: str | None = None,
) -> MCPGatewaySessionsTerminateResponse:
"""Force-close every live stateful session on this worker matching the selector.
The transport is terminated (open streams close), all per-session
tracking is dropped, and the id is remembered so a client that keeps
sending it receives 404 and has to ``initialize`` again, which re-runs
admission. Only sessions held by this worker process are affected.
"""
now: Final = time.monotonic()
_forget_expired_admin_terminated_session_ids(now)
server_instances: Final = _stateful_server_instances()
targets: Final = tuple(
(session_id, auth_user)
for session_id, auth_user in tuple(_stateful_session_auth_contexts.items())
if session_id in server_instances
and _session_matches_admin_selector(session_id, auth_user, session_id_prefix, user_id)
)
terminated: Final = tuple(_gateway_session_for(session_id, auth_user, now) for session_id, auth_user in targets)
for session_id, _ in targets:
_admin_terminated_session_ids[session_id] = now
transport = server_instances.pop(session_id, None)
_remove_stateful_session_tracking(session_id)
if transport is not None:
await transport.terminate()
verbose_logger.warning("MCP session '%s' terminated by an administrator.", session_id)
return MCPGatewaySessionsTerminateResponse(
worker_pid=os.getpid(),
terminated_sessions=len(terminated),
sessions=terminated,
)
async def _read_request_body_for_routing(
receive: Receive,
) -> tuple[list[Message], bytes]:
@ -4009,6 +4063,17 @@ if MCP_AVAILABLE:
await success_response(scope, receive, send)
return True
if _is_admin_terminated_session_id(_session_id, time.monotonic()):
terminated_response: Final = JSONResponse(
status_code=404,
content={ # mutable-ok: JSONResponse content must be a plain dict
"error": "Not Found",
"details": "mcp-session-id was terminated by an administrator. Send initialize to start a new session.",
},
)
await terminated_response(scope, receive, send)
return True
# Non-DELETE: strip stale session ID to allow new session creation
verbose_logger.warning(
"MCP session ID '%s' not found in this worker's memory. "

View file

@ -196,6 +196,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
"/assemblyai/",
"/azure/",
"/azure_ai/",
"/azure_speech/",
"/bedrock/",
"/cohere/",
"/comprehendmedical",

View file

@ -3050,6 +3050,18 @@
},
"DailySpendMetadata": {
"properties": {
"api_key_limit": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
],
"description": "When set, api_keys and every api_key_breakdown list at most this many keys, ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.",
"title": "Api Key Limit"
},
"has_more": {
"default": false,
"title": "Has More",
@ -3060,6 +3072,18 @@
"title": "Page",
"type": "integer"
},
"total_api_keys": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
],
"description": "Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key lists are truncated to the highest-spend keys.",
"title": "Total Api Keys"
},
"total_api_requests": {
"default": 0,
"title": "Total Api Requests",
@ -10030,7 +10054,7 @@
},
"unreachable_fallback": {
"default": "fail_closed",
"description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.",
"description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.",
"enum": [
"fail_closed",
"fail_open"
@ -17157,6 +17181,228 @@
]
}
},
"/azure_speech/{endpoint}": {
"delete": {
"description": "Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.\n`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`\nwith the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.\n\nThe body is forwarded byte for byte and the proxy injects its own\n`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key\nand is never forwarded.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)",
"operationId": "azure_speech_proxy_route_azure_speech__endpoint__delete",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Azure Speech Proxy Route",
"tags": [
"llm_passthrough"
]
},
"get": {
"description": "Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.\n`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`\nwith the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.\n\nThe body is forwarded byte for byte and the proxy injects its own\n`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key\nand is never forwarded.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)",
"operationId": "azure_speech_proxy_route_azure_speech__endpoint__get",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Azure Speech Proxy Route",
"tags": [
"llm_passthrough"
]
},
"patch": {
"description": "Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.\n`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`\nwith the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.\n\nThe body is forwarded byte for byte and the proxy injects its own\n`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key\nand is never forwarded.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)",
"operationId": "azure_speech_proxy_route_azure_speech__endpoint__patch",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Azure Speech Proxy Route",
"tags": [
"llm_passthrough"
]
},
"post": {
"description": "Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.\n`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`\nwith the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.\n\nThe body is forwarded byte for byte and the proxy injects its own\n`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key\nand is never forwarded.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)",
"operationId": "azure_speech_proxy_route_azure_speech__endpoint__post",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Azure Speech Proxy Route",
"tags": [
"llm_passthrough"
]
},
"put": {
"description": "Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.\n`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`\nwith the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.\n\nThe body is forwarded byte for byte and the proxy injects its own\n`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key\nand is never forwarded.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)",
"operationId": "azure_speech_proxy_route_azure_speech__endpoint__put",
"parameters": [
{
"in": "path",
"name": "endpoint",
"required": true,
"schema": {
"title": "Endpoint",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Azure Speech Proxy Route",
"tags": [
"llm_passthrough"
]
}
},
"/bedrock/{endpoint}": {
"delete": {
"description": "This is the v1 passthrough for Bedrock.\nV2 is handled by the `/bedrock/v2` endpoint.\n[Docs](https://docs.litellm.ai/docs/pass_through/bedrock)",
@ -27965,6 +28211,32 @@
"title": "MCPGatewaySessionsResponse",
"type": "object"
},
"MCPGatewaySessionsTerminateResponse": {
"description": "Stateful sessions an administrator force-closed on this proxy worker.",
"properties": {
"sessions": {
"items": {
"$ref": "#/components/schemas/MCPGatewaySession"
},
"title": "Sessions",
"type": "array"
},
"terminated_sessions": {
"title": "Terminated Sessions",
"type": "integer"
},
"worker_pid": {
"title": "Worker Pid",
"type": "integer"
}
},
"required": [
"worker_pid",
"terminated_sessions"
],
"title": "MCPGatewaySessionsTerminateResponse",
"type": "object"
},
"MCPOAuthUserCredentialRequest": {
"description": "Stores a user's OAuth2 token for an OpenAPI MCP server.",
"properties": {
@ -28061,6 +28333,56 @@
"title": "MCPOAuthUserCredentialStatus",
"type": "object"
},
"MCPServerUserCredentialListItem": {
"description": "One user's stored credential for an MCP server, as an admin sees it. Never carries the secret.",
"properties": {
"connected_at": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Connected At"
},
"credential_type": {
"enum": [
"oauth2",
"byok"
],
"title": "Credential Type",
"type": "string"
},
"expires_at": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Expires At"
},
"updated_at": {
"title": "Updated At",
"type": "string"
},
"user_id": {
"title": "User Id",
"type": "string"
}
},
"required": [
"user_id",
"credential_type",
"updated_at"
],
"title": "MCPServerUserCredentialListItem",
"type": "object"
},
"MCPSubmissionsSummary": {
"properties": {
"active": {
@ -30237,7 +30559,7 @@
},
"/v1/mcp/server/{server_id}/oauth-user-credential": {
"delete": {
"description": "Revoke the calling user's stored OAuth2 token for an MCP server",
"description": "Revoke the calling user's stored OAuth2 token for an MCP server. A proxy admin may pass user_id to revoke another user's stored token.",
"operationId": "delete_mcp_oauth_user_credential_v1_mcp_server__server_id__oauth_user_credential_delete",
"parameters": [
{
@ -30248,6 +30570,23 @@
"title": "Server Id",
"type": "string"
}
},
{
"in": "query",
"name": "user_id",
"required": false,
"schema": {
"anyOf": [
{
"minLength": 1,
"type": "string"
},
{
"type": "null"
}
],
"title": "User Id"
}
}
],
"responses": {
@ -30447,7 +30786,7 @@
},
"/v1/mcp/server/{server_id}/user-credential": {
"delete": {
"description": "Delete the calling user's stored API key for a BYOK MCP server",
"description": "Delete the calling user's stored API key for a BYOK MCP server. A proxy admin may pass user_id to revoke another user's stored key.",
"operationId": "delete_mcp_user_credential_v1_mcp_server__server_id__user_credential_delete",
"parameters": [
{
@ -30458,6 +30797,23 @@
"title": "Server Id",
"type": "string"
}
},
{
"in": "query",
"name": "user_id",
"required": false,
"schema": {
"anyOf": [
{
"minLength": 1,
"type": "string"
},
{
"type": "null"
}
],
"title": "User Id"
}
}
],
"responses": {
@ -30549,6 +30905,58 @@
]
}
},
"/v1/mcp/server/{server_id}/user-credentials": {
"get": {
"description": "List every user's stored BYOK or OAuth2 credential for an MCP server (admin only, no secrets)",
"operationId": "list_mcp_server_user_credentials_v1_mcp_server__server_id__user_credentials_get",
"parameters": [
{
"in": "path",
"name": "server_id",
"required": true,
"schema": {
"title": "Server Id",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"items": {
"$ref": "#/components/schemas/MCPServerUserCredentialListItem"
},
"title": "Response List Mcp Server User Credentials V1 Mcp Server Server Id User Credentials Get",
"type": "array"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "List Mcp Server User Credentials",
"tags": [
"mcp_management"
]
}
},
"/v1/mcp/server/{server_id}/user-env-vars": {
"delete": {
"description": "Clear the calling user's per-user MCP env var values for this server.",
@ -30700,6 +31108,77 @@
}
},
"/v1/mcp/sessions": {
"delete": {
"description": "Force-close live stateful MCP gateway sessions on this proxy worker, selected by session id prefix and/or by the LiteLLM user that opened them (proxy admin only).",
"operationId": "delete_mcp_gateway_sessions_v1_mcp_sessions_delete",
"parameters": [
{
"in": "query",
"name": "session_id_prefix",
"required": false,
"schema": {
"anyOf": [
{
"minLength": 8,
"type": "string"
},
{
"type": "null"
}
],
"title": "Session Id Prefix"
}
},
{
"in": "query",
"name": "user_id",
"required": false,
"schema": {
"anyOf": [
{
"minLength": 1,
"type": "string"
},
{
"type": "null"
}
],
"title": "User Id"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/MCPGatewaySessionsTerminateResponse"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Delete Mcp Gateway Sessions",
"tags": [
"mcp_management"
]
},
"get": {
"description": "Live stateful MCP gateway sessions on this proxy worker, grouped by AI client and by user.",
"operationId": "get_mcp_gateway_sessions_v1_mcp_sessions_get",
@ -38542,6 +39021,7 @@
"type": "object"
},
"SCIMMultiValuedAttribute": {
"additionalProperties": true,
"properties": {
"display": {
"anyOf": [
@ -38577,13 +39057,17 @@
"title": "Type"
},
"value": {
"title": "Value",
"type": "string"
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Value"
}
},
"required": [
"value"
],
"title": "SCIMMultiValuedAttribute",
"type": "object"
},

View file

@ -469,6 +469,7 @@ class LiteLLMRoutes(enum.Enum):
mapped_pass_through_routes = [
"/bedrock",
"/comprehendmedical",
"/azure_speech",
"/transcribe",
"/vertex-ai",
"/vertex_ai",
@ -662,6 +663,11 @@ class LiteLLMRoutes(enum.Enum):
KeyManagementRoutes.AUTO_ROUTER_MANAGE.value,
]
team_service_account_key_routes = (
KeyManagementRoutes.KEY_GENERATE.value,
KeyManagementRoutes.KEY_UPDATE.value,
)
management_routes = (
[
# user
@ -1721,6 +1727,16 @@ class MCPUserCredentialListItem(LiteLLMPydanticObjectBase):
connected_at: str | None = None # ISO-8601
class MCPServerUserCredentialListItem(LiteLLMPydanticObjectBase):
"""One user's stored credential for an MCP server, as an admin sees it. Never carries the secret."""
user_id: str
credential_type: Literal["oauth2", "byok"]
expires_at: str | None = None
connected_at: str | None = None
updated_at: str
class MCPUserEnvVarsRequest(LiteLLMPydanticObjectBase):
"""Payload for storing the calling user's per-user env var values."""
@ -2763,6 +2779,25 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
description="sends alerts if requests hang for 5min+",
)
ui_access_mode: Literal["admin_only", "all"] | None = Field("all", description="Control access to the Proxy UI")
max_failed_login_attempts_per_source: int | None = Field(
None,
ge=1,
description="Failed Admin UI sign-in attempts allowed from one source address, across every username, within `failed_login_window_seconds`. One more blocks that address for `failed_login_block_seconds`. Half this value, rounded down but at least 1, is the allowance for one username from that address; one more blocks that address for that username only, and its further failures stop counting toward the address limit, so a script stuck on one account does not block everyone behind a shared address. The per-address limit is only enforced when `trusted_proxy_ranges` is set: to the proxies in front of LiteLLM, or to an empty list when clients connect directly. Left unset, the peer address may be a shared ingress and only the per-username half runs. IPv6 addresses are grouped by /64. Set under `general_settings` in config.yaml. Defaults to 10",
)
max_failed_login_attempts_per_source_overrides: dict[str, int] | None = Field(
None,
description="Per-address overrides of `max_failed_login_attempts_per_source`, keyed by IP address or CIDR range, e.g. {'1.2.3.4': 200, '5.6.0.0/24': 500}. The most specific matching range wins (between equivalent keys such as '1.2.3.4' and '1.2.3.4/32', an exemption wins, then the higher limit), and the per-username allowance for that address follows as half the override. A value of 0 exempts the address from both limits. Set under `general_settings` in config.yaml",
)
failed_login_window_seconds: int | None = Field(
None,
ge=1,
description="Fixed window in seconds over which failed Admin UI sign-in attempts are counted. The window starts at the first failure and is not extended by later ones. Set under `general_settings` in config.yaml. Defaults to 60",
)
failed_login_block_seconds: int | None = Field(
None,
ge=1,
description="How long a blocked source address, or source address and username, stays blocked. Every attempt from a blocked key, right or wrong, is refused with 429 before the password is checked; the block is not extended by refused attempts. Set under `general_settings` in config.yaml. Defaults to 300",
)
allowed_routes: list | None = Field(None, description="Proxy API Endpoints you want users to be able to access")
reject_clientside_metadata_tags: bool | None = Field(
None,
@ -2876,7 +2911,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
)
trusted_proxy_ranges: list[str] | None = Field(
None,
description="CIDR ranges of trusted reverse proxies allowed to provide identity headers for header-based auth paths such as enable_oauth2_proxy_auth and custom_ui_sso_sign_in_handler.",
description="CIDR ranges of trusted reverse proxies allowed to provide identity headers for header-based auth paths such as enable_oauth2_proxy_auth and custom_ui_sso_sign_in_handler, and whose X-Forwarded-For is used to attribute Admin UI sign-in attempts to a source address. Set it to an empty list when clients connect directly, so the peer address is the source. Left unset, or containing an entry that is not an address or CIDR range, the per-source sign-in limit is off.",
)
store_model_in_db: bool | None = Field(
None,
@ -3279,6 +3314,15 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
user_role=LitellmUserRoles.PROXY_ADMIN,
)
@property
def is_team_service_account(self) -> bool:
return (
self.user_id is None
and self.team_id is not None
and bool(self.metadata)
and self.metadata.get("service_account_id") is not None
)
def user_api_key_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool:
"""Return True if the caller's role grants unscoped read access to all

View file

@ -0,0 +1,445 @@
"""Failed-login accounting for the Admin UI sign-in path.
Wrong passwords are counted over a short window per source address and per source-and-username
pair; too many in one window blocks that key for a fixed time. While a key is blocked every attempt
from it, right or wrong, is refused with 429 before the password is checked. A blocked pair stops
counting against its source, so one script stuck on one account does not block the whole office.
Recovery is the master key over the API, which never passes through here, or waiting out the block.
"""
from __future__ import annotations
import asyncio
import hashlib
import ipaddress
import math
import time
from collections.abc import Mapping
from dataclasses import dataclass
from functools import cache
from typing import Final, Literal, NamedTuple, Protocol, TypeAlias
from fastapi import Request, status
from pydantic import TypeAdapter, ValidationError
from redis.exceptions import RedisError
from litellm._logging import verbose_proxy_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError
from litellm.constants import (
EMPTY_MAPPING,
LOGIN_THROTTLE_CACHE_KEY_PREFIX,
LOGIN_THROTTLE_MAX_TRACKED_BLOCKS,
LOGIN_THROTTLE_MAX_TRACKED_COUNTERS,
LOGIN_THROTTLE_NOT_BLOCKED,
LOGIN_THROTTLE_UNKNOWN_SOURCE,
)
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.auth.network import TrustedProxyConfig, resolve_client_ip
from litellm.secret_managers.main import get_secret_bool
DEFAULT_MAX_FAILED_LOGIN_ATTEMPTS_PER_SOURCE: Final = 10
DEFAULT_FAILED_LOGIN_WINDOW_SECONDS: Final = 60
DEFAULT_FAILED_LOGIN_BLOCK_SECONDS: Final = 300
IPV6_SOURCE_PREFIX_LENGTH: Final = 64
EXEMPT: Final = 0
SOURCE_LIMIT_KEY: Final = "max_failed_login_attempts_per_source"
SOURCE_LIMIT_OVERRIDES_KEY: Final = "max_failed_login_attempts_per_source_overrides"
WINDOW_KEY: Final = "failed_login_window_seconds"
BLOCK_KEY: Final = "failed_login_block_seconds"
TRUSTED_PROXY_RANGES_KEY: Final = "trusted_proxy_ranges"
_REDIS_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, asyncio.TimeoutError)
_LOCAL_BLOCK_EXPIRY: Final = TypeAdapter[float | None](float | None)
_SOURCE_LIMIT_OVERRIDES: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object])
_RANGE_ENTRIES: Final = TypeAdapter[tuple[object, ...]](tuple[object, ...])
Scope: TypeAlias = Literal["user", "source"]
_BlockTtls: TypeAlias = tuple[int, int]
_LUA_BLOCK_TTLS: Final = TypeAdapter[_BlockTtls](_BlockTtls)
_Network: TypeAlias = ipaddress.IPv4Network | ipaddress.IPv6Network
class LocalStore(Protocol):
"""The per-worker store behind the counters and blocks; ``InMemoryCache`` satisfies it."""
def get_cache(self, key: str) -> object: ...
def set_cache(self, key: str, value: float, *, ttl: int) -> None: ...
def increment_cache(self, key: str, value: float, *, ttl: int) -> float: ...
def delete_cache(self, key: str) -> None: ...
# KEYS: pair counter, pair block, source counter, source block (one cluster slot via the source hash tag)
# ARGV: pair limit, source limit (0 = source scope off), window seconds, block seconds
# Both scripts return {pair block TTL, source block TTL}; 0 or below means not blocked
_BLOCK_TTLS_LUA: Final = "return {redis.call('TTL', KEYS[2]), redis.call('TTL', KEYS[4])}"
_RECORD_FAILURE_LUA: Final = (
"local function bump(count_key, block_key, limit) "
"local blocked = redis.call('TTL', block_key) "
"if blocked > 0 then return blocked end "
"local count = redis.call('INCR', count_key) "
"if redis.call('TTL', count_key) < 0 then redis.call('EXPIRE', count_key, ARGV[3]) end "
"if count > limit then redis.call('SET', block_key, '1', 'EX', ARGV[4]) return tonumber(ARGV[4]) end "
"return 0 end "
"local user_block = bump(KEYS[1], KEYS[2], tonumber(ARGV[1])) "
"local source_block = 0 "
"if tonumber(ARGV[2]) > 0 and user_block == 0 then "
"source_block = bump(KEYS[3], KEYS[4], tonumber(ARGV[2])) end "
"return {user_block, source_block}"
)
_COUNTERS: Final = InMemoryCache(
max_size_in_memory=LOGIN_THROTTLE_MAX_TRACKED_COUNTERS, default_ttl=DEFAULT_FAILED_LOGIN_WINDOW_SECONDS
)
_BLOCKS: Final = InMemoryCache(
max_size_in_memory=LOGIN_THROTTLE_MAX_TRACKED_BLOCKS, default_ttl=DEFAULT_FAILED_LOGIN_BLOCK_SECONDS
)
@cache
def _rate_limit_disabled() -> bool:
return get_secret_bool("LITELLM_DISABLE_LOGIN_RATE_LIMIT", default_value=False) is True
@cache
def warn_login_counters_are_per_worker(num_workers: str) -> None:
verbose_proxy_logger.warning(
"Running %s workers but Redis is not configured. Failed Admin UI sign-in attempts are counted "
"per worker, so the effective limits are %s times the configured values. Configure Redis "
"to share one count across workers.",
num_workers,
num_workers,
)
@cache
def warn_source_login_limit_is_off() -> None:
verbose_proxy_logger.warning(
"%s is not set or not a valid list of ranges, so failed Admin UI sign-in attempts are limited per "
"source address and username only. Set it to the address ranges of the proxies in front of LiteLLM, "
"or to an empty list when clients connect directly, to also limit each source address across usernames.",
TRUSTED_PROXY_RANGES_KEY,
)
def declared_proxy_ranges(settings: Mapping[str, object]) -> tuple[str, ...] | None:
"""What the operator says fronts LiteLLM: the proxy ranges, an empty tuple for none, None when unsaid.
Only a declared topology makes the source address trustworthy enough to limit across usernames.
An unset key, a value that is not a list of ranges, or a list with an entry that is not an address
or range leaves it unknown and the source scope off.
"""
entries: Final = _configured_range_entries(settings.get(TRUSTED_PROXY_RANGES_KEY))
if entries is None or any(_parse_network(entry, TRUSTED_PROXY_RANGES_KEY) is None for entry in entries):
return None
return entries
def _configured_range_entries(raw_ranges: object) -> tuple[str, ...] | None:
"""Every configured entry, blanks included, so a stray empty string fails validation like any other typo."""
if raw_ranges is None:
return None
if isinstance(raw_ranges, str):
return tuple(part.strip() for part in raw_ranges.split(","))
try:
return tuple(str(entry).strip() for entry in _RANGE_ENTRIES.validate_python(raw_ranges))
except ValidationError:
verbose_proxy_logger.warning(
"Invalid %s value: expected a list of address ranges, got %s",
TRUSTED_PROXY_RANGES_KEY,
type(raw_ranges).__name__,
)
return None
def _positive_int(raw: object, key: str, default: int) -> int:
if raw is None:
return default
try:
value: Final = int(str(raw))
except (TypeError, ValueError):
verbose_proxy_logger.warning("Invalid %s value %r; using %s", key, raw, default)
return default
if value < 1:
verbose_proxy_logger.warning("Invalid %s value %s (must be >= 1); using %s", key, value, default)
return default
return value
def _int_setting(settings: Mapping[str, object], key: str, default: int) -> int:
return _positive_int(settings.get(key), key, default)
def _override_limit(raw: object, default: int) -> int:
"""A per-address override: a limit of 1 or more, or ``EXEMPT`` (0) to leave that address unlimited."""
if str(raw).strip() == str(EXEMPT):
return EXEMPT
return _positive_int(raw, SOURCE_LIMIT_OVERRIDES_KEY, default)
def _parse_address(client_ip: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address | None:
"""The address as it is limited and counted: an IPv4-mapped IPv6 address is its IPv4 address."""
try:
address: Final = ipaddress.ip_address(client_ip)
except ValueError:
return None
if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped is not None:
return address.ipv4_mapped
return address
def _parse_network(raw_range: str, setting_name: str = SOURCE_LIMIT_OVERRIDES_KEY) -> _Network | None:
try:
return ipaddress.ip_network(raw_range.strip(), strict=False)
except ValueError:
verbose_proxy_logger.warning("Invalid address or range %r in %s; skipping", raw_range, setting_name)
return None
def _precedence(network: _Network, limit: int) -> tuple[int, bool, int]:
"""Sort key for competing overrides: the longest prefix wins, then an exemption, then the higher limit."""
return (network.prefixlen, limit == EXEMPT, limit)
def _source_limit(settings: Mapping[str, object], client_ip: str) -> int:
"""Failure allowance for this address: the most specific configured range containing it, else the default.
``EXEMPT`` (0) means the operator opted this address out of both limits. Between equivalent keys such as
``1.2.3.4`` and ``1.2.3.4/32`` an exemption wins, then the higher limit.
"""
default: Final = _int_setting(settings, SOURCE_LIMIT_KEY, DEFAULT_MAX_FAILED_LOGIN_ATTEMPTS_PER_SOURCE)
raw_overrides: Final = settings.get(SOURCE_LIMIT_OVERRIDES_KEY)
if raw_overrides is None:
return default
try:
overrides: Final = _SOURCE_LIMIT_OVERRIDES.validate_python(raw_overrides)
except ValidationError:
verbose_proxy_logger.warning(
"Invalid %s value; expected a mapping of address or range to limit", SOURCE_LIMIT_OVERRIDES_KEY
)
return default
address: Final = _parse_address(client_ip)
if address is None:
return default
matches: Final = sorted(
_precedence(network, _override_limit(raw_limit, default))
for raw_range, raw_limit in overrides.items()
if (network := _parse_network(raw_range)) is not None and address in network
)
return matches[-1][-1] if matches else default
def user_limit_for(source_limit: int) -> int:
"""Failures allowed for one username from one address: half the address allowance, rounded down, at least 1."""
return max(source_limit // 2, 1)
def source_group(client_ip: str) -> str:
"""The bucket an address is counted in: IPv4 as is, IPv6 by its /64, so one prefix holder cannot rotate."""
address: Final = _parse_address(client_ip)
if address is None:
return client_ip
if isinstance(address, ipaddress.IPv6Address):
return str(ipaddress.ip_network((address, IPV6_SOURCE_PREFIX_LENGTH), strict=False))
return str(address)
class _Keys(NamedTuple):
pair_counter: str
pair_block: str
source_counter: str
source_block: str
@dataclass(frozen=True, slots=True)
class Block:
scope: Scope
retry_after: int
@dataclass(frozen=True, slots=True)
class LoginThrottle:
"""Failed-login limits for one request's source address.
``source_limit`` is None when the source scope is off: ``trusted_proxy_ranges`` is unset, so the peer
address may be a shared ingress. An empty list means clients connect directly and the peer is the source.
``user_limit`` is derived from the address allowance either way, see ``user_limit_for``. An address whose
override is ``EXEMPT`` gets a disabled throttle: nothing is counted or blocked for it.
"""
client_ip: str
source_limit: int | None
user_limit: int
window_seconds: int
block_seconds: int
counters: LocalStore
blocks: LocalStore
redis_cache: RedisCache | None = None
enabled: bool = True
@classmethod
def from_request(
cls,
request: Request,
general_settings: Mapping[str, object] | None,
redis_cache: RedisCache | None,
) -> LoginThrottle:
settings: Final[Mapping[str, object]] = general_settings if general_settings is not None else EMPTY_MAPPING
proxies: Final = declared_proxy_ranges(settings)
resolved, _ = resolve_client_ip(
request, TrustedProxyConfig(use_forwarded_for=bool(proxies), trusted_proxy_cidrs=proxies or ())
)
source_limit: Final = _source_limit(settings, resolved or LOGIN_THROTTLE_UNKNOWN_SOURCE)
exempt: Final = source_limit == EXEMPT
return cls(
client_ip=resolved or LOGIN_THROTTLE_UNKNOWN_SOURCE,
source_limit=source_limit if proxies is not None and resolved is not None and not exempt else None,
user_limit=user_limit_for(source_limit),
window_seconds=_int_setting(settings, WINDOW_KEY, DEFAULT_FAILED_LOGIN_WINDOW_SECONDS),
block_seconds=_int_setting(settings, BLOCK_KEY, DEFAULT_FAILED_LOGIN_BLOCK_SECONDS),
counters=_COUNTERS,
blocks=_BLOCKS,
redis_cache=redis_cache,
enabled=not exempt and not _rate_limit_disabled(),
)
def _keys(self, username: str) -> _Keys:
group: Final = source_group(self.client_ip)
user: Final = hashlib.sha256(username.casefold().encode("utf-8")).hexdigest()
return _Keys(
pair_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:user:{user}",
pair_block=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:block:user:{user}",
source_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:source",
source_block=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:block:source",
)
async def attempt(self, username: str) -> LoginAttempt:
"""Refuses a blocked key before any credential is looked at; otherwise hands back the attempt to settle."""
if not self.enabled:
return LoginAttempt(throttle=self, username=username)
block: Final = await self._active_block(self._keys(username))
if block is None:
return LoginAttempt(throttle=self, username=username)
verbose_proxy_logger.warning(
"Admin UI sign-in refused: the %s is blocked for %s more seconds; username=%r source=%s",
block.scope,
block.retry_after,
username,
self.client_ip,
)
raise self.refused(block.retry_after)
async def _active_block(self, keys: _Keys) -> Block | None:
local: Final = self._local_block_ttls(keys)
shared: Final = await self._shared_block_ttls(keys)
user_ttl: Final = max(local[0], shared[0])
source_ttl: Final = max(local[1], shared[1])
if self.source_limit is not None and source_ttl > 0:
return Block(scope="source", retry_after=source_ttl)
if user_ttl > 0:
return Block(scope="user", retry_after=user_ttl)
return None
async def _shared_block_ttls(self, keys: _Keys) -> _BlockTtls:
if self.redis_cache is None:
return LOGIN_THROTTLE_NOT_BLOCKED
try:
return _LUA_BLOCK_TTLS.validate_python(
await self.redis_cache.async_register_script(_BLOCK_TTLS_LUA)(keys, ())
)
except _REDIS_FAILURES as err:
self._warn_redis(err)
return LOGIN_THROTTLE_NOT_BLOCKED
def _local_block_ttls(self, keys: _Keys) -> _BlockTtls:
return self._local_block_ttl(keys.pair_block), self._local_block_ttl(keys.source_block)
def _local_block_ttl(self, block_key: str) -> int:
expires_at: Final = _LOCAL_BLOCK_EXPIRY.validate_python(self.blocks.get_cache(block_key))
if expires_at is None:
return 0
return max(math.ceil(expires_at - time.time()), 0)
async def record_failure(self, username: str) -> _BlockTtls:
keys: Final = self._keys(username)
source_limit: Final = self.source_limit or 0
if self.redis_cache is not None:
try:
return _LUA_BLOCK_TTLS.validate_python(
await self.redis_cache.async_register_script(_RECORD_FAILURE_LUA)(
keys, (self.user_limit, source_limit, self.window_seconds, self.block_seconds)
)
)
except _REDIS_FAILURES as err:
self._warn_redis(err)
user_block: Final = self._local_bump(keys.pair_counter, keys.pair_block, self.user_limit)
if source_limit == 0 or user_block > 0:
return user_block, 0
return user_block, self._local_bump(keys.source_counter, keys.source_block, source_limit)
def _local_bump(self, count_key: str, block_key: str, limit: int) -> int:
blocked: Final = self._local_block_ttl(block_key)
if blocked > 0:
return blocked
count: Final = int(self.counters.increment_cache(count_key, 1, ttl=self.window_seconds))
if count <= limit:
return 0
self.blocks.set_cache(block_key, time.time() + self.block_seconds, ttl=self.block_seconds)
return self.block_seconds
async def clear_pair(self, username: str) -> None:
pair_counter: Final = self._keys(username).pair_counter
if self.redis_cache is not None:
try:
await self.redis_cache.async_delete_cache(pair_counter)
except _REDIS_FAILURES as err:
self._warn_redis(err)
self.counters.delete_cache(pair_counter)
def _warn_redis(self, err: Exception) -> None:
verbose_proxy_logger.warning(
"Redis failed while counting Admin UI sign-in attempts; using this worker's own counters "
"until it recovers: %s",
err,
)
@staticmethod
def refused(retry_after: int) -> ProxyException:
return ProxyException(
message="Too many failed sign-in attempts. Try again later.",
type=ProxyErrorTypes.auth_error,
param="username",
code=status.HTTP_429_TOO_MANY_REQUESTS,
headers={"Retry-After": str(retry_after)}, # mutable-ok: ProxyException writes into its headers dict
)
@dataclass(frozen=True, slots=True)
class LoginAttempt:
throttle: LoginThrottle
username: str
async def succeeded(self) -> None:
if not self.throttle.enabled:
return
await self.throttle.clear_pair(self.username)
async def failed(self) -> None:
if not self.throttle.enabled:
return
user_block, source_block = await self.throttle.record_failure(self.username)
if user_block == 0 and source_block == 0:
return
verbose_proxy_logger.warning(
"Admin UI sign-in blocked for %s seconds after too many failures; scope=%s username=%r source=%s",
user_block or source_block,
"user" if user_block else "source",
self.username,
self.throttle.client_ip,
)

View file

@ -27,6 +27,7 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured
from litellm.proxy.auth.login_throttle import LoginAttempt, LoginThrottle
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_key_helper_fn,
@ -44,6 +45,11 @@ from litellm.repositories.user_repository import UserRepository
from litellm.secret_managers.main import get_secret_bool
from litellm.types.proxy.ui_sso import ReturnedUITokenObject
INVALID_UI_CREDENTIALS_MESSAGE: Final = (
"Invalid credentials used to access UI. Check 'UI_USERNAME' and 'UI_PASSWORD', or the password set for your user"
)
INVALID_USER_PASSWORD_MESSAGE: Final = "Invalid credentials used to access UI. Check the password set for your user"
async def _rehash_password_if_needed(user_id: str, password: str, stored: str) -> None:
"""Rehash legacy password (SHA256) to scrypt on successful login."""
@ -92,6 +98,21 @@ def _matches_env_credentials(username: str, password: str, master_key: str | Non
)
def _admin_credentials_match(
username: str, password: str, master_key: str, general_settings: Mapping[str, object]
) -> bool:
return general_settings.get("disable_env_credential_login") is not True and _matches_env_credentials(
username, password, master_key
)
def _invalid_credentials_message(general_settings: Mapping[str, object]) -> str:
"""One rejection message for unknown usernames and wrong passwords alike, so neither can be enumerated."""
if is_env_credential_login_enabled(general_settings):
return INVALID_UI_CREDENTIALS_MESSAGE
return INVALID_USER_PASSWORD_MESSAGE
def is_env_credential_login_enabled(general_settings: Mapping[str, object]) -> bool:
"""Whether a login with UI_USERNAME/UI_PASSWORD (or the master-key fallback) can succeed.
@ -137,6 +158,7 @@ async def authenticate_user(
password: str,
master_key: str | None,
prisma_client: PrismaClient | None,
throttle: LoginThrottle,
general_settings: Mapping[str, object] = MappingProxyType({}),
) -> LoginResult:
"""
@ -151,6 +173,7 @@ async def authenticate_user(
password: Password from the login form
master_key: Master key for the proxy (required)
prisma_client: Prisma database client (optional)
throttle: Failed sign-in accounting for this request's source address
general_settings: Proxy general_settings, checked for
`disable_password_login_when_sso_enabled` and
`disable_env_credential_login`
@ -163,9 +186,11 @@ async def authenticate_user(
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
`disable_password_login_when_sso_enabled`, or by the failed sign-in block in
`throttle`, can still administer the proxy over the API with the master key
(Authorization: Bearer <master_key>), which never goes through this function.
No credential, the env admin credentials and the master key included, is
exempt from the block. 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.
@ -194,6 +219,19 @@ async def authenticate_user(
code=500,
)
attempt: Final = await throttle.attempt(username)
return await _sign_in(username, password, master_key, prisma_client, attempt, general_settings)
async def _sign_in(
username: str,
password: str,
master_key: str,
prisma_client: PrismaClient | None,
attempt: LoginAttempt,
general_settings: Mapping[str, object],
) -> LoginResult:
admin_credentials_match: Final = _admin_credentials_match(username, password, master_key, general_settings)
# Check if we can find the `username` in the db. On the UI, users can enter username=their email
_user_row: LiteLLM_UserTable | None = None
user_role: (
@ -219,20 +257,13 @@ async def authenticate_user(
- Login with UI_USERNAME and UI_PASSWORD
- Login with Invite Link `user_email` and `password` combination
"""
if general_settings.get("disable_env_credential_login") is not True and _matches_env_credentials(
username, password, master_key
):
if admin_credentials_match:
# Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin
user_role = LitellmUserRoles.PROXY_ADMIN
user_id = LITELLM_PROXY_ADMIN_NAME
# we want the key created to have PROXY_ADMIN_PERMISSIONS
key_user_id = LITELLM_PROXY_ADMIN_NAME
if (
os.getenv("PROXY_ADMIN_ID", None) is not None and os.environ["PROXY_ADMIN_ID"] == user_id
) or user_id == LITELLM_PROXY_ADMIN_NAME:
# checks if user is admin
key_user_id = os.getenv("PROXY_ADMIN_ID", LITELLM_PROXY_ADMIN_NAME)
key_user_id: Final = os.getenv("PROXY_ADMIN_ID", LITELLM_PROXY_ADMIN_NAME)
# Admin is Authe'd in - generate key for the UI to access Proxy
@ -294,6 +325,8 @@ async def authenticate_user(
key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(user_info)
await attempt.succeeded()
return LoginResult(
user_id=user_id,
key=key,
@ -349,6 +382,8 @@ async def authenticate_user(
key = response["token"]
await attempt.succeeded()
return LoginResult(
user_id=user_id,
key=key,
@ -357,20 +392,17 @@ async def authenticate_user(
login_method="username_password",
)
else:
await attempt.failed()
raise ProxyException(
message=f"Invalid credentials used to access UI.\nNot valid credentials for {username}",
message=_invalid_credentials_message(general_settings),
type=ProxyErrorTypes.auth_error,
param="invalid_credentials",
code=401,
)
else:
env_credentials_hint: Final = (
"\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file"
if is_env_credential_login_enabled(general_settings)
else ""
)
await attempt.failed()
raise ProxyException(
message=f"Invalid credentials used to access UI.{env_credentials_hint}",
message=_invalid_credentials_message(general_settings),
type=ProxyErrorTypes.auth_error,
param="invalid_credentials",
code=401,

View file

@ -1,6 +1,7 @@
from __future__ import annotations
import ipaddress
from collections.abc import Sequence
from typing import Any, Final
from fastapi import Request
@ -19,7 +20,7 @@ class NetworkContext(BaseModel):
class TrustedProxyConfig(BaseModel):
use_forwarded_for: bool = False
trusted_proxy_cidrs: list[str] = Field(default_factory=list)
trusted_proxy_cidrs: Sequence[str] = Field(default_factory=tuple)
def normalize_cidr_ranges(configured_ranges: Any, *, setting_name: str = "trusted_proxy_cidrs") -> list[str]:
@ -49,6 +50,12 @@ def parse_trusted_proxy_ranges(
return networks
def _unmapped(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> ipaddress.IPv4Address | ipaddress.IPv6Address:
if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None:
return addr.ipv4_mapped
return addr
def ip_in_networks(client_ip: str | None, networks: list[TrustedProxyNetwork]) -> bool:
if not client_ip or not networks:
return False
@ -56,7 +63,8 @@ def ip_in_networks(client_ip: str | None, networks: list[TrustedProxyNetwork]) -
addr: Final = ipaddress.ip_address(client_ip.strip())
except ValueError:
return False
return any(addr in network for network in networks)
candidates: Final = (addr, _unmapped(addr))
return any(candidate in network for candidate in candidates for network in networks)
def _is_valid_ip(value: str) -> bool:

View file

@ -326,7 +326,12 @@ class RouteChecks:
pass
elif route.startswith("/v1/mcp/") or route.startswith("/mcp-rest/"):
pass # authN/authZ handled by api itself
elif RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token):
elif RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token) or (
valid_token.is_team_service_account
and RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.team_service_account_key_routes.value
)
):
pass
elif valid_token.allowed_routes is not None:
# check if route is in allowed_routes (exact match or prefix match)

View file

@ -107,6 +107,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
_safe_get_request_headers,
_safe_get_request_query_params,
_safe_set_request_parsed_body,
is_opaque_audio_pass_through_request,
populate_request_with_path_params,
read_raw_json_body,
rewrite_request_model,
@ -1356,6 +1357,12 @@ async def _read_request_body_deferring_parse_failure(
must run (resolving identity onto the request's trace) before the 400 goes
out; the caller re-raises the returned exception once identity is seeded.
"""
if is_opaque_audio_pass_through_request(
route=get_request_route(request=request),
content_type=_safe_get_request_headers(request=request).get("content-type", ""),
):
_safe_set_request_parsed_body(request=request, parsed_body={}) # mutable-ok: the body cache stores a plain dict
return {}, None # mutable-ok: request_data is a plain dict across the whole auth path
try:
parsed_body: Final = await _read_request_body(request=request)
except ProxyException as parse_exception:

View file

@ -9,7 +9,11 @@ from fastapi import Request, UploadFile, status
from typing_extensions import NotRequired, ReadOnly, Required
from litellm._logging import verbose_proxy_logger
from litellm.constants import CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB
from litellm.constants import (
AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX,
CLIENT_REQUESTED_MODEL_SCOPE_KEY,
MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB,
)
from litellm.proxy._types import ProxyException
from litellm.proxy.common_utils.callback_utils import (
get_metadata_variable_name_from_kwargs,
@ -214,6 +218,14 @@ async def _read_request_body(request: Request | None) -> dict:
return {}
def is_opaque_audio_pass_through_request(route: str, content_type: str) -> bool:
"""Azure Speech bodies (raw audio, multipart uploads) are forwarded byte for byte, so auth must not consume them."""
media_type: Final = _normalize_media_type(content_type)
return route.startswith(f"{AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX}/") and (
media_type.startswith("audio/") or media_type == "multipart/form-data"
)
async def read_raw_json_body(request: Request | None) -> bytes | None:
if request is None or _safe_get_request_parsed_body(request=request) is None:
return None

View file

@ -87,6 +87,12 @@ class SettingsStore(MutableMapping[str, JsonValue]):
)
self._deleted_runtime_keys = self._deleted_runtime_keys | frozenset((key,))
def clear(self) -> None:
self._deleted_runtime_keys = frozenset(key for key in self._keys() if not self.owned_by_config(key))
self._runtime_values = MappingProxyType(
{key: value for key, value in self._runtime_values.items() if self.owned_by_config(key)}
)
def __iter__(self) -> Iterator[str]:
return iter(
key

View file

@ -18,7 +18,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload
from urllib.parse import quote, unquote
from typing_extensions import ReadOnly, TypedDict
from typing_extensions import LiteralString, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
@ -136,6 +136,8 @@ class _SpendBatchManager(Protocol):
class _SpendTransaction(Protocol):
def batch_(self) -> _SpendBatchManager: ...
async def execute_raw(self, query: LiteralString, *args: object) -> int: ...
class _SpendTransactionManager(Protocol):
async def __aenter__(self) -> _SpendTransaction: ...
@ -161,6 +163,44 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager:
return tx
# The per-team advisory lock the team endpoints hold while changing a roster (TEAM_ADVISORY_LOCK_SQL),
# so the roster check below cannot interleave with their writes. A row lock would deadlock with the
# access-group endpoints, which lock a team row after an access-group lock.
_TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
# One statement adds every member's cost to their membership row. A missing row is created only
# while the user is still on the team's roster, so a spend flush landing after a removal never
# recreates the member.
_TEAM_MEMBER_SPEND_SQL: Final = """
INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend)
SELECT p.user_id, p.team_id, p.cost, p.cost
FROM unnest($1::text[], $2::text[], $3::float8[]) AS p(user_id, team_id, cost)
WHERE EXISTS (
SELECT 1 FROM "LiteLLM_TeamTable" t
WHERE t.team_id = p.team_id
AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id))
)
OR EXISTS (SELECT 1 FROM "LiteLLM_TeamMembership" m WHERE m.user_id = p.user_id AND m.team_id = p.team_id)
ON CONFLICT (user_id, team_id) DO UPDATE
SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend,
total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend
"""
async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None:
# key is "team_id::<value>::user_id::<value>"; locks are taken in sorted team_id order like the team endpoints
rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items())
team_ids: Final = tuple(team_id for team_id, _user_id, _cost in rows)
for team_id in dict.fromkeys(team_ids):
_ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id)
_ = await transaction.execute_raw(
_TEAM_MEMBER_SPEND_SQL,
tuple(user_id for _team_id, user_id, _cost in rows),
team_ids,
tuple(cost for _team_id, _user_id, cost in rows),
)
def get_llm_router():
"""The proxy's router, or None outside a running proxy.
@ -1685,21 +1725,7 @@ class DBSpendUpdateWriter:
start_time = time.time()
try:
async with _spend_update_tx(prisma_client) as transaction:
async with transaction.batch_() as batcher:
# Sort by composite key for consistent lock ordering across pods to prevent deadlocks.
# Key format "team_id::<v>::user_id::<v>" makes the string sort equivalent to sorting by (team_id, user_id).
for key, response_cost in sorted(team_member_list_transactions.items()):
# key is "team_id::<value>::user_id::<value>"
team_id = key.split("::")[1]
user_id = key.split("::")[3]
batcher.litellm_teammembership.update_many( # 'update_many' prevents error from being raised if no row exists
where={"team_id": team_id, "user_id": user_id},
data={
"spend": {"increment": response_cost},
"total_spend": {"increment": response_cost},
},
)
await _write_team_member_spend(transaction, team_member_list_transactions)
# Transaction succeeded, break out of retry loop
break
except Exception as e:

View file

@ -21,7 +21,7 @@ if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.router import Router
COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"})
COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr", "typesafe"})
_NO_COMPRESSION: Final = "none"
# A ContextVar, not metadata: metadata reaches spend logs the caller can read, and a

View file

@ -0,0 +1,74 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Final
from pydantic import BaseModel
from litellm.types.guardrails import (
GuardrailEventHooks,
Mode,
SupportedGuardrailIntegrations,
)
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
TypeSafeGuardrailOptionalParams,
)
from .typesafe import TypeSafeGuardrail
if TYPE_CHECKING:
from litellm.types.guardrails import Guardrail, LitellmParams
def _coerce_event_hook(
mode: str | list[str] | Mode,
) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode:
if isinstance(mode, Mode):
return mode
if isinstance(mode, list):
return [ # mutable-ok: CustomGuardrail event_hook contract wants a list
GuardrailEventHooks(item) for item in mode
]
return GuardrailEventHooks(mode)
def _optional_params(litellm_params: LitellmParams) -> TypeSafeGuardrailOptionalParams:
value: Final = litellm_params.optional_params
if isinstance(value, TypeSafeGuardrailOptionalParams):
return value
if isinstance(value, BaseModel):
return TypeSafeGuardrailOptionalParams.model_validate(value.model_dump())
return TypeSafeGuardrailOptionalParams()
def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> TypeSafeGuardrail:
import litellm
optional_params: Final = _optional_params(litellm_params)
_callback: Final = TypeSafeGuardrail(
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
model=litellm_params.model,
relevance_threshold=optional_params.relevance_threshold,
min_chars_to_evaluate=optional_params.min_chars_to_evaluate,
max_result_chars_in_state=optional_params.max_result_chars_in_state,
guardrail_name=guardrail["guardrail_name"],
event_hook=_coerce_event_hook(litellm_params.mode),
default_on=litellm_params.default_on or False,
unreachable_fallback=(
litellm_params.unreachable_fallback if "unreachable_fallback" in litellm_params.model_fields_set else None
),
)
litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped
_callback
)
return _callback
guardrail_initializer_registry: Final = { # mutable-ok: guardrail_registry discovery checks isinstance(registry, dict)
SupportedGuardrailIntegrations.TYPESAFE.value: initialize_guardrail,
}
guardrail_class_registry: Final = { # mutable-ok: guardrail_registry discovery checks isinstance(registry, dict)
SupportedGuardrailIntegrations.TYPESAFE.value: TypeSafeGuardrail,
}

View file

@ -0,0 +1,416 @@
"""TypeSafe (Jev) relevance-based compaction guardrail.
Instead of summarizing tool output, the guardrail asks TypeSafe's Jev model
one yes/no question per completed tool exchange ("is this result still needed
for the current task?") over ``POST {api_base}/v1/systemone`` and blanks the
tool results Jev judges no longer relevant.
"""
from __future__ import annotations
import asyncio
import time
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Annotated, Final, Literal
import httpx
from fastapi import HTTPException
from httpx import Response as HttpxResponse
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.compression.compress import get_protected_indices
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # decorator is untyped in custom_guardrail
)
from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # helper is untyped in http_handler
httpxSpecialProvider,
)
from litellm.proxy.guardrails.guardrail_hooks.content_text import content_to_text
from litellm.secret_managers.main import get_secret_str
from litellm.types.guardrails import GuardrailEventHooks, Mode
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
TypeSafeGuardrailConfigModel,
)
DEFAULT_API_BASE: Final = "https://api.typesafe.ai"
DEFAULT_MODEL: Final = "jev-latest"
DEFAULT_RELEVANCE_THRESHOLD: Final = 0.2
DEFAULT_MIN_CHARS_TO_EVALUATE: Final = 200
DEFAULT_MAX_RESULT_CHARS_IN_STATE: Final = 4000
_MAX_EXCHANGES_EVALUATED: Final = 200
_JEV_TIMEOUT_SECONDS: Final = 30.0
DROPPED_RESULT_TEXT: Final = (
"[Tool result removed by TypeSafe compaction: judged no longer relevant to the current task]"
)
_ELISION_MARKER: Final = "\n... [middle truncated] ...\n"
_STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
def _as_str_object_dict(value: object) -> dict[str, object] | None:
try:
return _STR_OBJECT_DICT_ADAPTER.validate_python(value)
except ValidationError:
return None
def _as_object_list(value: object) -> list[object] | None:
try:
return _OBJECT_LIST_ADAPTER.validate_python(value)
except ValidationError:
return None
def _safe_response_text(response: HttpxResponse | None, limit: int = 500) -> str:
if response is None:
return ""
try:
text: Final = response.text
except httpx.DecodingError:
return "<undecodable response body>"
return (text or "")[:limit]
class _JevNoulAnswer(BaseModel):
model_config = ConfigDict(frozen=True, allow_inf_nan=False)
type: Literal["noul"]
noul: Annotated[float, Field(ge=0.0, le=1.0)]
class _JevSystemOneResponse(BaseModel):
model_config = ConfigDict(frozen=True)
answers: Mapping[str, _JevNoulAnswer]
_JEV_RESPONSE_ADAPTER: Final = TypeAdapter(_JevSystemOneResponse)
def _truncate_for_state(text: str, max_chars: int) -> str:
"""Keeps the head and tail within ``max_chars`` so Jev sees both ends of a long result."""
if len(text) <= max_chars:
return text
if max_chars <= len(_ELISION_MARKER):
return text[:max_chars]
budget: Final = max_chars - len(_ELISION_MARKER)
head: Final = budget // 2
return text[:head] + _ELISION_MARKER + text[len(text) - (budget - head) :]
def _question_instructions(question_id: str) -> str:
return (
f"Is tool exchange `{question_id}` in `tool_exchanges` still needed by the assistant to "
"complete `task`? Answer yes if its result contains information the assistant has not yet "
"fully used or will need again; answer no if it is off-topic, superseded, or already "
"incorporated into later messages."
)
def _tool_call_entry(tool_call: object) -> dict[str, object] | None:
parsed_call = _as_str_object_dict(tool_call)
if parsed_call is None:
return None
function = _as_str_object_dict(parsed_call.get("function"))
fn = function if function is not None else parsed_call
return {"name": fn.get("name"), "arguments": fn.get("arguments")} # mutable-ok: serialized to JSON
def _tool_call_entries(assistant_message: Mapping[str, object]) -> tuple[dict[str, object], ...]:
tool_calls: Final = _as_object_list(assistant_message.get("tool_calls"))
if tool_calls is None:
return ()
return tuple(entry for tool_call in tool_calls if (entry := _tool_call_entry(tool_call)) is not None)
def _protected_indices(messages: Sequence[Mapping[str, object]]) -> frozenset[int]:
"""``get_protected_indices`` expanded over whole tool exchanges, so the most recent exchange is never evaluated."""
protected: Final = frozenset(get_protected_indices(messages))
return protected | frozenset(
index
for group in group_tool_exchanges(messages)
if any(member in protected for member in group)
for index in group
)
class TypeSafeGuardrail(CustomGuardrail):
def __init__(
self,
api_base: str | None = None,
api_key: str | None = None,
model: str | None = None,
relevance_threshold: float | None = None,
min_chars_to_evaluate: int | None = None,
max_result_chars_in_state: int | None = None,
unreachable_fallback: str | None = None,
guardrail_name: str | None = None,
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None,
default_on: bool = False,
async_handler: AsyncHTTPHandler | None = None,
) -> None:
raw_api_base: Final = (api_base or get_secret_str("TYPESAFE_API_BASE") or DEFAULT_API_BASE).rstrip("/")
self.typesafe_api_base = raw_api_base
self.typesafe_api_key = api_key or get_secret_str("TYPESAFE_API_KEY")
if not self.typesafe_api_key:
raise ValueError(
"TypeSafe guardrail requires an API key. Set `api_key` in the "
"guardrail config or the TYPESAFE_API_KEY env var."
)
self.jev_model = model or DEFAULT_MODEL
self.relevance_threshold = DEFAULT_RELEVANCE_THRESHOLD if relevance_threshold is None else relevance_threshold
self.min_chars_to_evaluate = (
DEFAULT_MIN_CHARS_TO_EVALUATE if min_chars_to_evaluate is None else min_chars_to_evaluate
)
self.max_result_chars_in_state = (
DEFAULT_MAX_RESULT_CHARS_IN_STATE if max_result_chars_in_state is None else max_result_chars_in_state
)
self.unreachable_fallback: Literal["fail_closed", "fail_open"] = (
"fail_closed" if unreachable_fallback == "fail_closed" else "fail_open"
)
self.async_handler: AsyncHTTPHandler = async_handler or get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback,
)
super().__init__( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.__init__ is untyped
guardrail_name=guardrail_name,
event_hook=event_hook,
default_on=default_on,
)
def _handle_failure(self, error: str, log_detail: dict[str, object]) -> None:
"""fail_open logs and returns; fail_closed raises a generic 502 (upstream bodies stay in server logs)."""
if self.unreachable_fallback == "fail_open":
verbose_proxy_logger.warning(
"TypeSafe: %s; fail_open configured, forwarding request uncompacted. detail=%s",
error,
log_detail,
)
return
verbose_proxy_logger.error("TypeSafe: %s. detail=%s", error, log_detail)
raise HTTPException(status_code=502, detail={"error": error}) # mutable-ok: FastAPI wants a dict detail
def _candidate_exchanges(self, messages: Sequence[dict[str, object]]) -> tuple[tuple[int, ...], ...]:
"""Completed tool exchanges eligible for evaluation: unprotected, and long enough to be worth a call."""
protected: Final = _protected_indices(messages)
candidates: Final = tuple(
group
for group in group_tool_exchanges(messages)
if len(group) >= 2
and messages[group[0]].get("role") == "assistant"
and not any(member in protected for member in group)
and len(self._exchange_tool_text(messages, group)) >= self.min_chars_to_evaluate
)
return candidates[-_MAX_EXCHANGES_EVALUATED:]
@staticmethod
def _exchange_tool_text(messages: Sequence[dict[str, object]], group: tuple[int, ...]) -> str:
return "".join(
content_to_text(messages[index].get("content"))
for index in group[1:]
if messages[index].get("role") in ("tool", "function")
)
def _build_state(
self, messages: Sequence[dict[str, object]], candidates: tuple[tuple[int, ...], ...]
) -> dict[str, object]:
task: Final = next(
(
content_to_text(messages[index].get("content"))
for index in range(len(messages) - 1, -1, -1)
if messages[index].get("role") == "user"
),
"",
)
system: Final = "\n\n".join(
content_to_text(message.get("content")) for message in messages if message.get("role") == "system"
)
tool_exchanges: Final = { # mutable-ok: accumulated once, serialized to JSON
f"e{ordinal}": { # mutable-ok: serialized to JSON
"tool_calls": _tool_call_entries(messages[group[0]]),
"result": _truncate_for_state(
self._exchange_tool_text(messages, group), self.max_result_chars_in_state
),
}
for ordinal, group in enumerate(candidates)
}
return {"task": task, "system": system, "tool_exchanges": tool_exchanges} # mutable-ok: serialized to JSON
async def _call_systemone(
self, state: dict[str, object], question_ids: Sequence[str]
) -> _JevSystemOneResponse | None:
"""Returns the response, or None when the service failed and fail_open applies."""
payload: Final[dict[str, object]] = { # mutable-ok: serialized to JSON by httpx
"model": self.jev_model,
"state": state,
"questions": { # mutable-ok: serialized to JSON
question_id: { # mutable-ok: serialized to JSON
"type": "noul",
"instructions": _question_instructions(question_id),
}
for question_id in question_ids
},
}
try:
raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped
url=f"{self.typesafe_api_base}/v1/systemone",
json=payload,
headers={ # mutable-ok: httpx header contract is a dict
"Authorization": f"Bearer {self.typesafe_api_key}",
"Content-Type": "application/json",
},
timeout=_JEV_TIMEOUT_SECONDS,
)
except asyncio.CancelledError:
raise
except Exception as e:
detail: Final[dict[str, object]] = (
{ # mutable-ok: log detail record
"error_type": type(e).__name__,
"detail": str(e),
"status_code": e.response.status_code,
"body": _safe_response_text(e.response),
}
if isinstance(e, httpx.HTTPStatusError)
else {"error_type": type(e).__name__, "detail": str(e)} # mutable-ok: log detail record
)
self._handle_failure("TypeSafe evaluation service request failed", detail)
return None
if not 200 <= raw_response.status_code < 300:
self._handle_failure(
"TypeSafe evaluation service returned an error",
{ # mutable-ok: log detail record
"status_code": raw_response.status_code,
"body": _safe_response_text(raw_response),
},
)
return None
try:
body: Final[object] = raw_response.json() # pyright: ignore[reportAny] # httpx Response.json() is untyped
except (ValueError, httpx.DecodingError, RecursionError):
self._handle_failure(
"TypeSafe evaluation service returned an unreadable response",
{"body": _safe_response_text(raw_response)}, # mutable-ok: log detail record
)
return None
try:
return _JEV_RESPONSE_ADAPTER.validate_python(body)
except ValidationError:
self._handle_failure(
"TypeSafe evaluation service returned unexpected response shape",
{"body": _safe_response_text(raw_response)}, # mutable-ok: log detail record
)
return None
@log_guardrail_information
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
if input_type != "request":
return inputs
structured_messages: Final = _as_object_list(inputs.get("structured_messages"))
if not structured_messages:
return inputs
parsed_messages: Final = tuple(_as_str_object_dict(m) for m in structured_messages)
if any(m is None for m in parsed_messages):
return inputs
messages: Final = tuple(m for m in parsed_messages if m is not None)
candidates: Final = self._candidate_exchanges(messages)
if not candidates:
verbose_proxy_logger.debug("TypeSafe: no completed tool exchanges eligible for evaluation")
return inputs
question_ids: Final = tuple(f"e{ordinal}" for ordinal in range(len(candidates)))
state: Final = self._build_state(messages, candidates)
start_time: Final = time.monotonic()
response: Final = await self._call_systemone(state, question_ids)
end_time: Final = time.monotonic()
if response is None:
self.add_standard_logging_guardrail_information_to_request_data( # pyright: ignore[reportUnknownMemberType] # untyped base helper
guardrail_json_response={ # mutable-ok: must stay JSON-serializable for shared logging
"error": "TypeSafe evaluation unavailable; request forwarded uncompacted",
"model": self.jev_model,
},
request_data=request_data,
guardrail_status="guardrail_failed_to_respond",
guardrail_provider="typesafe",
start_time=start_time,
end_time=end_time,
duration=end_time - start_time,
)
return inputs
dropped_ordinals: Final = frozenset(
ordinal
for ordinal in range(len(candidates))
if (answer := response.answers.get(f"e{ordinal}")) is not None and answer.noul < self.relevance_threshold
)
dropped_tool_indices: Final[frozenset[int]] = frozenset(
index
for ordinal in dropped_ordinals
for index in candidates[ordinal][1:]
if messages[index].get("role") in ("tool", "function")
)
if not dropped_tool_indices:
verbose_proxy_logger.debug("TypeSafe: all evaluated exchanges still relevant; request unchanged")
return inputs
compacted_messages: Final = [ # mutable-ok: structured_messages contract is a list of dicts
{**message, "content": DROPPED_RESULT_TEXT} # mutable-ok: JSON message row
if index in dropped_tool_indices
else message
for index, message in enumerate(messages)
]
chars_removed: Final = sum(
len(content_to_text(messages[index].get("content"))) - len(DROPPED_RESULT_TEXT)
for index in dropped_tool_indices
)
exchanges_dropped: Final = len(dropped_ordinals)
verbose_proxy_logger.info(
"TypeSafe: evaluated %s tool exchange(s), dropped %s, ~%s chars removed",
len(candidates),
exchanges_dropped,
chars_removed,
)
self.add_standard_logging_guardrail_information_to_request_data( # pyright: ignore[reportUnknownMemberType] # untyped base helper
guardrail_json_response={ # mutable-ok: must stay JSON-serializable for shared logging
"exchanges_evaluated": len(candidates),
"exchanges_dropped": exchanges_dropped,
"chars_removed": chars_removed,
"model": self.jev_model,
},
request_data=request_data,
guardrail_status="success",
guardrail_provider="typesafe",
start_time=start_time,
end_time=end_time,
duration=end_time - start_time,
)
return {**inputs, "structured_messages": compacted_messages} # pyright: ignore[reportReturnType] # mutable-ok: inputs protocol is a plain dict # plain dicts satisfy AllMessageValues at runtime
@staticmethod
def get_config_model() -> type[TypeSafeGuardrailConfigModel] | None:
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
TypeSafeGuardrailConfigModel,
)
return TypeSafeGuardrailConfigModel

View file

@ -9,7 +9,7 @@ from fastapi import HTTPException, status
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.constants import PTU_SENTINEL_API_KEY
from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.spend_tracking.key_metadata_recovery import (
attach_user_emails,
@ -146,15 +146,9 @@ class _AggregatedSpendData(TypedDict):
totals: SpendMetrics
class _GroupingSetsRow(SimpleNamespace):
class _RollupMetricsRow(SimpleNamespace):
date: str
api_key: str | None
model: str | None
model_group: str | None
custom_llm_provider: str | None
mcp_namespaced_tool_name: str | None
endpoint: str | None
group_level: int
spend: float | None
prompt_tokens: int | None
completion_tokens: int | None
@ -172,12 +166,46 @@ class _GroupingSetsRow(SimpleNamespace):
timed_requests: int | None
class _EntityRollupRow(_GroupingSetsRow):
class _GroupingSetsRow(_RollupMetricsRow):
model: str | None
model_group: str | None
custom_llm_provider: str | None
mcp_namespaced_tool_name: str | None
endpoint: str | None
group_level: int
distinct_api_keys: int | None
class _EntityRollupRow(_RollupMetricsRow):
entity_id: str | None
api_key_rolled: int
def _reported_flat_cost(record: DailySpendRecord | _GroupingSetsRow) -> float:
class _AggregatedQueryKwargs(TypedDict):
table_name: ReadOnly[str]
entity_id_field: ReadOnly[str]
entity_id: ReadOnly[str | list[str] | None]
start_date: ReadOnly[str]
end_date: ReadOnly[str]
model: ReadOnly[str | None]
api_key: ReadOnly[str | list[str] | None]
exclude_entity_ids: ReadOnly[list[str] | None]
timezone_offset_minutes: ReadOnly[int | None]
include_current_utc_day: ReadOnly[bool]
_SqlQuery = tuple[str, list[str]]
async def _query_raw_optional(
prisma_client: PrismaClient, query: _SqlQuery | None
) -> list[dict[str, object]] | None: # mutable-ok: prisma query_raw return shape
if query is None:
return None
return await prisma_client.db.query_raw(query[0], *query[1])
def _reported_flat_cost(record: DailySpendRecord | _RollupMetricsRow) -> float:
"""Flat cost a daily row reports, which is zero unless PTU cost attribution is enabled.
Both read paths funnel through here: the paginated path reads the ``ptu_flat_cost``
@ -699,71 +727,8 @@ def _ptu_flat_cost_select(table_name: str) -> str:
return "0::float AS ptu_flat_cost"
def _build_aggregated_sql_query(
*,
table_name: str,
entity_id_field: str,
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
start_date: str,
end_date: str,
model: str | None,
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path
timezone_offset_minutes: int | None = None,
include_current_utc_day: bool = False,
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
"""Build a parameterized SQL GROUP BY query for aggregated daily activity.
Groups by (date, api_key, model, model_group, custom_llm_provider,
mcp_namespaced_tool_name, endpoint) with SUMs on all metric columns.
The entity_id column is intentionally omitted from GROUP BY to collapse
rows across entities this is where the biggest row reduction comes from.
Returns:
Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw().
"""
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
if pg_table is None:
raise ValueError(f"Unknown table name: {table_name}")
adjusted_start, adjusted_end = _adjust_dates_for_timezone(
start_date, end_date, timezone_offset_minutes, include_current_utc_day
)
where_clause, sql_params = _build_aggregated_where_clause(
entity_id_field=entity_id_field,
entity_id=entity_id,
adjusted_start=adjusted_start,
adjusted_end=adjusted_end,
model=model,
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
)
# Postgres computes every rollup level the response needs — per-date
# totals, per-(date, model), per-(date, model, api_key), per-provider,
# etc. — in a single pass via GROUPING SETS. The GROUPING() bitmask
# encodes which level a row belongs to so Python can dispatch rows
# straight into their buckets without re-summing. The leaf grouping
# is omitted on purpose: nothing in the response shape needs it once
# all the rollups are present.
#
# TODO: drop the successful_requests/failed_requests aggregates (and the
# total_successful_requests metadata they feed) once the admin UI reads SGR
# only from LiteLLM_DailyGatewayRequests. The remaining spend, token and
# api_requests rollups are still served from here.
sql_query: Final = f"""
SELECT
date,
api_key,
model,
COALESCE(NULLIF(model_group, ''), model) AS model_group,
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
GROUPING(date, api_key, model, COALESCE(NULLIF(model_group, ''), model),
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
def _rollup_metric_select(table_name: str) -> str:
return f"""
SUM(spend)::float AS spend,
{_ptu_flat_cost_select(table_name)},
SUM(prompt_tokens)::bigint AS prompt_tokens,
@ -779,27 +744,113 @@ def _build_aggregated_sql_query(
SUM(successful_requests)::bigint AS successful_requests,
SUM(failed_requests)::bigint AS failed_requests,
SUM(total_response_time_ms)::bigint AS total_response_time_ms,
SUM(timed_requests)::bigint AS timed_requests
SUM(timed_requests)::bigint AS timed_requests"""
_MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)"
def _build_aggregated_sql_query(
*,
table_name: str,
entity_id_field: str,
entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
start_date: str,
end_date: str,
model: str | None,
api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path
exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path
timezone_offset_minutes: int | None = None,
include_current_utc_day: bool = False,
) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params
"""Build the GROUPING SETS query for aggregated daily activity.
Returns:
Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw().
"""
pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name)
if pg_table is None:
raise ValueError(f"Unknown table name: {table_name}")
adjusted_start, adjusted_end = _adjust_dates_for_timezone(
start_date, end_date, timezone_offset_minutes, include_current_utc_day
)
where_clause, where_params = _build_aggregated_where_clause(
entity_id_field=entity_id_field,
entity_id=entity_id,
adjusted_start=adjusted_start,
adjusted_end=adjusted_end,
model=model,
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
)
sentinel_param: Final = f"${len(where_params) + 1}"
metric_select: Final = _rollup_metric_select(table_name)
# TODO: drop the successful_requests/failed_requests aggregates (and the
# total_successful_requests metadata they feed) once the admin UI reads SGR
# only from LiteLLM_DailyGatewayRequests. The remaining spend, token and
# api_requests rollups are still served from here.
sql_query: Final = f"""
(SELECT
date,
NULL::text AS api_key,
model,
{_MODEL_GROUP_EXPR} AS model_group,
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
(GROUPING(date) << 6) | {_API_KEY_ROLLED_UP_BIT}
| GROUPING(model, {_MODEL_GROUP_EXPR},
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
NULL::bigint AS distinct_api_keys,{metric_select}
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY GROUPING SETS (
(date),
(date, api_key),
(date, model),
(date, model, api_key),
(date, COALESCE(NULLIF(model_group, ''), model)),
(date, COALESCE(NULLIF(model_group, ''), model), api_key),
(date, {_MODEL_GROUP_EXPR}),
(date, custom_llm_provider),
(date, custom_llm_provider, api_key),
(date, mcp_namespaced_tool_name),
(date, mcp_namespaced_tool_name, api_key),
(date, endpoint),
(date, endpoint, api_key),
()
))
UNION ALL
(WITH top_api_keys AS (
SELECT api_key, COUNT(*) OVER () AS distinct_api_keys
FROM "{pg_table}"
WHERE {where_clause} AND api_key <> {sentinel_param}
GROUP BY api_key
ORDER BY SUM(spend) DESC, api_key
LIMIT {USAGE_TOP_API_KEYS_LIMIT}
)
SELECT
date,
api_key,
model,
{_MODEL_GROUP_EXPR} AS model_group,
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
GROUPING(date, api_key, model, {_MODEL_GROUP_EXPR},
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
MAX(top_api_keys.distinct_api_keys) AS distinct_api_keys,{metric_select}
FROM "{pg_table}" JOIN top_api_keys USING (api_key)
WHERE {where_clause}
GROUP BY GROUPING SETS (
(date, api_key),
(date, model, api_key),
(date, {_MODEL_GROUP_EXPR}, api_key),
(date, custom_llm_provider, api_key),
(date, mcp_namespaced_tool_name, api_key),
(date, endpoint, api_key)
))
"""
return sql_query, sql_params
return sql_query, [*where_params, PTU_SENTINEL_API_KEY]
def _build_entity_rollup_sql_query(
@ -844,23 +895,7 @@ def _build_entity_rollup_sql_query(
"{entity_id_field}" AS entity_id,
date,
api_key,
GROUPING(api_key) AS api_key_rolled,
SUM(spend)::float AS spend,
{_ptu_flat_cost_select(table_name)},
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens,
SUM(cache_creation_input_tokens)::bigint AS cache_creation_input_tokens,
SUM(compression_saved_tokens)::bigint AS compression_saved_tokens,
SUM(compression_savings_spend)::float AS compression_savings_spend,
SUM(prompt_caching_savings_spend)::float AS prompt_caching_savings_spend,
SUM(gateway_injected_caching_savings_spend)::float AS gateway_injected_caching_savings_spend,
SUM(autorouter_savings_spend)::float AS autorouter_savings_spend,
SUM(api_requests)::bigint AS api_requests,
SUM(successful_requests)::bigint AS successful_requests,
SUM(failed_requests)::bigint AS failed_requests,
SUM(total_response_time_ms)::bigint AS total_response_time_ms,
SUM(timed_requests)::bigint AS timed_requests
GROUPING(api_key) AS api_key_rolled,{_rollup_metric_select(table_name)}
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY GROUPING SETS (
@ -962,6 +997,7 @@ async def _aggregate_spend_records(
# current grouping set's key), 0 when the column is part of the key.
_GROUP_GRAND_TOTAL: Final = 127 # 0b1111111 — all rolled up
_GROUP_DATE: Final = 63 # 0b0111111 — only date kept
_API_KEY_ROLLED_UP_BIT: Final = 32 # 0b0100000
_GROUP_DATE_API_KEY: Final = 31 # 0b0011111
_GROUP_DATE_MODEL: Final = 47 # 0b0101111
_GROUP_DATE_MODEL_API_KEY: Final = 15 # 0b0001111
@ -975,7 +1011,7 @@ _GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110
_GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110
def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
def _record_to_spend_metrics(record: _RollupMetricsRow) -> SpendMetrics:
"""Build a SpendMetrics directly from one already-aggregated rollup row.
SUM() over zero rows is SQL NULL, so rollup rows (notably the grand-total
@ -1329,10 +1365,6 @@ async def get_daily_activity_aggregated(
) -> SpendAnalyticsPaginatedResponse:
"""Aggregated variant that returns the full result set (no pagination).
Uses SQL GROUP BY to aggregate rows in the database rather than fetching
all individual rows into Python. This collapses rows across entities
(users/teams/orgs), reducing ~150k rows to ~2-3k grouped rows.
include_entity_breakdown runs a small companion rollup query and folds
`breakdown.entities` onto the response, as entity-scoped views like Team Usage need.
@ -1351,7 +1383,7 @@ async def get_daily_activity_aggregated(
)
try:
sql_query, sql_params = _build_aggregated_sql_query(
query_kwargs: Final = _AggregatedQueryKwargs(
table_name=table_name,
entity_id_field=entity_id_field,
entity_id=entity_id,
@ -1363,36 +1395,16 @@ async def get_daily_activity_aggregated(
timezone_offset_minutes=timezone_offset_minutes,
include_current_utc_day=include_current_utc_day,
)
sql_query, sql_params = _build_aggregated_sql_query(**query_kwargs)
entity_query: Final = _build_entity_rollup_sql_query(**query_kwargs) if include_entity_breakdown else None
entity_query: Final = (
_build_entity_rollup_sql_query(
table_name=table_name,
entity_id_field=entity_id_field,
entity_id=entity_id,
start_date=start_date,
end_date=end_date,
model=model,
api_key=api_key,
exclude_entity_ids=exclude_entity_ids,
timezone_offset_minutes=timezone_offset_minutes,
include_current_utc_day=include_current_utc_day,
)
if include_entity_breakdown
else None
raw_rows, raw_entity_rows = await asyncio.gather(
prisma_client.db.query_raw(sql_query, *sql_params),
_query_raw_optional(prisma_client, entity_query),
)
# Execute the GROUPING SETS query (one row per rollup level), alongside
# the per-entity companion rollup when the caller wants entities.
raw_rows, raw_entity_rows = (
await asyncio.gather(
prisma_client.db.query_raw(sql_query, *sql_params),
prisma_client.db.query_raw(entity_query[0], *entity_query[1]),
)
if entity_query is not None
else (await prisma_client.db.query_raw(sql_query, *sql_params), None)
)
records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or [])]
records: Final = [_GroupingSetsRow(**row) for row in (raw_rows or ())]
total_api_keys: Final = next((r.distinct_api_keys for r in records if r.distinct_api_keys is not None), 0)
# The grouping-sets dispatcher places each row directly in its bucket
# using the row's GROUPING() bitmask. No Python-side summing needed.
@ -1446,6 +1458,8 @@ async def get_daily_activity_aggregated(
page=1,
total_pages=1,
has_more=False,
api_key_limit=USAGE_TOP_API_KEYS_LIMIT,
total_api_keys=total_api_keys,
),
)

View file

@ -510,6 +510,16 @@ def _get_user_in_team(team_table: LiteLLM_TeamTableCachedObj, user_id: str | Non
return None
def _get_caller_team_role(
team_table: LiteLLM_TeamTableCachedObj,
user_api_key_dict: UserAPIKeyAuth,
) -> Literal["admin", "user"] | None:
if user_api_key_dict.is_team_service_account and user_api_key_dict.team_id == team_table.team_id:
return "user"
member: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
return None if member is None else member.role
def _calculate_key_rotation_time(rotation_interval: str) -> datetime:
"""
Helper function to calculate the next rotation time for a key based on the rotation interval.
@ -604,7 +614,7 @@ def _team_key_operation_team_member_check(
detail=f"User={assigned_user_id} not assigned to team={team_table.team_id}",
)
team_member_object: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
caller_team_role: Final = _get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict)
is_admin: Final = (
user_api_key_dict.user_role is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
@ -612,22 +622,22 @@ def _team_key_operation_team_member_check(
if is_admin:
return True
elif team_member_object is None:
elif caller_team_role is None:
raise HTTPException(
status_code=400,
detail=f"User={user_api_key_dict.user_id} not assigned to team={team_table.team_id}",
)
elif (
"allowed_team_member_roles" in team_key_generation
and team_member_object.role not in team_key_generation["allowed_team_member_roles"]
and caller_team_role not in team_key_generation["allowed_team_member_roles"]
):
raise HTTPException(
status_code=400,
detail=f"Team member role {team_member_object.role} not in allowed_team_member_roles={team_key_generation['allowed_team_member_roles']}",
detail=f"Team member role {caller_team_role} not in allowed_team_member_roles={team_key_generation['allowed_team_member_roles']}",
)
TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint(
team_member_object=team_member_object,
team_member_role=caller_team_role,
team_table=team_table,
route=route,
)
@ -748,6 +758,12 @@ def key_generation_check(
Check if admin has restricted key creation to certain roles for teams or individuals
"""
if user_api_key_dict.is_team_service_account and data.team_id != user_api_key_dict.team_id:
raise HTTPException(
status_code=403,
detail=f"Service account keys can only create keys for their own team. team_id={user_api_key_dict.team_id}",
)
## check if key is for team or individual
is_team_key: Final = _is_team_key(data=data)
_is_admin: Final = (
@ -2233,6 +2249,14 @@ async def generate_service_account_key_fn(
prisma_client=prisma_client,
)
if data.metadata is None or data.metadata.get("service_account_id") is None:
service_account_id: Final = data.key_alias or str(uuid.uuid4())
stamped_metadata: Final = { # mutable-ok: GenerateKeyRequest.metadata is a plain dict field
**(data.metadata or MappingProxyType({})),
"service_account_id": service_account_id,
}
data.metadata = stamped_metadata # rebind-ok: the request carries the stamp so it persists on the key
verbose_proxy_logger.debug("entered /key/generate")
custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_generate_hook(
@ -3891,8 +3915,10 @@ async def validate_key_team_change(
detail=f"Key={key.token} has a rpm_limit={key.rpm_limit} which is greater than the team's rpm_limit={team.rpm_limit}.",
)
team_table: Final = cast(LiteLLM_TeamTableCachedObj, team)
# Check if the key's user_id is a member of the team
member_object: Final = _get_user_in_team(team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id)
member_object: Final = _get_user_in_team(team_table=team_table, user_id=key.user_id)
if key.user_id is not None:
if not member_object:
raise HTTPException(
@ -3908,8 +3934,8 @@ async def validate_key_team_change(
team_obj=team,
)
or TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint(
team_member_object=member_object,
team_table=cast(LiteLLM_TeamTableCachedObj, team),
team_member_role=None if member_object is None else member_object.role,
team_table=team_table,
route=KeyManagementRoutes.KEY_UPDATE.value,
)
):

View file

@ -47,7 +47,7 @@ except ImportError:
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.constants import LITELLM_PROXY_ADMIN_NAME, MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH
from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_DESCRIPTION,
LITELLM_MCP_SERVER_NAME,
@ -145,6 +145,7 @@ if MCP_AVAILABLE:
get_user_env_vars,
get_user_env_vars_bulk,
get_user_oauth_credential,
list_server_user_credentials,
list_user_oauth_credentials,
mcp_oauth_token_identity,
merge_user_env_vars,
@ -180,6 +181,7 @@ if MCP_AVAILABLE:
MCPApprovalStatus,
MCPOAuthUserCredentialRequest,
MCPOAuthUserCredentialStatus,
MCPServerUserCredentialListItem,
MCPSubmissionsSummary,
MCPTransport,
MCPUserCredentialListItem,
@ -221,6 +223,7 @@ if MCP_AVAILABLE:
MCPAuth,
MCPCredentials,
MCPGatewaySessionsResponse,
MCPGatewaySessionsTerminateResponse,
normalize_upstream_header_name,
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -662,6 +665,31 @@ if MCP_AVAILABLE:
"""
return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
def _resolve_credential_target_user_id(user_api_key_dict: UserAPIKeyAuth, requested_user_id: str | None) -> str:
"""The user whose stored MCP credential a request acts on.
Defaults to the caller. Naming another user is a revocation and needs
``PROXY_ADMIN``; a read-only admin or a regular user gets 403.
"""
caller_user_id: Final = user_api_key_dict.user_id or ""
if requested_user_id is not None and requested_user_id != caller_user_id:
if not _user_is_full_admin(user_api_key_dict):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
"error": "Proxy admin access required to revoke another user's MCP credential.",
},
)
return requested_user_id
if not caller_user_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": "User ID not found in token"
}, # mutable-ok: FastAPI HTTPException detail requires a plain dict
)
return caller_user_id
def _is_restricted_virtual_key_request(user_api_key_dict: UserAPIKeyAuth) -> bool:
"""Best-effort detection for route-restricted virtual keys.
@ -1373,6 +1401,41 @@ if MCP_AVAILABLE:
return get_mcp_gateway_sessions_report()
@router.delete(
"/sessions",
description=(
"Force-close live stateful MCP gateway sessions on this proxy worker, selected by session id prefix "
"and/or by the LiteLLM user that opened them (proxy admin only)."
),
dependencies=(Depends(user_api_key_auth),),
response_model=MCPGatewaySessionsTerminateResponse,
)
@management_endpoint_wrapper
async def delete_mcp_gateway_sessions(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
session_id_prefix: Annotated[str | None, Query(min_length=MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH)] = None,
user_id: Annotated[str | None, Query(min_length=1)] = None,
) -> MCPGatewaySessionsTerminateResponse:
if not _user_is_full_admin(user_api_key_dict):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
"error": "Proxy admin access required to terminate MCP gateway sessions.",
},
)
if session_id_prefix is None and user_id is None:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
"error": "Provide session_id_prefix and/or user_id to select the sessions to terminate.",
},
)
from litellm.proxy._experimental.mcp_server.server import (
terminate_mcp_gateway_sessions,
)
return await terminate_mcp_gateway_sessions(session_id_prefix=session_id_prefix, user_id=user_id)
@router.get(
"/server/submissions",
description="Returns all MCP servers submitted by non-admin users (admin review queue). Mirrors GET /guardrails/submissions.",
@ -2254,14 +2317,17 @@ if MCP_AVAILABLE:
_invalidate_byok_cred_cache,
)
_invalidate_byok_cred_cache(user_id, server_id)
await _invalidate_byok_cred_cache(user_id, server_id)
return MCPUserCredentialResponse(server_id=server_id, has_credential=True)
# save=False: credential not persisted
return MCPUserCredentialResponse(server_id=server_id, has_credential=False)
@router.delete(
"/server/{server_id}/user-credential",
description="Delete the calling user's stored API key for a BYOK MCP server",
description=(
"Delete the calling user's stored API key for a BYOK MCP server. "
"A proxy admin may pass user_id to revoke another user's stored key."
),
dependencies=[Depends(user_api_key_auth)],
response_model=MCPUserCredentialResponse,
)
@ -2269,24 +2335,20 @@ if MCP_AVAILABLE:
async def delete_mcp_user_credential(
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_id: Annotated[str | None, Query(min_length=1)] = None,
):
"""Remove the calling user's BYOK credential."""
"""Remove the target user's BYOK credential (the caller unless an admin names another user)."""
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
user_id: Final = user_api_key_dict.user_id or ""
if not user_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "User ID not found in token"},
)
target_user_id: Final = _resolve_credential_target_user_id(user_api_key_dict, user_id)
try:
await delete_user_credential(prisma_client, user_id, server_id)
await delete_user_credential(prisma_client, target_user_id, server_id)
except RecordNotFoundError:
pass # Already deleted or didn't exist
from litellm.proxy._experimental.mcp_server.server import (
_invalidate_byok_cred_cache,
)
_invalidate_byok_cred_cache(user_id, server_id)
await _invalidate_byok_cred_cache(target_user_id, server_id)
return MCPUserCredentialResponse(server_id=server_id, has_credential=False)
# ── OAuth2 user-credential endpoints ──────────────────────────────────────
@ -2362,7 +2424,10 @@ if MCP_AVAILABLE:
@router.delete(
"/server/{server_id}/oauth-user-credential",
description="Revoke the calling user's stored OAuth2 token for an MCP server",
description=(
"Revoke the calling user's stored OAuth2 token for an MCP server. "
"A proxy admin may pass user_id to revoke another user's stored token."
),
dependencies=[Depends(user_api_key_auth)],
response_model=MCPOAuthUserCredentialStatus,
)
@ -2370,29 +2435,25 @@ if MCP_AVAILABLE:
async def delete_mcp_oauth_user_credential(
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_id: Annotated[str | None, Query(min_length=1)] = None,
):
"""Revoke/delete the user's OAuth2 credential."""
"""Revoke the target user's OAuth2 credential (the caller unless an admin names another user)."""
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
user_id: Final = user_api_key_dict.user_id or ""
if not user_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "User ID not found in token"},
)
target_user_id: Final = _resolve_credential_target_user_id(user_api_key_dict, user_id)
# Only delete if the stored credential is actually an OAuth2 token.
# This prevents accidentally deleting a BYOK credential if one exists
# for the same (user_id, server_id) pair.
cred_to_delete: Final = await get_user_oauth_credential(prisma_client, user_id, server_id)
cred_to_delete: Final = await get_user_oauth_credential(prisma_client, target_user_id, server_id)
if cred_to_delete is not None:
try:
await delete_user_credential(prisma_client, user_id, server_id)
await delete_user_credential(prisma_client, target_user_id, server_id)
except RecordNotFoundError:
pass # Already gone — treat as a successful delete
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415
global_mcp_server_manager,
)
await global_mcp_server_manager.invalidate_user_oauth_token_cache(user_id, server_id)
await global_mcp_server_manager.invalidate_user_oauth_token_cache(target_user_id, server_id)
return MCPOAuthUserCredentialStatus(
server_id=server_id,
has_credential=False,
@ -2481,6 +2542,30 @@ if MCP_AVAILABLE:
)
return items
@router.get(
"/server/{server_id}/user-credentials",
description="List every user's stored BYOK or OAuth2 credential for an MCP server (admin only, no secrets)",
dependencies=(Depends(user_api_key_auth),),
response_model=list[MCPServerUserCredentialListItem],
)
@management_endpoint_wrapper
async def list_mcp_server_user_credentials(
server_id: str,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> tuple[MCPServerUserCredentialListItem, ...]:
if user_api_key_dict.user_role not in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
"error": "Admin access required to view MCP server user credentials.",
},
)
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
return await list_server_user_credentials(prisma_client, server_id)
# ── Per-user MCP env var endpoints ────────────────────────────────────────
async def _authorize_and_fetch_mcp_server(

View file

@ -2107,7 +2107,7 @@ def _handle_multi_valued_attribute_update(path: str, op_type: str, value: object
except ValidationError:
raise HTTPException(
status_code=400,
detail={"error": f"Invalid value for {base}: expected a list of objects with a 'value' sub-attribute"},
detail={"error": f"Invalid value for {base}: expected a list of objects or strings"},
)
dumped: Final = [attr.model_dump(exclude_none=True) for attr in attrs]

View file

@ -2904,10 +2904,15 @@ async def _process_team_members(
if member_allowed_models is None and team_default_member_models:
member_allowed_models = team_default_member_models
if isinstance(data.member, Member):
requested_members: Final[Sequence[Member]] = (
(data.member,) if isinstance(data.member, Member) else tuple(data.member)
)
for m in requested_members:
if _member_already_in_team(m, complete_team_data):
continue
try:
updated_user, updated_tm = await add_new_member(
new_member=data.member,
new_member=m,
max_budget_in_team=data.max_budget_in_team,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
@ -2921,34 +2926,11 @@ async def _process_team_members(
except Exception as e:
raise HTTPException(
status_code=500,
detail={"error": f"Unable to add user - {data.member}, to team - {data.team_id}, for reason - {e}"},
detail={"error": f"Unable to add user - {m}, to team - {data.team_id}, for reason - {e}"},
)
updated_users.append(updated_user)
if updated_tm is not None:
updated_team_memberships.append(updated_tm)
elif isinstance(data.member, list):
for m in data.member:
try:
updated_user, updated_tm = await add_new_member(
new_member=m,
max_budget_in_team=data.max_budget_in_team,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
team_id=data.team_id,
default_team_budget_id=default_team_budget_id,
allowed_models=member_allowed_models,
budget_duration=data.budget_duration,
tx=tx,
)
except Exception as e:
raise HTTPException(
status_code=500,
detail={"error": f"Unable to add user - {m}, to team - {data.team_id}, for reason - {e}"},
)
updated_users.append(updated_user)
if updated_tm is not None:
updated_team_memberships.append(updated_tm)
return updated_users, updated_team_memberships

View file

@ -1,4 +1,4 @@
from typing import Final
from typing import Final, Literal
from litellm.proxy._types import (
KeyManagementRoutes,
@ -6,7 +6,6 @@ from litellm.proxy._types import (
LiteLLM_VerificationToken,
LiteLLMRoutes,
LitellmUserRoles,
Member,
ProxyErrorTypes,
ProxyException,
UserAPIKeyAuth,
@ -27,7 +26,6 @@ DEFAULT_TEAM_MEMBER_PERMISSIONS: Final = BASELINE_TEAM_MEMBER_PERMISSIONS
class TeamMemberPermissionChecks:
@staticmethod
def get_permissions_for_team_member(
team_member_object: Member,
team_table: LiteLLM_TeamTableCachedObj,
) -> list[KeyManagementRoutes]:
"""
@ -67,7 +65,7 @@ class TeamMemberPermissionChecks:
Main handler for checking if a team member can update a key
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
_get_user_in_team,
_get_caller_team_role,
)
# 1. Don't execute these checks if the user role is proxy admin
@ -87,12 +85,11 @@ class TeamMemberPermissionChecks:
check_db_only=True,
)
# 4. Extract `Member` object from `team_table`
key_assigned_user_in_team: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
caller_team_role: Final = _get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict)
# 5. Check if the team member has permissions for the endpoint
# 4. Check if the team member has permissions for the endpoint
has_permission: Final = TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint(
team_member_object=key_assigned_user_in_team,
team_member_role=caller_team_role,
team_table=team_table,
route=route,
)
@ -106,7 +103,7 @@ class TeamMemberPermissionChecks:
@staticmethod
def does_team_member_have_permissions_for_endpoint(
team_member_object: Member | None,
team_member_role: Literal["admin", "user"] | None,
team_table: LiteLLM_TeamTableCachedObj,
route: str,
) -> bool | None:
@ -116,13 +113,12 @@ class TeamMemberPermissionChecks:
# permission checks only run for non-admin users
# Non-Admin user trying to access information about a team's key
if team_member_object is None:
if team_member_role is None:
return False
if team_member_object.role == "admin":
if team_member_role == "admin":
return True
_team_member_permissions: Final = TeamMemberPermissionChecks.get_permissions_for_team_member(
team_member_object=team_member_object,
team_table=team_table,
)
team_member_permissions = TeamMemberPermissionChecks._get_list_of_route_enum_as_str(_team_member_permissions)
@ -156,7 +152,7 @@ class TeamMemberPermissionChecks:
from fastapi import HTTPException
from litellm.proxy.management_endpoints.key_management_endpoints import (
_get_user_in_team,
_get_caller_team_role,
)
# No-op when the request does not assign any access groups.
@ -177,20 +173,19 @@ class TeamMemberPermissionChecks:
),
)
team_member_object: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
caller_team_role: Final = _get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict)
# Team admins always bypass (consistent with other member-permission checks).
if team_member_object is not None and team_member_object.role == "admin":
if caller_team_role == "admin":
return
permissions: Final = (
TeamMemberPermissionChecks._get_list_of_route_enum_as_str(
TeamMemberPermissionChecks.get_permissions_for_team_member(
team_member_object=team_member_object,
team_table=team_table,
)
)
if team_member_object is not None
if caller_team_role is not None
else []
)
@ -214,7 +209,7 @@ class TeamMemberPermissionChecks:
Returns True if the user belongs to the team that the key is assigned to
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
_get_user_in_team,
_get_caller_team_role,
)
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
@ -228,9 +223,8 @@ class TeamMemberPermissionChecks:
check_db_only=True,
)
# 4. Extract `Member` object from `team_table`
team_member_object: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
return team_member_object is not None
caller_team_role: Final = _get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict)
return caller_team_role is not None
@staticmethod
def get_all_available_team_member_permissions() -> list[str]:

View file

@ -86,7 +86,9 @@ class _PrismaUserTable(Protocol):
class _PrismaTeamMembershipTable(Protocol):
"""Team membership table actions the management helpers issue."""
async def create(self, *, data: Mapping[str, object], include: Mapping[str, bool]) -> _PrismaRecord: ...
async def upsert(
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]], include: Mapping[str, bool]
) -> _PrismaRecord: ...
class MemberWriteTx(Protocol):
@ -348,7 +350,7 @@ async def _resolve_member_budget_id(
default member budget is cloned (with ``budget_duration`` overriding its
reset window while keeping its other limits). A lone ``budget_duration``
with no team default creates a window-only budget. With nothing set the
member gets no budget.
member gets no budget, though ``add_new_member`` still writes its membership row.
"""
has_explicit_limit: Final = max_budget_in_team is not None or allowed_models is not None
@ -415,9 +417,9 @@ async def add_new_member(
Add a new member to a team
- add team id to user table
- add team member w/ budget to team member table
- add team member to team member table, linked to a budget when one resolves
Returns created/existing user + team membership w/ budget id
Returns created/existing user + team membership (``budget_id`` is ``None`` when no budget applies)
Callers already inside a transaction pass it as ``tx`` so every write here runs on that
connection instead of borrowing more from the pool while the caller's locks are held.
@ -471,14 +473,15 @@ async def add_new_member(
tx=tx,
)
if _budget_id and returned_user is not None and returned_user.user_id is not None:
if returned_user is not None and returned_user.user_id is not None:
membership_table: Final[_PrismaTeamMembershipTable] = _team_membership_table(prisma_client, tx)
_returned_team_membership: Final = await membership_table.create(
data={
"team_id": team_id,
"user_id": returned_user.user_id,
"budget_id": _budget_id,
},
membership_key: Final[Mapping[str, object]] = {"user_id": returned_user.user_id, "team_id": team_id}
budget_link: Final[Mapping[str, str]] = (
MappingProxyType({"budget_id": _budget_id}) if _budget_id is not None else MappingProxyType({})
)
_returned_team_membership: Final = await membership_table.upsert(
where={"user_id_team_id": membership_key},
data={"create": {**membership_key, **budget_link}, "update": {}},
include={"litellm_budget_table": True},
)

View file

@ -12,6 +12,7 @@ import hmac
import inspect
import json
import os
import posixpath
import re
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
from dataclasses import dataclass
@ -30,6 +31,14 @@ from litellm import get_llm_provider
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS,
AZURE_SPEECH_BATCH_PATH_PREFIX,
AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN,
AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
AZURE_SPEECH_FAST_TRANSCRIPTION_PATH,
AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX,
AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX,
AZURE_SPEECH_STT_DOMAIN,
AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER,
BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES,
)
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
@ -58,6 +67,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
get_request_body,
is_json_content_type,
)
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
from litellm.proxy.common_utils.sse_keepalive import (
wrap_passthrough_sse_bytes_with_keepalive_pings,
)
@ -1358,6 +1368,145 @@ async def comprehend_medical_sdk_proxy_route(
)
AZURE_SPEECH_FORWARDED_REQUEST_HEADERS: Final = ("content-type", "accept")
AZURE_SPEECH_ENDPOINT_FAMILY_DOMAINS: Final = MappingProxyType(
{
AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX: AZURE_SPEECH_STT_DOMAIN,
AZURE_SPEECH_BATCH_PATH_PREFIX: AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN,
}
)
def resolve_azure_speech_base_url(endpoint_path: str, api_base: str | None, region: str | None) -> httpx.URL | None:
"""
Azure AI Speech serves the two REST families from different regional hosts: short-audio
recognition under ``{region}.stt.speech.microsoft.com`` and batch transcription under
``{region}.api.cognitive.microsoft.com``. An operator-configured ``api_base`` (custom
domain or private endpoint) serves both and wins over the region. Returns ``None`` when
the path is outside both families so the operator key is never sent for an unknown API.
"""
domain: Final = next(
(
family_domain
for family_prefix, family_domain in AZURE_SPEECH_ENDPOINT_FAMILY_DOMAINS.items()
if endpoint_path.startswith(family_prefix)
),
None,
)
if domain is None:
return None
if api_base:
return httpx.URL(api_base)
if not region:
return None
return httpx.URL(f"https://{region}.{domain}")
def azure_speech_path_manages_shared_resources(endpoint_path: str) -> bool:
return (
endpoint_path.startswith(AZURE_SPEECH_BATCH_PATH_PREFIX)
and endpoint_path != AZURE_SPEECH_FAST_TRANSCRIPTION_PATH
)
def canonical_azure_speech_endpoint_path(endpoint: str) -> str:
"""
The path Azure will actually serve, with ``.`` and ``..`` segments resolved, so the
endpoint family and the admin guard are decided on the same path the upstream request uses.
"""
raw_path: Final = httpx.URL(endpoint).path
resolved_path: Final = posixpath.normpath(f"/{raw_path.lstrip('/')}")
if raw_path.endswith("/") and resolved_path != "/":
return f"{resolved_path}/"
return resolved_path
@router.api_route(
f"{AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX}/{{endpoint:path}}",
methods=["GET", "POST", "PUT", "DELETE", "PATCH"], # mutable-ok: fastapi route methods must be a list
tags=["Azure AI Speech Pass-through", "pass-through"], # mutable-ok: fastapi route tags must be a list
)
async def azure_speech_proxy_route(
endpoint: str,
request: Request,
fastapi_response: Response,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Pass-through for the Azure AI Speech REST APIs (speech to text), e.g.
`POST /azure_speech/speech/recognition/conversation/cognitiveservices/v1?language=en-US`
with the raw audio as the body, or `POST /azure_speech/speechtotext/v3.2/transcriptions`.
The body is forwarded byte for byte and the proxy injects its own
`Ocp-Apim-Subscription-Key`; the caller's `Authorization` header is the LiteLLM key
and is never forwarded.
[Docs](https://docs.litellm.ai/docs/pass_through/azure_speech)
"""
normalized_endpoint_path: Final = canonical_azure_speech_endpoint_path(endpoint)
base_url: Final = resolve_azure_speech_base_url(
endpoint_path=normalized_endpoint_path,
api_base=get_secret_str(secret_name="AZURE_SPEECH_API_BASE"),
region=get_secret_str(secret_name="AZURE_SPEECH_REGION"),
)
if base_url is None:
raise HTTPException(
status_code=400,
detail=(
f"Unsupported Azure Speech path: {normalized_endpoint_path}. Supported prefixes are "
f"{AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX} and {AZURE_SPEECH_BATCH_PATH_PREFIX}; set "
"AZURE_SPEECH_REGION or AZURE_SPEECH_API_BASE in the proxy environment."
),
)
if azure_speech_path_manages_shared_resources(normalized_endpoint_path) and not is_proxy_admin(user_api_key_dict):
raise HTTPException(
status_code=403,
detail=(
f"{request.method} {normalized_endpoint_path} manages batch transcription resources that belong to "
"the proxy's Azure Speech subscription and whose cost is unknown at request time, so it is limited "
f"to proxy admin keys. Use {AZURE_SPEECH_FAST_TRANSCRIPTION_PATH} for transcription that is priced "
"per request."
),
)
azure_speech_api_key: Final = passthrough_endpoint_router.get_credentials(
custom_llm_provider=AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
region_name=None,
)
if azure_speech_api_key is None:
raise HTTPException(
status_code=400,
detail="Azure Speech credentials not found. Set AZURE_SPEECH_API_KEY in the proxy environment.",
)
target_url: Final = base_url.copy_with(
path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, normalized_endpoint_path)
)
request_headers: Final = _safe_get_request_headers(request)
upstream_headers: Final = MappingProxyType(
{
header_name: header_value
for header_name, header_value in (
*(
(header_name, request_headers[header_name])
for header_name in AZURE_SPEECH_FORWARDED_REQUEST_HEADERS
if header_name in request_headers
),
(AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER, azure_speech_api_key),
)
}
)
raw_body: Final = await request.body()
endpoint_func: Final = create_pass_through_route(
endpoint=endpoint,
target=str(target_url),
custom_headers=upstream_headers,
custom_llm_provider=AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
)
setattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, raw_body)
return await endpoint_func(request, fastapi_response, user_api_key_dict)
@router.post(
"/transcribe/{operation}",
tags=["Amazon Transcribe Pass-through", "pass-through"], # mutable-ok: fastapi route tags must be a list

View file

@ -0,0 +1,170 @@
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import Final
from urllib.parse import urlparse
import httpx
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
AZURE_SPEECH_BATCH_MODEL,
AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL,
AZURE_SPEECH_FAST_TRANSCRIPTION_PATH,
AZURE_SPEECH_MILLISECONDS_PER_SECOND,
AZURE_SPEECH_PRICING_MODEL,
AZURE_SPEECH_SHORT_AUDIO_MODEL,
AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX,
AZURE_SPEECH_TICKS_PER_SECOND,
)
from litellm.cost_calculator import transcription_cost
from litellm.litellm_core_utils.audio_utils.utils import calculate_request_duration
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import (
get_standard_logging_object_payload,
)
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.types.utils import StandardPassThroughResponseObject
class AzureSpeechPassthroughLoggingHandler:
@staticmethod
def _is_short_audio_route(url_route: str) -> bool:
return urlparse(url_route).path.startswith(AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX)
@staticmethod
def _is_fast_transcription_route(url_route: str) -> bool:
return urlparse(url_route).path.endswith(AZURE_SPEECH_FAST_TRANSCRIPTION_PATH)
@staticmethod
def _model_from_url_route(url_route: str) -> str:
if AzureSpeechPassthroughLoggingHandler._is_short_audio_route(url_route):
return f"{AZURE_SPEECH_CUSTOM_LLM_PROVIDER}/{AZURE_SPEECH_SHORT_AUDIO_MODEL}"
if AzureSpeechPassthroughLoggingHandler._is_fast_transcription_route(url_route):
return f"{AZURE_SPEECH_CUSTOM_LLM_PROVIDER}/{AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL}"
return f"{AZURE_SPEECH_CUSTOM_LLM_PROVIDER}/{AZURE_SPEECH_BATCH_MODEL}"
@staticmethod
def _recognized_audio_seconds(response_body: Mapping[str, object] | Sequence[object] | None) -> float:
if not isinstance(response_body, Mapping):
return 0.0
offset: Final = response_body.get("Offset")
duration: Final = response_body.get("Duration")
if not isinstance(offset, int) or not isinstance(duration, int):
return 0.0
return (offset + duration) / AZURE_SPEECH_TICKS_PER_SECOND
@staticmethod
def _uploaded_audio_seconds(httpx_response: httpx.Response) -> float:
try:
uploaded_audio: Final = httpx_response.request.content
except RuntimeError:
return 0.0
return calculate_request_duration(uploaded_audio) or 0.0
@staticmethod
def _short_audio_seconds(
httpx_response: httpx.Response, response_body: Mapping[str, object] | Sequence[object] | None
) -> float:
return max(
AzureSpeechPassthroughLoggingHandler._uploaded_audio_seconds(httpx_response),
AzureSpeechPassthroughLoggingHandler._recognized_audio_seconds(response_body),
)
@staticmethod
def _fast_transcription_audio_seconds(response_body: Mapping[str, object] | Sequence[object] | None) -> float:
if not isinstance(response_body, Mapping):
return 0.0
duration_milliseconds: Final = response_body.get("durationMilliseconds")
if not isinstance(duration_milliseconds, int):
return 0.0
return duration_milliseconds / AZURE_SPEECH_MILLISECONDS_PER_SECOND
@staticmethod
def _billed_audio_seconds(
url_route: str,
httpx_response: httpx.Response,
response_body: Mapping[str, object] | Sequence[object] | None,
) -> float:
if AzureSpeechPassthroughLoggingHandler._is_short_audio_route(url_route):
return AzureSpeechPassthroughLoggingHandler._short_audio_seconds(httpx_response, response_body)
if AzureSpeechPassthroughLoggingHandler._is_fast_transcription_route(url_route):
return AzureSpeechPassthroughLoggingHandler._fast_transcription_audio_seconds(response_body)
return 0.0
@staticmethod
def _response_cost(
url_route: str,
httpx_response: httpx.Response,
response_body: Mapping[str, object] | Sequence[object] | None,
) -> float:
audio_seconds: Final = AzureSpeechPassthroughLoggingHandler._billed_audio_seconds(
url_route, httpx_response, response_body
)
if audio_seconds <= 0.0:
return 0.0
try:
prompt_cost, completion_cost = transcription_cost(
model=AZURE_SPEECH_PRICING_MODEL,
custom_llm_provider="azure",
duration=audio_seconds,
)
except Exception as e: # noqa: BLE001 # a missing price entry must not drop the spend log row
verbose_proxy_logger.warning(
"No price for %s, logging Azure Speech call at zero cost: %s", AZURE_SPEECH_PRICING_MODEL, e
)
return 0.0
return prompt_cost + completion_cost
@staticmethod
def azure_speech_passthrough_handler(
httpx_response: httpx.Response,
response_body: Mapping[str, object] | Sequence[object] | None,
logging_obj: LiteLLMLoggingObj,
url_route: str,
result: str,
start_time: datetime,
end_time: datetime,
cache_hit: bool,
request_body: Mapping[str, object],
**kwargs: object, # kwargs-ok: the passthrough logging dispatch forwards shared logging kwargs to every handler
) -> PassThroughEndpointLoggingTypedDict:
try:
model_name: Final = AzureSpeechPassthroughLoggingHandler._model_from_url_route(url_route)
response_cost: Final = AzureSpeechPassthroughLoggingHandler._response_cost(
url_route, httpx_response, response_body
)
updated_kwargs: Final = { # mutable-ok: the logging pipeline requires a plain kwargs dict
**kwargs,
"model": model_name,
"custom_llm_provider": AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
"response_cost": response_cost,
}
logging_obj.model_call_details.update(
model=model_name,
custom_llm_provider=AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
response_cost=response_cost,
)
standard_logging_object: Final = get_standard_logging_object_payload(
kwargs=updated_kwargs,
init_response_obj=StandardPassThroughResponseObject(response=result),
start_time=start_time,
end_time=end_time,
logging_obj=logging_obj,
status="success",
)
handler_payload: Final[PassThroughEndpointLoggingTypedDict] = {
"result": StandardPassThroughResponseObject(response=result),
"kwargs": {**updated_kwargs, "standard_logging_object": standard_logging_object},
}
except Exception as e: # noqa: BLE001 # logging must never fail the forwarded request
verbose_proxy_logger.exception("Error in Azure Speech passthrough logging handler: %s", e)
fallback_payload: Final[PassThroughEndpointLoggingTypedDict] = {
"result": StandardPassThroughResponseObject(response=result),
"kwargs": kwargs,
}
return fallback_payload
return handler_payload

View file

@ -53,6 +53,7 @@ from litellm.litellm_core_utils.core_helpers import (
from litellm.litellm_core_utils.initialize_dynamic_callback_params import validate_no_callback_env_reference
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.base_llm.managed_resources.utils import (
@ -1024,7 +1025,7 @@ async def pass_through_request(
verbose_proxy_logger.debug(
"Pass through endpoint sending request to \nURL %s\nheaders: %s\nbody: %s\n",
url,
upstream_headers,
_get_masked_values(upstream_headers),
_parsed_body,
)

View file

@ -6,6 +6,7 @@ from urllib.parse import urlparse
import httpx
from litellm.constants import AZURE_SPEECH_CUSTOM_LLM_PROVIDER
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import PassThroughEndpointLoggingResultValues
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
@ -274,6 +275,25 @@ class PassThroughEndpointLogging:
)
standard_logging_response_object = comprehend_medical_handler_result["result"] # rebind-ok: elif-chain
kwargs = comprehend_medical_handler_result["kwargs"] # rebind-ok: elif-chain contract
elif self.is_azure_speech_route(custom_llm_provider):
from .llm_provider_handlers.azure_speech_passthrough_logging_handler import (
AzureSpeechPassthroughLoggingHandler,
)
azure_speech_handler_result: Final = AzureSpeechPassthroughLoggingHandler.azure_speech_passthrough_handler(
httpx_response=httpx_response,
response_body=response_body,
logging_obj=logging_obj,
url_route=url_route,
result=result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
request_body=request_body,
**kwargs,
)
standard_logging_response_object = azure_speech_handler_result["result"] # rebind-ok: elif-chain
kwargs = azure_speech_handler_result["kwargs"] # rebind-ok: elif-chain contract
elif self.is_transcribe_route(custom_llm_provider):
transcribe_handler_result: Final = TranscribePassthroughLoggingHandler.transcribe_passthrough_handler(
httpx_response=httpx_response,
@ -351,7 +371,7 @@ class PassThroughEndpointLogging:
):
standard_logging_response_object: PassThroughEndpointLoggingResultValues | None = None
logging_obj.model_call_details["passthrough_logging_payload"] = passthrough_logging_payload
if self.is_assemblyai_route(url_route):
if self.is_assemblyai_route(url_route) and not self.is_azure_speech_route(custom_llm_provider):
if AssemblyAIPassthroughLoggingHandler._should_log_request(httpx_response.request.method) is not True:
return
self.assemblyai_passthrough_logging_handler.assemblyai_passthrough_logging_handler(
@ -458,6 +478,9 @@ class PassThroughEndpointLogging:
def is_comprehend_medical_route(self, custom_llm_provider: str | None) -> bool:
return custom_llm_provider == "comprehendmedical"
def is_azure_speech_route(self, custom_llm_provider: str | None) -> bool:
return custom_llm_provider == AZURE_SPEECH_CUSTOM_LLM_PROVIDER
def is_transcribe_route(self, custom_llm_provider: str | None) -> bool:
return custom_llm_provider == TRANSCRIBE_CUSTOM_LLM_PROVIDER

View file

@ -1411,6 +1411,8 @@ def run_server(
# DO NOT DELETE - enables global variables to work across files
from litellm.proxy.proxy_server import app
os.environ["NUM_WORKERS"] = str(num_workers)
# Auto-create PROMETHEUS_MULTIPROC_DIR for multi-worker setups
prometheus_multiproc_dir: Final = ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
num_workers=num_workers,

View file

@ -154,7 +154,7 @@ from litellm.types.utils import (
TextCompletionResponse,
TokenCountResponse,
)
from litellm.utils import load_credentials_from_list
from litellm.utils import cost_map_omits_token_price, load_credentials_from_list
if TYPE_CHECKING:
from aiohttp import ClientSession
@ -309,6 +309,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.proxy._experimental.mcp_server.byok_credential_cache import byok_credential_cache
from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot
from litellm.proxy._types import *
from litellm.proxy.analytics_endpoints.analytics_endpoints import (
@ -330,6 +331,12 @@ from litellm.proxy.auth.fallback_budget import router_fallback_budget_check
from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck
from litellm.proxy.auth.login_throttle import (
LoginThrottle,
declared_proxy_ranges,
warn_login_counters_are_per_worker,
warn_source_login_limit_is_off,
)
from litellm.proxy.auth.model_checks import (
expand_wildcard_deployments_for_model_info,
get_all_fallbacks,
@ -827,6 +834,7 @@ from fastapi.openapi.docs import get_swagger_ui_html
from fastapi.openapi.utils import get_openapi
from fastapi.responses import (
FileResponse,
HTMLResponse,
JSONResponse,
ORJSONResponse,
RedirectResponse,
@ -6050,6 +6058,12 @@ class ProxyConfig:
general_settings = config.get("general_settings", {})
if general_settings is None:
general_settings = {}
if os.getenv("NUM_WORKERS", "1") != "1" and redis_usage_cache is None:
warn_login_counters_are_per_worker(os.getenv("NUM_WORKERS", "1"))
if declared_proxy_ranges(general_settings) is None:
warn_source_login_limit_is_off()
_bg_hc_model_groups: Final = parse_background_health_check_model_groups(general_settings)
_enable_hc_routing = False
_hc_staleness = None
@ -7537,7 +7551,7 @@ class ProxyConfig:
subscriber: Final = AuthCacheInvalidationSubscriber(
redis_cache=redis_cache,
user_api_key_cache=user_api_key_cache,
additional_in_memory_caches=(spend_counter_cache.in_memory_cache,),
additional_in_memory_caches=(spend_counter_cache.in_memory_cache, byok_credential_cache),
)
self.auth_cache_invalidation_subscriber = subscriber
subscriber.start()
@ -13618,9 +13632,10 @@ def _enrich_model_info_with_litellm_data(
discovered_model_info: Final = (
llm_router.get_discovered_model_info(model_info.get("id")) if llm_router is not None else MappingProxyType({})
)
unpriced: Final = cost_map_omits_token_price(model_info.get("id"), litellm_model_info.get("key"))
for k, v in MappingProxyType({**litellm_model_info, **discovered_model_info}).items():
if k not in model_info or (model_info[k] is None and k in discovered_model_info):
model_info[k] = v
model_info[k] = None if unpriced and k in ("input_cost_per_token", "output_cost_per_token") else v
model["model_info"] = model_info
# don't return the api key / vertex credentials
# don't return the llm credentials
@ -15056,45 +15071,7 @@ def _translate_model_name_for_response(model: dict) -> dict:
def _get_proxy_model_info(model: dict) -> dict:
# provided model_info in config.yaml
model_info: Final = model.get("model_info", {})
# read litellm model_prices_and_context_window.json to get the following:
# input_cost_per_token, output_cost_per_token, max_tokens
litellm_model_info = get_litellm_model_info(model=model)
# 2nd pass on the model, try seeing if we can find model in litellm model_cost map
if litellm_model_info == {}:
# use litellm_param model_name to get model_info
litellm_params = model.get("litellm_params", {})
litellm_model = litellm_params.get("model", None)
try:
litellm_model_info = litellm.get_model_info(model=litellm_model)
except Exception:
litellm_model_info = {}
# 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map
if litellm_model_info == {}:
# use litellm_param model_name to get model_info
litellm_params = model.get("litellm_params", {})
litellm_model = litellm_params.get("model", None)
split_model: Final = litellm_model.split("/")
if len(split_model) > 0:
litellm_model = split_model[-1]
try:
litellm_model_info = litellm.get_model_info(model=litellm_model, custom_llm_provider=split_model[0])
except Exception:
litellm_model_info = {}
discovered_model_info: Final = (
llm_router.get_discovered_model_info(model_info.get("id")) if llm_router is not None else MappingProxyType({})
)
for k, v in MappingProxyType({**litellm_model_info, **discovered_model_info}).items():
if k not in model_info or (model_info[k] is None and k in discovered_model_info):
model_info[k] = v
model["model_info"] = model_info
# don't return the llm credentials
model = remove_sensitive_info_from_deployment(deployment_dict=model, excluded_keys={"litellm_credential_name"})
return _translate_model_name_for_response(model)
return _translate_model_name_for_response(_enrich_model_info_with_litellm_data(model=model, llm_router=llm_router))
def _model_info_json_response(data: Sequence[Mapping[str, object]] | Mapping[str, object]) -> Response:
@ -15901,8 +15878,6 @@ async def fallback_login(request: Request):
else:
redirect_url += "/sso/callback"
from fastapi.responses import HTMLResponse
hide_default_credentials_hint: Final = should_hide_default_credentials_hint(general_settings)
return HTMLResponse(
content=build_ui_login_form(
@ -15924,13 +15899,27 @@ async def login(request: Request):
password: Final = str(form.get("password"))
# Authenticate user and get login result
login_result: Final = await authenticate_user(
username=username,
password=password,
master_key=master_key,
prisma_client=prisma_client,
general_settings=general_settings,
)
try:
login_result: Final = await authenticate_user(
username=username,
password=password,
master_key=master_key,
prisma_client=prisma_client,
throttle=LoginThrottle.from_request(request, general_settings, redis_usage_cache),
general_settings=general_settings,
)
except ProxyException as exc:
if int(exc.code) != status.HTTP_429_TOO_MANY_REQUESTS:
raise
retry_after: Final = exc.headers.get("Retry-After", "30")
return HTMLResponse(
content=(
"<html><body><h1>Too many sign-in attempts</h1>"
f"<p>Try again in about {retry_after} seconds</p></body></html>"
),
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
headers=exc.headers,
)
# Create UI token object
returned_ui_token_object: Final = create_ui_token_object(
@ -16009,6 +15998,7 @@ async def login_v2(request: Request):
password=password,
master_key=master_key,
prisma_client=prisma_client,
throttle=LoginThrottle.from_request(request, general_settings, redis_usage_cache),
general_settings=general_settings,
)
@ -16080,6 +16070,7 @@ async def login_v3(request: Request):
password=password,
master_key=master_key,
prisma_client=prisma_client,
throttle=LoginThrottle.from_request(request, general_settings, redis_usage_cache),
general_settings=general_settings,
)

View file

@ -586,6 +586,24 @@
],
"default_model_placeholder": "azure_ai/command-r-plus"
},
{
"provider": "Azure_Speech",
"provider_display_name": "Azure AI Speech",
"litellm_provider": "azure_speech",
"credential_fields": [
{
"key": "api_key",
"label": "Azure AI Speech Subscription Key",
"placeholder": null,
"tooltip": "The Ocp-Apim-Subscription-Key for your Azure AI Speech resource. The proxy injects it on every /azure_speech/* pass-through request. Region and API base come from AZURE_SPEECH_REGION / AZURE_SPEECH_API_BASE",
"required": true,
"field_type": "password",
"options": null,
"default_value": null
}
],
"default_model_placeholder": "azure_speech/short-audio"
},
{
"provider": "AZURE_TEXT",
"provider_display_name": "Azure Text",

View file

@ -2589,8 +2589,7 @@ class ManagedResponsesWebSocketHandler:
if "litellm_metadata" not in call_kwargs:
call_kwargs["litellm_metadata"] = {}
call_kwargs["litellm_metadata"]["proxy_server_request"] = proxy_server_request
call_kwargs.setdefault("litellm_params", {})
call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request
call_kwargs["proxy_server_request"] = proxy_server_request
async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> _MutableJsonObject | None:
"""

View file

@ -3940,12 +3940,24 @@ class Router:
)
_router_timeout: Final = (
float(self._explicit_timeout) if isinstance(self._explicit_timeout, (int, float)) else None
self.request_timeout
if self.request_timeout is not None
else float(self._explicit_timeout)
if isinstance(self._explicit_timeout, (int, float))
else None
)
_router_stream_timeout: Final = (
self.stream_timeout
if self.stream_timeout is not None
else self.request_timeout
if self.request_timeout is not None
else self.default_litellm_params.get("stream_timeout")
)
kwargs["timeout"] = resolve_llm_passthrough_timeout(
kwargs=kwargs,
litellm_params=deployment["litellm_params"],
router_timeout=_router_timeout,
router_stream_timeout=_router_stream_timeout,
)
else:
kwargs["timeout"] = self._get_timeout(kwargs=kwargs, data=deployment["litellm_params"])
@ -10575,11 +10587,12 @@ class Router:
2. If not, check if litellm model name is in model info
3. If not, return None
"""
from litellm.utils import _update_dictionary
from litellm.utils import _update_dictionary, cost_map_omits_token_price
model_info: ModelInfo | None = None
custom_model_info: dict | None = None
litellm_model_name_model_info: ModelInfo | None = None
base_model_key: str | None = None
try:
custom_model_info = (
@ -10606,6 +10619,7 @@ class Router:
## update litellm model info with base model info
base_model_info: Final = copy.deepcopy(litellm.get_model_info(model=base_model))
if base_model_info is not None:
base_model_key = base_model_info.get("key")
# Base model provides defaults, custom model info overrides
custom_model_info = _update_dictionary(
cast(dict, base_model_info),
@ -10633,6 +10647,15 @@ class Router:
# custom_model_info already includes base_model defaults at this point, if applicable
model_info = cast(ModelInfo, custom_model_info)
if model_info is None:
return None
builtin_key: Final = (
litellm_model_name_model_info.get("key") if litellm_model_name_model_info is not None else None
)
if cost_map_omits_token_price(model_id, builtin_key, base_model_key):
return cast( # cast-ok: TypedDict spread with overridden keys loses its type
ModelInfo, {**model_info, "input_cost_per_token": None, "output_cost_per_token": None}
)
return model_info
def _set_model_group_info(self, model_group: str, user_facing_model_group_name: str) -> ModelGroupInfo | None:

View file

@ -62,6 +62,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.singulr import (
from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
ToolPermissionGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
TypeSafeGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import (
VigilGuardGuardrailConfigModel,
)
@ -138,6 +141,7 @@ class SupportedGuardrailIntegrations(Enum):
SINGULR = "singulr"
HEADROOM = "headroom"
COMPRESR = "compresr"
TYPESAFE = "typesafe"
STRAIKER = "straiker"
ALICE = "alice"
AGENT_365 = "agent_365"
@ -1055,7 +1059,7 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
default="fail_closed",
description=(
"Behavior when a guardrail endpoint is unreachable due to network errors. "
"Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. "
"Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. "
"'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed."
),
)
@ -1171,6 +1175,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o
LakeraV2GuardrailConfigModel,
HeadroomGuardrailConfigModel,
CompresrGuardrailConfigModel,
TypeSafeGuardrailConfigModel,
RepelloAIGuardrailConfigModel,
LassoGuardrailConfigModel,
DeepKeepGuardrailConfigModel,

View file

@ -1160,6 +1160,10 @@ OpenAIImageGenerationOptionalParams = Literal[
"image_url",
"image_prompt_strength",
"aspect_ratio",
"width",
"height",
"guidance",
"steps",
"imageConfig",
]

View file

@ -464,3 +464,11 @@ class MCPGatewaySessionsResponse(BaseModel):
by_client: list[MCPGatewaySessionGroupCount] = Field(default_factory=list)
by_user: list[MCPGatewaySessionGroupCount] = Field(default_factory=list)
sessions: list[MCPGatewaySession] = Field(default_factory=list)
class MCPGatewaySessionsTerminateResponse(BaseModel):
"""Stateful sessions an administrator force-closed on this proxy worker."""
worker_pid: int
terminated_sessions: int
sessions: list[MCPGatewaySession] = Field(default_factory=list)

View file

@ -0,0 +1,63 @@
from typing import Literal
from pydantic import BaseModel, Field
from .base import GuardrailConfigModel
class TypeSafeGuardrailOptionalParams(BaseModel):
"""Optional tuning knobs for the TypeSafe (Jev) compaction guardrail."""
relevance_threshold: float | None = Field(
default=None,
ge=0.0,
le=1.0,
description=(
"Relevance cutoff in [0, 1]. A completed tool exchange is dropped when Jev "
"scores the probability that it is still needed below this value. Defaults to 0.2."
),
)
min_chars_to_evaluate: int | None = Field(
default=None,
ge=0,
description=(
"Skip tool exchanges whose combined tool-result text is shorter than this many characters. Defaults to 200."
),
)
max_result_chars_in_state: int | None = Field(
default=None,
ge=1,
description=(
"Tool result text is truncated to this many characters when sent to the Jev evaluator, "
"keeping the head and tail. Defaults to 4000."
),
)
class TypeSafeGuardrailConfigModel(GuardrailConfigModel[TypeSafeGuardrailOptionalParams]):
api_key: str | None = Field(
default=None,
description="TypeSafe API key, sent as a Bearer token. Falls back to the TYPESAFE_API_KEY env var.",
)
api_base: str | None = Field(
default=None,
description=(
"Base URL of the TypeSafe API. Falls back to the TYPESAFE_API_BASE env var, then https://api.typesafe.ai."
),
)
model: str | None = Field(
default=None,
description="TypeSafe evaluation model (not the LLM). Defaults to 'jev-latest'.",
)
unreachable_fallback: Literal["fail_closed", "fail_open"] = Field(
default="fail_open",
description=(
"Behavior when the TypeSafe evaluation service is unreachable or errors. "
"'fail_open' (default) forwards the request uncompacted. 'fail_closed' "
"raises an error instead."
),
)
@staticmethod
def ui_friendly_name() -> str:
return "TypeSafe (Jev) Compaction"

View file

@ -100,6 +100,16 @@ class DailySpendMetadata(BaseModel):
page: int = Field(default=1)
total_pages: int = Field(default=1)
has_more: bool = Field(default=False)
api_key_limit: int | None = Field(
default=None,
description="When set, api_keys and every api_key_breakdown list at most this many keys, "
"ranked by spend. Totals and the model, provider, mcp and endpoint rollups still cover every key.",
)
total_api_keys: int | None = Field(
default=None,
description="Distinct API keys matching the filters. When this exceeds api_key_limit, the per-key "
"lists are truncated to the highest-spend keys.",
)
class SpendAnalyticsPaginatedResponse(BaseModel):

View file

@ -61,7 +61,9 @@ class SCIMUserGroup(BaseModel):
class SCIMMultiValuedAttribute(BaseModel):
value: str
model_config = ConfigDict(extra="allow")
value: str | None = None
display: str | None = None
type: str | None = None
primary: bool | None = None

View file

@ -4074,6 +4074,7 @@ class LlmProviders(str, Enum):
TOPAZ = "topaz"
SAP_GENERATIVE_AI_HUB = "sap"
ASSEMBLYAI = "assemblyai"
AZURE_SPEECH = "azure_speech"
CHARITY_ENGINE = "charity_engine"
GITHUB_COPILOT = "github_copilot"
SNOWFLAKE = "snowflake"

View file

@ -3156,6 +3156,22 @@ def reapply_runtime_model_cost_registrations() -> None:
register_model(model_cost=dict(_runtime_registered_model_cost)) # mutable-ok: snapshot, replay rewrites it
def cost_map_omits_token_price(*keys: object) -> bool:
"""Whether the raw ``litellm.model_cost`` entries under ``keys`` exist but none carries a per-token price.
``get_model_info`` substitutes 0 for a missing price, which reads exactly like a declared
zero. Surfaces that report pricing use this to keep an unpriced deployment at ``None``.
"""
entries: Final = tuple(
entry
for entry in (litellm.model_cost.get(key) for key in keys if isinstance(key, str))
if isinstance(entry, dict)
)
return len(entries) > 0 and not any(
"input_cost_per_token" in entry or "output_cost_per_token" in entry for entry in entries
)
def register_model(
model_cost: str | dict,
*,

View file

@ -10820,6 +10820,25 @@
"/v1/images/generations"
]
},
"azure_ai/FLUX.2-flex": {
"input_cost_per_pixel": 5e-08,
"litellm_provider": "azure_ai",
"max_input_tokens": 32000,
"max_tokens": 32000,
"mode": "image_generation",
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/",
"supported_endpoints": [
"/v1/images/generations",
"/v1/images/edits"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"image"
]
},
"azure_ai/FW-DeepSeek-V3.2": {
"deprecation_date": "2027-07-01",
"cache_read_input_token_cost": 3.1e-07,
@ -16690,6 +16709,46 @@
"supports_tool_choice": true,
"supports_vision": true
},
"dashscope/qwen3.8-flash": {
"cache_creation_input_token_cost": 2e-07,
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "dashscope",
"max_input_tokens": 991808,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4.7e-07,
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true,
"supports_web_search": true
},
"dashscope/qwen3.8-omni-flash": {
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "dashscope",
"max_input_tokens": 991808,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4.7e-07,
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true,
"supports_web_search": true
},
"dashscope/qwq-plus": {
"input_cost_per_token": 8e-07,
"litellm_provider": "dashscope",
@ -18594,6 +18653,46 @@
"supports_tool_choice": true,
"supports_vision": true
},
"qwen_ai_platform/qwen3.8-flash": {
"cache_creation_input_token_cost": 2e-07,
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "qwen_ai_platform",
"max_input_tokens": 991808,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4.7e-07,
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true,
"supports_web_search": true
},
"qwen_ai_platform/qwen3.8-omni-flash": {
"cache_read_input_token_cost": 1.6e-08,
"input_cost_per_token": 1.5e-07,
"litellm_provider": "qwen_ai_platform",
"max_input_tokens": 991808,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 4.7e-07,
"source": "https://docs.modelstudio.console.alibabacloud.com/en/model-studio/model-pricing",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true,
"supports_web_search": true
},
"qwen_ai_platform/qwq-plus": {
"input_cost_per_token": 8e-07,
"litellm_provider": "qwen_ai_platform",
@ -22090,8 +22189,8 @@
"embed-english-light-v3.0": {
"input_cost_per_token": 1e-07,
"litellm_provider": "cohere",
"max_input_tokens": 1024,
"max_tokens": 1024,
"max_input_tokens": 512,
"max_tokens": 512,
"mode": "embedding",
"output_cost_per_token": 0.0
},
@ -22108,8 +22207,8 @@
"input_cost_per_image": 0.0001,
"input_cost_per_token": 1e-07,
"litellm_provider": "cohere",
"max_input_tokens": 1024,
"max_tokens": 1024,
"max_input_tokens": 512,
"max_tokens": 512,
"metadata": {
"notes": "'supports_image_input' is a deprecated field. Use 'supports_embedding_image_input' instead."
},
@ -22130,8 +22229,8 @@
"embed-multilingual-v3.0": {
"input_cost_per_token": 1e-07,
"litellm_provider": "cohere",
"max_input_tokens": 1024,
"max_tokens": 1024,
"max_input_tokens": 512,
"max_tokens": 512,
"mode": "embedding",
"output_cost_per_token": 0.0,
"supports_embedding_image_input": true
@ -22139,8 +22238,8 @@
"embed-multilingual-light-v3.0": {
"input_cost_per_token": 0.0001,
"litellm_provider": "cohere",
"max_input_tokens": 1024,
"max_tokens": 1024,
"max_input_tokens": 512,
"max_tokens": 512,
"mode": "embedding",
"output_cost_per_token": 0.0,
"supports_embedding_image_input": true
@ -57910,14 +58009,14 @@
"supports_tool_choice": true
},
"bedrock_mantle/openai.gpt-5.6-sol": {
"input_cost_per_token": 5.5e-06,
"input_cost_per_token_above_272k_tokens": 1.1e-05,
"cache_creation_input_token_cost": 6.875e-06,
"cache_creation_input_token_cost_above_272k_tokens": 1.375e-05,
"cache_read_input_token_cost": 5.5e-07,
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"input_cost_per_token": 4.4e-06,
"input_cost_per_token_above_272k_tokens": 8.8e-06,
"cache_creation_input_token_cost": 5.5e-06,
"cache_creation_input_token_cost_above_272k_tokens": 1.1e-05,
"cache_read_input_token_cost": 4.4e-07,
"cache_read_input_token_cost_above_272k_tokens": 8.8e-07,
"output_cost_per_token": 2.2e-05,
"output_cost_per_token_above_272k_tokens": 3.3e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
@ -65816,9 +65915,9 @@
"supports_web_search": false
},
"openrouter/z-ai/glm-5.3": {
"input_cost_per_token": 1.4e-06,
"output_cost_per_token": 4.4e-06,
"cache_read_input_token_cost": 2.6e-07,
"input_cost_per_token": 9.1e-07,
"output_cost_per_token": 2.86e-06,
"cache_read_input_token_cost": 1.69e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1310720,
"max_output_tokens": 943717,
@ -70539,14 +70638,14 @@
"supports_web_search": false
},
"openrouter/~deepseek/deepseek-flash-latest": {
"cache_read_input_token_cost": 1.5e-08,
"input_cost_per_token": 1.5e-07,
"cache_read_input_token_cost": 4.2e-09,
"input_cost_per_token": 1.4e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 393216,
"max_tokens": 393216,
"mode": "chat",
"output_cost_per_token": 6e-07,
"output_cost_per_token": 4.2e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -70831,14 +70930,14 @@
"supports_web_search": false
},
"openrouter/~z-ai/glm-latest": {
"cache_read_input_token_cost": 1.5e-07,
"cache_read_input_token_cost": 1.46625e-07,
"input_cost_per_token": 9e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1310720,
"max_output_tokens": 235929,
"max_tokens": 235929,
"mode": "chat",
"output_cost_per_token": 3e-06,
"output_cost_per_token": 2.805e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -73982,14 +74081,14 @@
"supports_web_search": false
},
"openrouter/tencent/hy3": {
"cache_read_input_token_cost": 3.3e-08,
"input_cost_per_token": 1.32e-07,
"cache_read_input_token_cost": 2.0625e-08,
"input_cost_per_token": 8.25e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 262144,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 5.28e-07,
"output_cost_per_token": 3.3e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,

View file

@ -86,7 +86,7 @@ locals {
"/queue/chat/*",
"/v1beta/*",
"/interactions/*",
"/anthropic/*", "/azure/*", "/azure_ai/*", "/aws/*", "/bedrock/*", "/comprehendmedical*", "/transcribe*",
"/anthropic/*", "/azure/*", "/azure_ai/*", "/azure_speech/*", "/aws/*", "/bedrock/*", "/comprehendmedical*", "/transcribe*",
"/cohere/*", "/gemini/*", "/google/*",
"/vertex_ai/*", "/vertex-ai/*",
"/assemblyai/*", "/eu.assemblyai/*",

View file

@ -55,7 +55,7 @@ locals {
"/queue/chat/*",
"/v1beta/*",
"/interactions/*",
"/anthropic/*", "/azure/*", "/azure_ai/*", "/aws/*", "/bedrock/*", "/comprehendmedical*", "/transcribe*",
"/anthropic/*", "/azure/*", "/azure_ai/*", "/azure_speech/*", "/aws/*", "/bedrock/*", "/comprehendmedical*", "/transcribe*",
"/cohere/*", "/gemini/*", "/google/*",
"/vertex_ai/*", "/vertex-ai/*",
"/assemblyai/*", "/eu.assemblyai/*",

View file

@ -51,9 +51,7 @@ try:
if general_settings_section:
# Extract the table rows, which contain the documented keys
table_content = general_settings_section.group(1)
doc_key_pattern = re.compile(
r"\|\s*([^\|]+?)\s*\|"
) # Capture the key from each row of the table
doc_key_pattern = re.compile(r"^\|\s*([^\|]+?)\s*\|", re.MULTILINE)
documented_keys.update(doc_key_pattern.findall(table_content))
except Exception as e:
raise Exception(

View file

@ -331,9 +331,7 @@ class TestMCPPerUserTokenCache:
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache):
await cache.delete("alice", "slack-test")
mock_dual_cache.async_delete_cache.assert_called_once_with(
"mcp:per_user_token:alice:slack-test"
)
mock_dual_cache.async_delete_cache.assert_called_once_with(key="mcp:per_user_token:alice:slack-test")
mock_dual_cache.async_set_cache.assert_not_called()
@pytest.mark.asyncio

View file

@ -1372,6 +1372,7 @@ async def test_create_team_member_add_team_admin(
from fastapi import Request
from litellm.proxy._types import (
LiteLLM_TeamMembership,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
Member,
@ -1454,6 +1455,10 @@ async def test_create_team_member_add_team_admin(
team_mock_client.update = AsyncMock(
return_value=LiteLLM_TeamTableCachedObj(team_id="1234")
)
membership_mock_client = AsyncMock()
membership_mock_client.upsert = AsyncMock(
return_value=LiteLLM_TeamMembership(user_id="1234", team_id=_team_id)
)
tx_cm = _member_add_tx_cm(team_mock_client)
@ -1463,6 +1468,11 @@ async def test_create_team_member_add_team_admin(
"litellm_teamtable",
team_mock_client,
),
patch.object( # test-quality-ok: legacy test swaps the prisma table on the module-level client
litellm.proxy.proxy_server.prisma_client.db,
"litellm_teammembership",
membership_mock_client,
),
patch.object(
litellm.proxy.proxy_server.prisma_client,
"tx",

View file

@ -0,0 +1,20 @@
from litellm.llms.azure.vector_stores.transformation import AzureOpenAIVectorStoreConfig
def test_transform_search_vector_store_request_preserves_azure_query_string():
config = AzureOpenAIVectorStoreConfig()
api_base = config.get_complete_url(
api_base="https://x.openai.azure.com",
litellm_params={"api_version": "2024-10-21"},
)
url, _ = config.transform_search_vector_store_request(
vector_store_id="vs_1",
query="hello",
vector_store_search_optional_params={},
api_base=api_base,
litellm_logging_obj=None,
litellm_params={"api_version": "2024-10-21"},
)
assert url == "https://x.openai.azure.com/openai/vector_stores/vs_1/search?api-version=2024-10-21"

View file

@ -1,12 +1,20 @@
import base64
import json
from collections.abc import Mapping
from typing import Final
import httpx
import pytest
import litellm
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.azure_ai.image_edit.flux2_transformation import (
AzureFoundryFlux2ImageEditConfig,
)
from litellm.llms.azure_ai.image_edit.transformation import (
AzureFoundryFluxImageEditConfig,
)
from litellm.llms.custom_httpx.http_handler import HTTPHandler
def test_azure_ai_validate_environment():
@ -60,3 +68,122 @@ def test_flux2_validate_environment_with_entra_token(monkeypatch):
assert headers["Authorization"] == "Bearer entra-token"
assert headers["Content-Type"] == "application/json"
def test_flux2_image_edit_maps_openai_and_provider_parameters():
config = AzureFoundryFlux2ImageEditConfig()
requested_params = ImageEditRequestUtils.get_requested_image_edit_optional_param(
{
"n": 2,
"size": "1536x1024",
"guidance": 4.5,
"steps": 32,
"unrelated": "discarded",
},
provider_supported_params=config.get_supported_openai_params("FLUX.2-flex"),
)
mapped_params = config.map_openai_params(
image_edit_optional_params=requested_params,
model="FLUX.2-flex",
drop_params=False,
)
assert mapped_params == {
"num_images": 2,
"width": 1536,
"height": 1024,
"guidance": 4.5,
"steps": 32,
}
@pytest.mark.parametrize(
("model", "max_reference_images"),
[
("FLUX.2-flex", 10),
("FLUX.2-pro", 8),
],
)
def test_flux2_image_edit_uses_all_reference_fields(model: str, max_reference_images: int):
images = [f"image-{index}".encode() for index in range(1, max_reference_images + 1)]
request, files = AzureFoundryFlux2ImageEditConfig().transform_image_edit_request(
model=model,
prompt="Blend every reference",
image=images,
image_edit_optional_request_params={"guidance": 4.5, "steps": 20},
litellm_params={},
headers={},
)
assert files == []
assert request["input_image"] == base64.b64encode(images[0]).decode()
assert request[f"input_image_{max_reference_images}"] == base64.b64encode(images[-1]).decode()
assert "input_image_1" not in request
assert "image" not in request
assert len([key for key in request if key.startswith("input_image")]) == max_reference_images
assert request["guidance"] == 4.5
assert request["steps"] == 20
@pytest.mark.parametrize(
("model", "reference_images"),
[
("FLUX.2-flex", 11),
("FLUX.2-pro", 9),
],
)
def test_flux2_image_edit_rejects_too_many_references(model: str, reference_images: int):
with pytest.raises(ValueError, match=f"at most {reference_images - 1} reference images"):
AzureFoundryFlux2ImageEditConfig().transform_image_edit_request(
model=model,
prompt="Blend every reference",
image=[b"image"] * reference_images,
image_edit_optional_request_params={},
litellm_params={},
headers={},
)
@pytest.mark.parametrize("dimensions", ({"size": "2048x1024"}, {"width": 2048, "height": 1024}, {"width": "2048", "height": "1024"}))
@pytest.mark.usefixtures("local_model_cost_map")
def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[str, int | str]):
def respond(request: httpx.Request) -> httpx.Response:
body: Final = json.loads(request.content)
assert body == {
"model": "FLUX.2-flex",
"prompt": "Add a hat",
"input_image": base64.b64encode(b"image").decode(),
"num_images": 2,
"width": 2048,
"height": 1024,
"guidance": 4.5,
"steps": 32,
}
return httpx.Response(200, json={"data": [{"b64_json": "aW1n"}, {"b64_json": "aW1n"}]})
client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond)))
response: Final = litellm.image_edit(
model="azure_ai/FLUX.2-flex",
image=b"image",
prompt="Add a hat",
api_key="test-key",
api_base="https://example.services.ai.azure.com",
client=client,
n=2,
guidance="4.5",
steps="32",
**dimensions,
)
assert response._hidden_params["response_cost"] == pytest.approx(5e-08 * 2048 * 1024 * 2)
def test_flux2_image_edit_accepts_and_drops_openai_only_parameters():
optional_params: Final = ImageEditRequestUtils.get_optional_params_image_edit(
model="FLUX.2-pro",
image_edit_provider_config=AzureFoundryFlux2ImageEditConfig(),
image_edit_optional_params={"n": 1, "size": "auto", "quality": "high", "user": "end-user-1"},
drop_params=False,
)
assert optional_params == {"num_images": 1}

View file

@ -0,0 +1,223 @@
from collections.abc import Mapping
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
import litellm
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
from litellm.llms.azure.azure import AzureChatCompletion
from litellm.llms.azure.image_generation import get_azure_image_generation_config
from litellm.llms.azure.image_generation.http_utils import azure_deployment_image_generation_json_body
from litellm.llms.azure_ai.image_generation.flux_transformation import (
AzureFoundryFluxImageGenerationConfig,
)
from litellm.types.utils import ImageObject, ImageResponse
from litellm.utils import _invalidate_model_cost_lowercase_map, get_optional_params_image_gen
@pytest.fixture(autouse=True)
def use_local_model_cost_map(monkeypatch):
monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map())
litellm.get_model_info.cache_clear()
_invalidate_model_cost_lowercase_map()
yield
litellm.get_model_info.cache_clear()
_invalidate_model_cost_lowercase_map()
@pytest.mark.parametrize(
("model", "provider_path"),
[
("FLUX.2-flex", "flux-2-flex"),
("FLUX.2-pro", "flux-2-pro"),
],
)
def test_flux2_uses_model_specific_provider_url(model: str, provider_path: str):
url = AzureChatCompletion().create_azure_base_url(
azure_client_params={
"azure_endpoint": "https://example.services.ai.azure.com/",
"api_version": "preview",
},
model=model,
)
assert (
url == f"https://example.services.ai.azure.com/providers/blackforestlabs/v1/{provider_path}?api-version=preview"
)
def test_flux2_flex_maps_openai_and_provider_parameters():
config = AzureFoundryFluxImageGenerationConfig()
mapped_params = config.map_openai_params(
non_default_params={
"n": 2,
"size": "1536x1024",
"guidance": 4.5,
"steps": 32,
"output_format": "jpeg",
},
optional_params={},
model="FLUX.2-flex",
drop_params=False,
)
url = config.get_flux2_image_generation_url(
api_base="https://example.services.ai.azure.com",
model="FLUX.2-flex",
api_version="preview",
)
request = azure_deployment_image_generation_json_body(
api_base=url,
data={"model": "FLUX.2-flex", "prompt": "A red fox", **mapped_params},
deployment_name="FLUX.2-flex",
)
assert request == {
"model": "FLUX.2-flex",
"prompt": "A red fox",
"num_images": 2,
"width": 1536,
"height": 1024,
"guidance": 4.5,
"steps": 32,
"output_format": "jpeg",
}
def test_flux2_flex_rejects_invalid_size_as_bad_request():
with pytest.raises(litellm.BadRequestError, match="Expected 'WxH'") as raised:
get_optional_params_image_gen(
model="FLUX.2-flex",
custom_llm_provider="azure_ai",
provider_config=AzureFoundryFluxImageGenerationConfig(),
size="large",
)
assert raised.value.status_code == 400
@pytest.mark.parametrize("model", ("FLUX.2-pro", "FLUX.2-flex"))
def test_flux2_accepts_and_drops_openai_only_image_parameters(model: str):
optional_params: Final = get_optional_params_image_gen(
model=model,
custom_llm_provider="azure_ai",
provider_config=AzureFoundryFluxImageGenerationConfig(),
n=1,
size="auto",
quality="high",
user="end-user-1",
background="transparent",
moderation="low",
output_compression=50,
)
assert optional_params == {"num_images": 1}
def test_flux2_flex_model_info():
model_info = litellm.get_model_info(
model="FLUX.2-flex",
custom_llm_provider="azure_ai",
)
catalog_info = litellm.model_cost["azure_ai/FLUX.2-flex"]
assert model_info["mode"] == "image_generation"
assert model_info["max_input_tokens"] == 32000
assert model_info["max_tokens"] == 32000
assert model_info["supported_endpoints"] == ["/v1/images/generations", "/v1/images/edits"]
assert catalog_info["input_cost_per_pixel"] == 5e-08
assert catalog_info["supported_modalities"] == ["text", "image"]
assert catalog_info["supported_output_modalities"] == ["image"]
def test_flux2_flex_cost_uses_generated_megapixels():
response = ImageResponse(
data=[
ImageObject(url="https://example.com/one.png"),
ImageObject(url="https://example.com/two.png"),
]
)
cost = CostCalculatorUtils.route_image_generation_cost_calculator(
model="FLUX.2-flex",
completion_response=response,
custom_llm_provider="azure_ai",
size="2048x1024",
call_type="image_generation",
)
assert cost == pytest.approx(5e-08 * 2048 * 1024 * 2)
@pytest.mark.parametrize("model", ("FLUX-1.1-pro", "FLUX.1-Kontext-pro"))
def test_flux1_preserves_existing_openai_parameters(model: str):
params: Final = {"n": 2, "size": "1536x1024", "quality": "high", "user": "test-user"}
mapped: Final = AzureFoundryFluxImageGenerationConfig().map_openai_params(
non_default_params=params,
optional_params={},
model=model,
drop_params=False,
)
assert mapped == params
@pytest.mark.parametrize("dimensions", ({"size": "2048x1024"}, {"width": 2048, "height": 1024}))
def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensions: Mapping[str, int | str]):
params: Final = AzureFoundryFluxImageGenerationConfig().map_openai_params(
non_default_params={"n": 2, **dimensions},
optional_params={},
model="FLUX.2-flex",
drop_params=False,
)
response: Final = get_azure_image_generation_config("FLUX.2-flex").transform_image_generation_response(
model="FLUX.2-flex",
raw_response=httpx.Response(200, json={"data": [{"b64_json": "aW1n"}, {"b64_json": "aW1n"}]}),
model_response=ImageResponse(),
logging_obj=MagicMock(),
request_data={"prompt": "A red fox", **params},
optional_params=params,
litellm_params={},
encoding=None,
)
assert litellm.completion_cost(
model="azure_ai/FLUX.2-flex",
completion_response=response,
optional_params=params,
call_type="image_generation",
) == pytest.approx(5e-08 * 2048 * 1024 * 2)
def test_flux2_flex_cost_accepts_lowercase_model_spelling():
response: Final = ImageResponse(data=[ImageObject(b64_json="aW1n"), ImageObject(b64_json="aW1n")])
cost: Final = litellm.completion_cost(
model="azure_ai/flux.2-flex",
completion_response=response,
optional_params={"width": 1536, "height": 1024, "num_images": 2},
call_type="image_generation",
)
assert cost == pytest.approx(5e-08 * 1536 * 1024 * 2)
def test_flux2_response_preserves_mapped_dimensions():
config = AzureFoundryFluxImageGenerationConfig()
params = config.map_openai_params(
non_default_params={"size": "2048x1024"}, optional_params={}, model="FLUX.2-flex", drop_params=False
)
response = config.transform_image_generation_response(
model="FLUX.2-flex",
raw_response=httpx.Response(200, json={"data": [{"b64_json": "aW1n"}]}),
model_response=ImageResponse(),
logging_obj=MagicMock(),
request_data={"prompt": "A landscape"},
optional_params=params,
litellm_params={},
encoding=None,
)
assert response.size == "2048x1024"

View file

@ -1839,6 +1839,27 @@ class TestBedrockMantleResponsesSigV4:
class TestBedrockMantleResponsesPricing:
@pytest.mark.parametrize(
"model",
["openai.gpt-5.6-sol", "openai.gpt-5.6-terra", "openai.gpt-5.6-luna"],
)
def test_mantle_matches_in_region_converse_pricing(self, local_cost_map, model):
"""bedrock-mantle serves these models In-Region only, and the AWS model
cards price In-Region and Geo CRIS identically -- so every cost field on
the mantle key must equal the `us.` converse key. A price change applied
to one namespace but not the other shows up here.
"""
mantle = litellm.model_cost[f"bedrock_mantle/{model}"]
converse = litellm.model_cost[f"us.{model}"]
cost_fields = [k for k in converse if "cost" in k and k != "search_context_cost_per_query"]
assert cost_fields, "expected cost fields on the converse entry"
for field in cost_fields:
assert mantle.get(field) == pytest.approx(converse[field]), (
f"{model}: {field} is {mantle.get(field)} on bedrock_mantle "
f"but {converse[field]} on us. (bedrock_converse)"
)
def test_models_registered(self, local_cost_map):
assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models
assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models

View file

@ -7,11 +7,8 @@ from litellm.types.vector_stores import (
class TestOpenAIVectorStoreAPIConfig:
@pytest.mark.parametrize("metadata", [{}, None])
def test_transform_create_vector_store_request_with_metadata_empty_or_none(
self, metadata
):
def test_transform_create_vector_store_request_with_metadata_empty_or_none(self, metadata):
"""
Test transform_create_vector_store_request when metadata is None or empty dict.
"""
@ -24,9 +21,7 @@ class TestOpenAIVectorStoreAPIConfig:
"metadata": metadata,
}
url, request_body = config.transform_create_vector_store_request(
vector_store_create_params, api_base
)
url, request_body = config.transform_create_vector_store_request(vector_store_create_params, api_base)
assert url == api_base
assert request_body["name"] == "test-vector-store"
@ -50,9 +45,7 @@ class TestOpenAIVectorStoreAPIConfig:
"metadata": large_metadata,
}
url, request_body = config.transform_create_vector_store_request(
vector_store_create_params, api_base
)
url, request_body = config.transform_create_vector_store_request(vector_store_create_params, api_base)
assert url == api_base
assert request_body["name"] == "test-vector-store"
@ -77,8 +70,19 @@ class TestOpenAIVectorStoreAPIConfig:
litellm_params={},
)
assert (
url
== "https://api.openai.com/v1/vector_stores/..%2F..%2Ffiles%3Fx%3D1%23frag/search"
)
assert url == "https://api.openai.com/v1/vector_stores/..%2F..%2Ffiles%3Fx%3D1%23frag/search"
assert request_body["query"] == "hello"
def test_transform_search_vector_store_request_preserves_query_string(self):
config = OpenAIVectorStoreConfig()
url, _ = config.transform_search_vector_store_request(
vector_store_id="vs_1",
query="hello",
vector_store_search_optional_params={},
api_base="https://x.openai.azure.com/openai/vector_stores?api-version=2024-10-21",
litellm_logging_obj=None,
litellm_params={},
)
assert url == "https://x.openai.azure.com/openai/vector_stores/vs_1/search?api-version=2024-10-21"

View file

@ -0,0 +1,57 @@
import json
import pytest
from litellm.proxy._experimental.mcp_server.byok_credential_cache import (
CachedByokCredential,
byok_credential_cache,
byok_credential_cache_key,
cache_byok_credential,
get_cached_byok_credential,
)
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
class _FakeRedisCache:
namespace = None
def init_async_client(self) -> object:
return object()
@pytest.fixture(autouse=True)
def _empty_cache():
byok_credential_cache.flush_cache()
yield
byok_credential_cache.flush_cache()
def test_a_cached_negative_lookup_is_distinguishable_from_a_miss():
assert get_cached_byok_credential("u-1", "srv-1") is None
cache_byok_credential("u-1", "srv-1", None)
assert get_cached_byok_credential("u-1", "srv-1") == CachedByokCredential(credential=None)
cache_byok_credential("u-1", "srv-1", "sk-stored")
assert get_cached_byok_credential("u-1", "srv-1") == CachedByokCredential(credential="sk-stored")
assert get_cached_byok_credential("u-1", "srv-2") is None
def test_peer_worker_invalidation_message_evicts_the_cached_credential():
"""The key a mutating worker broadcasts must be the key every other worker caches under."""
cache_byok_credential("mallory", "srv-byok", "sk-revoked")
cache_byok_credential("alice", "srv-byok", "sk-kept")
subscriber = AuthCacheInvalidationSubscriber(
redis_cache=_FakeRedisCache(), # pyright: ignore[reportArgumentType] # subscriber is never started; only its message handler runs
user_api_key_cache=UserApiKeyCache(),
additional_in_memory_caches=(byok_credential_cache,),
)
subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler
{
"type": "message",
"data": json.dumps({"cache_key": byok_credential_cache_key("mallory", "srv-byok")}).encode(),
}
)
assert get_cached_byok_credential("mallory", "srv-byok") is None
assert get_cached_byok_credential("alice", "srv-byok") == CachedByokCredential(credential="sk-kept")

View file

@ -592,7 +592,7 @@ async def test_check_byok_credential_missing_credential(monkeypatch):
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
monkeypatch.setattr(server_module, "_byok_cred_cache", {})
server_module.byok_credential_cache.flush_cache()
mock_prisma = MagicMock()
with (
@ -628,7 +628,7 @@ async def test_execute_byok_tool_missing_credential_advertises_api_key_flow(monk
from litellm.types.mcp_server.mcp_server_manager import MCPServer
monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy")
monkeypatch.setattr(mcp_module, "_byok_cred_cache", {})
mcp_module.byok_credential_cache.flush_cache()
server = MCPServer(server_id="byok-discovery", name="byok-discovery", transport=MCPTransport.http, is_byok=True)
prisma = MagicMock()
prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None)
@ -677,6 +677,40 @@ async def test_check_byok_credential_has_credential():
await _check_byok_credential(server, user_auth)
@pytest.mark.asyncio
async def test_invalidate_byok_cred_cache_evicts_locally_and_broadcasts_the_same_key():
"""A revoked credential must stop being served here and on every peer worker within the TTL."""
from litellm.proxy._experimental.mcp_server import server as server_module
from litellm.proxy._experimental.mcp_server.byok_credential_cache import byok_credential_cache_key
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(server_id="byok-revoke", name="byok-server", transport=MCPTransport.http, is_byok=True)
user_auth = UserAPIKeyAuth(user_id="mallory", api_key="sk-test")
server_module.byok_credential_cache.flush_cache()
db_lookup = AsyncMock(side_effect=["sk-before-revoke", None])
publish = AsyncMock()
with (
patch( # test-quality-ok: the DB row lookup is the only seam below the credential resolver; no Prisma fake exists
"litellm.proxy._experimental.mcp_server.db.get_user_credential", new=db_lookup
),
patch( # test-quality-ok: the resolver reads the module-level prisma_client singleton; the suite's only seam
"litellm.proxy.proxy_server.prisma_client", MagicMock()
),
patch.object( # test-quality-ok: the redis publisher is module-level; asserting the broadcast without a redis
server_module, "publish_auth_cache_invalidation", new=publish
),
):
assert await server_module._get_byok_credential(server, user_auth) == "sk-before-revoke"
assert await server_module._get_byok_credential(server, user_auth) == "sk-before-revoke"
await server_module._invalidate_byok_cred_cache("mallory", "byok-revoke")
assert await server_module._get_byok_credential(server, user_auth) is None
assert db_lookup.await_count == 2
publish.assert_awaited_once_with(cache_key=byok_credential_cache_key("mallory", "byok-revoke"))
@pytest.mark.asyncio
async def test_check_byok_credential_db_unavailable_fails_closed():
"""BYOK server with no prisma_client → 503, not silent pass.

View file

@ -212,6 +212,53 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_each_user():
assert set(invalidations) == {("alice", "srv-1"), ("bob", "srv-1")}
@pytest.mark.asyncio
async def test_list_server_user_credentials_types_each_row_without_leaking_the_secret():
"""The admin view of one server's stored credentials names the user and the kind of
credential (OAuth2 vs BYOK) and echoes OAuth expiry, but never the token or key itself."""
from litellm.proxy._experimental.mcp_server.db import list_server_user_credentials
oauth_row = _legacy_row(
json.dumps(
{
"type": "oauth2",
"access_token": "tok-alice",
"expires_at": "2026-12-31T00:00:00+00:00",
"connected_at": "2026-01-01T00:00:00+00:00",
}
)
)
oauth_row.user_id = "alice"
oauth_row.updated_at = datetime(2026, 1, 1, tzinfo=timezone.utc)
byok_row = _byok_row("carol")
byok_row.updated_at = datetime(2026, 2, 1, tzinfo=timezone.utc)
prisma = MagicMock()
prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[oauth_row, byok_row])
items = await list_server_user_credentials(prisma, "srv-1")
prisma.db.litellm_mcpusercredentials.find_many.assert_awaited_once_with(where={"server_id": "srv-1"})
assert [item.model_dump() for item in items] == [
{
"user_id": "alice",
"credential_type": "oauth2",
"expires_at": "2026-12-31T00:00:00+00:00",
"connected_at": "2026-01-01T00:00:00+00:00",
"updated_at": "2026-01-01T00:00:00+00:00",
},
{
"user_id": "carol",
"credential_type": "byok",
"expires_at": None,
"connected_at": None,
"updated_at": "2026-02-01T00:00:00+00:00",
},
]
serialized = "".join(item.model_dump_json() for item in items)
assert "tok-alice" not in serialized
assert "sk-byok-carol" not in serialized
@pytest.mark.asyncio
async def test_purge_user_oauth_credentials_for_server_spares_byok_rows():
"""Regression: the purge used to delete_many on server_id alone, wiping BYOK API keys that share

View file

@ -2871,6 +2871,255 @@ def test_remove_stateful_session_tracking_drops_client_info():
assert session_id not in mcp_server._stateful_session_client_info
def _admin_terminate_fixture(mcp_server):
def auth_user(user_id: str):
return mcp_server.MCPAuthenticatedUser(
user_api_key_auth=UserAPIKeyAuth(api_key=f"key-{user_id}", user_id=user_id),
)
contexts = {
"alice-session-1": auth_user("alice"),
"alice-session-2": auth_user("alice"),
"bob-session-1": auth_user("bob"),
"anon-session-1": mcp_server.MCPAuthenticatedUser(user_api_key_auth=None),
"gone-session-1": auth_user("alice"),
}
transports = {
session_id: MagicMock(terminate=AsyncMock())
for session_id in ("alice-session-1", "alice-session-2", "bob-session-1", "anon-session-1")
}
return contexts, transports
@pytest.mark.asyncio
async def test_terminate_mcp_gateway_sessions_by_user_closes_every_live_session_of_that_user():
try:
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy._experimental.mcp_server.server import session_manager_stateful
except ImportError:
pytest.skip("MCP server not available")
contexts, transports = _admin_terminate_fixture(mcp_server)
live_transports = dict(transports)
last_seen = {session_id: 100.0 for session_id in contexts}
locks = {session_id: asyncio.Lock() for session_id in contexts}
with (
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
session_manager_stateful, "_server_instances", live_transports
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_contexts, contexts, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_context_last_seen, last_seen, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_locks, locks, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_owners, {session_id: "owner" for session_id in contexts}, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_active_request_counts, {}, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_client_info, {}, clear=True
),
):
result = await mcp_server.terminate_mcp_gateway_sessions(user_id="alice")
assert set(live_transports) == {"bob-session-1", "anon-session-1"}
assert set(mcp_server._stateful_session_auth_contexts) == {"bob-session-1", "anon-session-1", "gone-session-1"}
assert set(mcp_server._stateful_session_locks) == {"bob-session-1", "anon-session-1", "gone-session-1"}
assert set(mcp_server._stateful_session_owners) == {"bob-session-1", "anon-session-1", "gone-session-1"}
assert set(mcp_server._stateful_session_auth_context_last_seen) == {
"bob-session-1",
"anon-session-1",
"gone-session-1",
}
transports["alice-session-1"].terminate.assert_awaited_once()
transports["alice-session-2"].terminate.assert_awaited_once()
transports["bob-session-1"].terminate.assert_not_awaited()
transports["anon-session-1"].terminate.assert_not_awaited()
assert result.terminated_sessions == 2
assert sorted(session.session_id_prefix for session in result.sessions) == ["alice-se", "alice-se"]
assert {session.user_id for session in result.sessions} == {"alice"}
assert "key-alice" not in result.model_dump_json()
assert "alice-session-1" not in result.model_dump_json()
@pytest.mark.asyncio
async def test_terminate_mcp_gateway_sessions_prefix_and_user_must_both_match():
try:
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy._experimental.mcp_server.server import session_manager_stateful
except ImportError:
pytest.skip("MCP server not available")
contexts, transports = _admin_terminate_fixture(mcp_server)
live_transports = dict(transports)
with (
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
session_manager_stateful, "_server_instances", live_transports
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_contexts, contexts, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_client_info, {}, clear=True
),
):
mismatch = await mcp_server.terminate_mcp_gateway_sessions(session_id_prefix="alice-session-1", user_id="bob")
assert mismatch.terminated_sessions == 0
assert set(live_transports) == set(transports)
stale = await mcp_server.terminate_mcp_gateway_sessions(session_id_prefix="gone-session-1")
assert stale.terminated_sessions == 0
exact = await mcp_server.terminate_mcp_gateway_sessions(session_id_prefix="alice-session-1", user_id="alice")
assert exact.terminated_sessions == 1
assert set(live_transports) == {"alice-session-2", "bob-session-1", "anon-session-1"}
@pytest.mark.asyncio
async def test_admin_terminated_session_id_gets_404_instead_of_a_fresh_stateless_session():
"""Once an admin closes a session, a client replaying its id must not be silently upgraded to a
new stateless session by the stale-header path; it gets 404 and has to initialize again."""
try:
from starlette.types import Scope
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy._experimental.mcp_server.server import session_manager_stateful
except ImportError:
pytest.skip("MCP server not available")
session_id = "admin-closed-session-1"
live_transports = {session_id: MagicMock(terminate=AsyncMock())}
contexts = {
session_id: mcp_server.MCPAuthenticatedUser(
user_api_key_auth=UserAPIKeyAuth(api_key="key-alice", user_id="alice"),
)
}
def scope_with_session_header() -> Scope:
return {
"type": "http",
"method": "POST",
"headers": [(b"content-type", b"application/json"), (b"mcp-session-id", session_id.encode())],
}
try:
with (
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
session_manager_stateful, "_server_instances", live_transports
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_contexts, contexts, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_client_info, {}, clear=True
),
):
await mcp_server.terminate_mcp_gateway_sessions(session_id_prefix=session_id)
terminated_scope = scope_with_session_header()
send = AsyncMock()
handled = await mcp_server._handle_stale_mcp_session(
terminated_scope, AsyncMock(), send, session_manager_stateful
)
assert handled is True
statuses = [m["status"] for (m,), _ in send.await_args_list if m["type"] == "http.response.start"]
assert statuses == [404]
assert [k for k, _ in terminated_scope["headers"]] == [b"content-type", b"mcp-session-id"]
unknown_scope = scope_with_session_header()
unknown_scope["headers"][1] = (b"mcp-session-id", b"never-seen-session")
assert (
await mcp_server._handle_stale_mcp_session(
unknown_scope, AsyncMock(), AsyncMock(), session_manager_stateful
)
is False
)
assert [k for k, _ in unknown_scope["headers"]] == [b"content-type"]
finally:
mcp_server._admin_terminated_session_ids.clear()
@pytest.mark.asyncio
async def test_admin_terminated_session_id_stays_refused_while_replayed_and_is_forgotten_like_an_idle_session():
"""The refusal window slides on every replay, so a client that keeps retrying is never silently
upgraded to a stateless session no matter how many other sessions an admin closes later; an id
nobody has replayed for a full idle timeout is dropped from the table by the idle sweep."""
try:
from starlette.types import Scope
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy._experimental.mcp_server.server import session_manager_stateful
except ImportError:
pytest.skip("MCP server not available")
idle_timeout = mcp_server._STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS
retrying_id, silent_id = "admin-closed-retrying", "admin-closed-silent"
contexts = {
session_id: mcp_server.MCPAuthenticatedUser(
user_api_key_auth=UserAPIKeyAuth(api_key="key-alice", user_id="alice"),
)
for session_id in (retrying_id, silent_id)
}
live_transports = {session_id: MagicMock(terminate=AsyncMock()) for session_id in contexts}
async def replay(session_id: str, now: float) -> tuple[bool, list[bytes]]:
scope: Scope = {
"type": "http",
"method": "POST",
"headers": [(b"content-type", b"application/json"), (b"mcp-session-id", session_id.encode())],
}
with patch.object( # test-quality-ok: the stale-session handler reads the clock directly; no injectable now
mcp_server.time, "monotonic", return_value=now
):
handled = await mcp_server._handle_stale_mcp_session(
scope, AsyncMock(), AsyncMock(), session_manager_stateful
)
return handled, [k for k, _ in scope["headers"]]
try:
with (
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
session_manager_stateful, "_server_instances", live_transports
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_contexts, contexts, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_client_info, {}, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_context_last_seen, {}, clear=True
),
):
with patch.object( # test-quality-ok: termination stamps the tombstone from the clock directly; no injectable now
mcp_server.time, "monotonic", return_value=1000.0
):
closed = await mcp_server.terminate_mcp_gateway_sessions(user_id="alice")
assert closed.terminated_sessions == 2
for elapsed in (idle_timeout - 1, 2 * idle_timeout - 2, 3 * idle_timeout - 3):
assert await replay(retrying_id, 1000.0 + elapsed) == (True, [b"content-type", b"mcp-session-id"])
await mcp_server._purge_expired_stateful_session_auth_contexts(now=1000.0 + idle_timeout)
assert set(mcp_server._admin_terminated_session_ids) == {retrying_id}
assert await replay(silent_id, 1000.0 + idle_timeout) == (False, [b"content-type"])
assert await replay(retrying_id, 1000.0 + 4 * idle_timeout) == (False, [b"content-type"])
assert mcp_server._admin_terminated_session_ids == {}
finally:
mcp_server._admin_terminated_session_ids.clear()
@pytest.mark.asyncio
async def test_initialize_request_with_existing_session_tracks_new_session():
try:

View file

@ -395,6 +395,34 @@ async def test_invalidate_clears_every_identity_for_a_server():
assert mock_client.post.call_count == 3
@pytest.mark.asyncio
async def test_per_user_token_delete_evicts_locally_and_broadcasts_to_peer_workers():
"""Revoking a user's OAuth token must not leave peer workers serving it from their in-memory layer."""
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import MCPPerUserTokenCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
local_cache = UserApiKeyCache()
publish = AsyncMock()
token_cache = MCPPerUserTokenCache()
key = token_cache._cache_key("mallory", "srv-oauth") # pyright: ignore[reportPrivateUsage] # asserting the broadcast names the stored key
local_cache.in_memory_cache.set_cache(key, "encrypted-token")
with (
patch.object( # test-quality-ok: the token cache reads the module-level user_api_key_cache singleton; the suite's only seam
proxy_server, "user_api_key_cache", local_cache
),
patch( # test-quality-ok: the redis publisher is module-level; asserting the broadcast without a redis
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new=publish,
),
):
await token_cache.delete("mallory", "srv-oauth")
assert local_cache.in_memory_cache.get_cache(key) is None
publish.assert_awaited_once_with(cache_key=key)
@pytest.mark.asyncio
async def test_m2m_mint_uses_admin_entered_token_url_when_issuer_yield_empties_resolved():
"""A pinned issuer empties the resolved token_url while configured_token_url keeps the

File diff suppressed because it is too large Load diff

View file

@ -57,6 +57,21 @@ def test_xff_honored_from_trusted_peer():
assert via_proxy is True
def test_ipv4_mapped_peer_and_hop_match_ipv4_trusted_ranges():
request = make_request(headers={"x-forwarded-for": "203.0.113.9, ::ffff:10.0.0.5"}, client=("::ffff:10.0.0.1", 1))
ip, via_proxy = resolve_client_ip(request, TRUSTED)
assert ip == "203.0.113.9"
assert via_proxy is True
def test_ipv4_mapped_peer_still_matches_mapped_notation_trusted_range():
config = TrustedProxyConfig(use_forwarded_for=True, trusted_proxy_cidrs=["::ffff:10.0.0.0/104"])
request = make_request(headers={"x-forwarded-for": "203.0.113.9"}, client=("::ffff:10.0.0.1", 1))
ip, via_proxy = resolve_client_ip(request, config)
assert ip == "203.0.113.9"
assert via_proxy is True
def test_spoofed_xff_from_untrusted_peer_is_ignored():
request = make_request(
headers={"x-forwarded-for": "203.0.113.9"}, client=("8.8.8.8", 1)

View file

@ -3967,3 +3967,74 @@ def test_auto_router_session_read_grant_rejects_other_methods_paths_and_scopes(
RouteChecks.should_call_route(route, valid_token, request)
assert error.value.status_code == 403
@pytest.mark.parametrize("route", ["/key/generate", "/key/update"])
def test_team_service_account_key_allowed_key_management_routes(route):
"""A service account key (user_id=None, team_id set, metadata.service_account_id)
can reach key-management routes; team scoping is enforced in the handlers."""
valid_token = UserAPIKeyAuth(
api_key="sk",
team_id="t1",
user_id=None,
metadata={"service_account_id": "ci"},
)
request = MagicMock(spec=Request)
request.query_params = {}
result = RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=None,
_user_role=None,
route=route,
request=request,
valid_token=valid_token,
request_data={},
)
assert result is None
@pytest.mark.parametrize("route", ["/team/new", "/spend/logs", "/key/delete", "/key/regenerate"])
def test_team_service_account_key_rejected_outside_generate_and_update(route):
"""The service account carve-out covers only /key/generate and /key/update; other
key-management routes lack team scoping for a userless caller and stay denied."""
valid_token = UserAPIKeyAuth(
api_key="sk",
team_id="t1",
user_id=None,
metadata={"service_account_id": "ci"},
)
request = MagicMock(spec=Request)
request.query_params = {}
with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update"):
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=None,
_user_role=None,
route=route,
request=request,
valid_token=valid_token,
request_data={},
)
def test_team_key_without_service_account_marker_still_rejected():
"""A team key without metadata.service_account_id is not a service account
and still cannot reach key-management routes."""
valid_token = UserAPIKeyAuth(
api_key="sk",
team_id="t1",
user_id=None,
metadata={},
)
request = MagicMock(spec=Request)
request.query_params = {}
with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update"):
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=None,
_user_role=None,
route="/key/generate",
request=request,
valid_token=valid_token,
request_data={},
)

View file

@ -1,6 +1,7 @@
from __future__ import annotations
from typing import Final
from unittest.mock import patch
import pytest
@ -168,6 +169,54 @@ def test_settings_store_refuses_a_runtime_write_to_a_config_owned_key() -> None:
assert store.source("max_parallel_requests") == "config"
@pytest.mark.timeout(10)
def test_settings_store_clear_removes_every_key_the_config_file_does_not_own() -> None:
store: Final = SettingsStore("general_settings")
store.load_yaml({"master_key": "os.environ/MASTER_KEY"})
store.apply_db_row("general_settings", {"max_parallel_requests": 3, "alerting": ["slack"]})
store.apply_runtime_values({"master_key": "sk-resolved", "alerting": ["slack"]})
store["allow_requests_on_db_unavailable"] = True
del store["alerting"]
store.clear()
assert dict(store) == {"master_key": "sk-resolved"}
assert "alerting" not in store
with pytest.raises(KeyError):
store["max_parallel_requests"]
@pytest.mark.timeout(10)
def test_settings_store_clear_then_refill_matches_a_plain_dict() -> None:
refilled: Final[dict[str, JsonValue]] = {"alerting": ["email"], "max_parallel_requests": 11}
store: Final = SettingsStore("general_settings")
store.update({"max_parallel_requests": 3, "alerting": ["slack"]})
store.clear()
store.update(refilled)
assert dict(store) == refilled
assert tuple(store) == tuple(refilled)
assert len(store) == len(refilled)
@pytest.mark.timeout(10)
@pytest.mark.parametrize("clear", (False, True))
def test_settings_store_survives_a_patch_dict_round_trip_when_the_config_file_owns_a_key(clear: bool) -> None:
store: Final = SettingsStore("general_settings")
store.load_yaml({"master_key": "os.environ/MASTER_KEY"})
store.apply_db_row("general_settings", {"max_parallel_requests": 3})
store.apply_runtime_values({"master_key": "sk-resolved", "max_parallel_requests": 3})
before: Final = dict(store)
with patch.dict(store, {"allow_requests_on_db_unavailable": True}, clear=clear):
assert store["allow_requests_on_db_unavailable"] is True
assert store["master_key"] == "sk-resolved"
assert ("max_parallel_requests" in store) is not clear
assert dict(store) == before
def test_settings_store_reports_the_config_owned_keys_a_write_would_change() -> None:
store: Final = SettingsStore("general_settings")
store.load_yaml({"max_parallel_requests": 3, "ui_access_mode": "admin_only"})

View file

@ -16,7 +16,11 @@ from redis.exceptions import DataError
import litellm
from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.db.db_spend_update_writer import (
_TEAM_ADVISORY_LOCK_SQL,
_TEAM_MEMBER_SPEND_SQL,
DBSpendUpdateWriter,
)
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
build_window_spend_transaction,
)
@ -913,79 +917,118 @@ async def test_commit_spend_updates_to_db_increments_agent_spend():
assert call_kwargs["data"] == {"spend": {"increment": response_cost}}
@pytest.mark.asyncio
async def test_commit_spend_updates_to_db_increments_team_member_spend_and_total_spend():
"""
Verify that _commit_spend_updates_to_db increments BOTH spend (cycle-scoped)
and total_spend (non-resetting) on LiteLLM_TeamMembership in a single
update_many call, using the same response_cost.
"""
db_writer = DBSpendUpdateWriter()
mock_batcher = MagicMock()
mock_batcher.litellm_verificationtoken = MagicMock()
mock_batcher.litellm_verificationtoken.update_many = MagicMock()
mock_batcher.litellm_usertable = MagicMock()
mock_batcher.litellm_usertable.update_many = MagicMock()
mock_batcher.litellm_teamtable = MagicMock()
mock_batcher.litellm_teamtable.update_many = MagicMock()
mock_batcher.litellm_teammembership = MagicMock()
mock_batcher.litellm_teammembership.update_many = MagicMock()
mock_batcher.litellm_organizationtable = MagicMock()
mock_batcher.litellm_organizationtable.update_many = MagicMock()
mock_batcher.litellm_tagtable = MagicMock()
mock_batcher.litellm_tagtable.update_many = MagicMock()
mock_batcher.litellm_agentstable = MagicMock()
mock_batcher.litellm_agentstable.update_many = MagicMock()
def _team_member_flush_fixtures() -> tuple[AsyncMock, MagicMock]:
"""A transaction and prisma client that record the raw statement the member spend flush runs."""
mock_transaction = AsyncMock()
mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
mock_transaction.__aexit__ = AsyncMock(return_value=False)
mock_transaction.batch_ = MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=mock_batcher),
__aexit__=AsyncMock(return_value=False),
)
)
mock_transaction.execute_raw = AsyncMock(return_value=1)
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
return mock_transaction, mock_prisma_client
mock_proxy_logging = MagicMock()
# Skip team-membership cache invalidation — out of scope for this test.
mock_proxy_logging.call_details.get = MagicMock(return_value=None)
team_id = "team-abc"
user_id = "user-xyz"
response_cost = 0.75
entity_id = f"team_id::{team_id}::user_id::{user_id}"
db_spend_update_transactions = {
def _team_member_only_transactions(spend_by_member_key: dict[str, float]) -> dict[str, dict[str, float]]:
return {
"user_list_transactions": {},
"end_user_list_transactions": {},
"key_list_transactions": {},
"team_list_transactions": {},
"team_member_list_transactions": {entity_id: response_cost},
"team_member_list_transactions": spend_by_member_key,
"org_list_transactions": {},
"tag_list_transactions": {},
"agent_list_transactions": {},
}
with patch("litellm.proxy.utils._raise_failed_update_spend_exception"):
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=mock_proxy_logging,
db_spend_update_transactions=db_spend_update_transactions,
)
mock_batcher.litellm_teammembership.update_many.assert_called_once()
call_kwargs = mock_batcher.litellm_teammembership.update_many.call_args[1]
assert call_kwargs["where"] == {"team_id": team_id, "user_id": user_id}
assert call_kwargs["data"] == {
"spend": {"increment": response_cost},
"total_spend": {"increment": response_cost},
}
@pytest.mark.asyncio
async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster_checked_upsert():
"""
Regression (LIT-5502): members added without a budget had no membership row, and the
previous update_many matched zero rows, so their spend was silently dropped.
The flush now takes the same per-team advisory lock the team endpoints hold, then runs
one INSERT ... ON CONFLICT statement for the whole batch that adds the cost to both spend
and total_spend and creates the missing row for a user still on the team roster, so no
per-team read can fail or time out ahead of the writes.
"""
db_writer = DBSpendUpdateWriter()
team_id = "team-abc"
user_id = "user-xyz"
response_cost = 0.75
mock_transaction, mock_prisma_client = _team_member_flush_fixtures()
mock_proxy_logging = MagicMock()
mock_proxy_logging.call_details.get = MagicMock(return_value=None)
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=mock_proxy_logging,
db_spend_update_transactions=_team_member_only_transactions(
{f"team_id::{team_id}::user_id::{user_id}": response_cost}
),
)
lock_call, spend_call = mock_transaction.execute_raw.await_args_list
lock_statement, locked_team_id = lock_call.args
assert lock_statement is _TEAM_ADVISORY_LOCK_SQL
assert locked_team_id == team_id
assert "pg_advisory_xact_lock(hashtext($1))" in lock_statement
statement, user_ids, team_ids, costs = spend_call.args
assert statement is _TEAM_MEMBER_SPEND_SQL
assert (list(user_ids), list(team_ids), list(costs)) == ([user_id], [team_id], [response_cost])
assert 'INSERT INTO "LiteLLM_TeamMembership"' in statement
assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id))" in statement
assert "ON CONFLICT (user_id, team_id) DO UPDATE" in statement
assert 'spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend' in statement
assert 'total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend' in statement
@pytest.mark.asyncio
async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_user():
"""
The member spend statement touches rows in the order of its input arrays, so the batch
is handed over sorted by (team_id, user_id), with each cost kept next to its member, and
each distinct team is locked once, in `sorted(team_ids)` order, the order /team/delete
locks in, so a concurrent flush and delete cannot deadlock. `eng` and `eng2` pin that:
sorting the composite keys instead would lock `eng2` first because `2` < `:`.
"""
db_writer = DBSpendUpdateWriter()
mock_transaction, mock_prisma_client = _team_member_flush_fixtures()
mock_proxy_logging = MagicMock()
mock_proxy_logging.call_details.get = MagicMock(return_value=None)
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=mock_proxy_logging,
db_spend_update_transactions=_team_member_only_transactions(
{
"team_id::eng2::user_id::user_x": 0.1,
"team_id::eng::user_id::user_y": 0.2,
"team_id::eng::user_id::user_x": 0.3,
"team_id::eng-b::user_id::user_x": 0.4,
}
),
)
*lock_calls, spend_call = mock_transaction.execute_raw.await_args_list
_statement, user_ids, team_ids, costs = spend_call.args
assert [lock_call.args for lock_call in lock_calls] == [
(_TEAM_ADVISORY_LOCK_SQL, "eng"),
(_TEAM_ADVISORY_LOCK_SQL, "eng-b"),
(_TEAM_ADVISORY_LOCK_SQL, "eng2"),
]
assert list(zip(team_ids, user_ids, costs)) == [
("eng", "user_x", 0.3),
("eng", "user_y", 0.2),
("eng-b", "user_x", 0.4),
("eng2", "user_x", 0.1),
]
@pytest.mark.asyncio
@ -2211,19 +2254,6 @@ async def test_commit_daily_tag_spend_no_requeue_on_success():
["team_a", "team_b", "team_c"],
id="team",
),
pytest.param(
"team_member_list_transactions",
{
"team_id::team_c::user_id::user_x": 0.1,
"team_id::team_a::user_id::user_x": 0.2,
"team_id::team_b::user_id::user_x": 0.3,
},
"litellm_teammembership",
"update_many",
"team_id",
["team_a", "team_b", "team_c"],
id="team_member",
),
pytest.param(
"org_list_transactions",
{"org_c": 0.1, "org_a": 0.2, "org_b": 0.3},
@ -2295,6 +2325,8 @@ async def test_commit_spend_updates_iterates_in_sorted_order(
)
)
mock_transaction.query_raw = AsyncMock(return_value=[])
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
@ -3054,6 +3086,7 @@ def _good_tx(mock_batcher):
tx = AsyncMock()
tx.__aenter__ = AsyncMock(return_value=tx)
tx.__aexit__ = AsyncMock(return_value=False)
tx.query_raw = AsyncMock(return_value=[])
tx.batch_ = MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=mock_batcher),

View file

@ -0,0 +1,409 @@
"""
Unit tests for the TypeSafe (Jev) compaction guardrail.
Tests cover:
- exchanges scored below relevance_threshold have their tool rows blanked while
assistant tool-call rows and kept exchanges pass through verbatim, without
mutating the caller's message list
- protected rows (system, last user, and the last tool exchange via the
last-assistant rule) are never sent to Jev even when long
- exchanges under min_chars_to_evaluate are skipped
- request shape: POST {api_base}/v1/systemone with Bearer auth, one noul
question per candidate keyed e<i>, task = last user text, results truncated
to max_result_chars_in_state
- identity return when there are no candidates or nothing is dropped
- fail_open forwards uncompacted on service failure; fail_closed raises
- response input_type passthrough and initialize_guardrail wiring
"""
from unittest.mock import AsyncMock, MagicMock, PropertyMock
import pytest
from fastapi import HTTPException
from litellm.proxy.guardrails.guardrail_hooks.typesafe import (
TypeSafeGuardrail,
guardrail_class_registry,
guardrail_initializer_registry,
initialize_guardrail,
)
from litellm.proxy.guardrails.guardrail_hooks.typesafe.typesafe import DROPPED_RESULT_TEXT
from litellm.types.guardrails import SupportedGuardrailIntegrations
from litellm.types.utils import GenericGuardrailAPIInputs
FAKE_API_BASE = "https://typesafe.example.com"
FAKE_API_KEY = "ts_test-key"
SYSTEM_TEXT = "You are a research assistant."
USER_TEXT = "Which 2026 EV has the longest range?"
TOOL_OUTPUT_LONG = "Result: EV range comparison. " * 40
TOOL_OUTPUT_SHORT = "short"
def _exchange(call_id: str, tool_text: str, name: str = "web_search") -> list[dict[str, object]]:
return [
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": call_id,
"type": "function",
"function": {"name": name, "arguments": '{"query": "ev"}'},
}
],
},
{"role": "tool", "tool_call_id": call_id, "name": name, "content": tool_text},
]
def _messages(*, tail: list[dict[str, object]] | None = None) -> list[dict[str, object]]:
base = [
{"role": "system", "content": SYSTEM_TEXT},
{"role": "user", "content": USER_TEXT},
]
return base + (tail or [])
def _make_guardrail(
handler: MagicMock | None = None,
*,
max_result_chars_in_state: int | None = None,
unreachable_fallback: str | None = None,
) -> TypeSafeGuardrail:
return TypeSafeGuardrail(
api_base=FAKE_API_BASE,
api_key=FAKE_API_KEY,
guardrail_name="typesafe",
default_on=True,
async_handler=handler or _make_handler({"e0": 0.9}),
max_result_chars_in_state=max_result_chars_in_state,
unreachable_fallback=unreachable_fallback,
)
def _make_handler(answers: dict[str, float], status: int = 200) -> MagicMock:
response = MagicMock()
response.status_code = status
response.json.return_value = {
"model": "jev-1.13.0",
"answers": {qid: {"type": "noul", "noul": score} for qid, score in answers.items()},
"usage": {"input_tokens": 10, "output_tokens": 1},
}
response.text = ""
handler = MagicMock()
handler.post = AsyncMock(return_value=response)
return handler
def _inputs(messages: list[dict[str, object]]) -> GenericGuardrailAPIInputs:
return GenericGuardrailAPIInputs(structured_messages=messages)
async def _apply(
guardrail: TypeSafeGuardrail, messages: list[dict[str, object]], input_type: str = "request"
) -> GenericGuardrailAPIInputs:
return await guardrail.apply_guardrail(
inputs=_inputs(messages),
request_data={},
input_type=input_type, # pyright: ignore[reportArgumentType] # test uses the same literal domain
logging_obj=None,
)
@pytest.mark.asyncio
async def test_low_noul_exchange_blanked_high_kept_and_input_not_mutated():
handler = _make_handler({"e0": 0.1, "e1": 0.95})
guardrail = _make_guardrail(handler)
messages = _messages(
tail=[
*_exchange("call_1", TOOL_OUTPUT_LONG),
*_exchange("call_2", TOOL_OUTPUT_LONG),
{"role": "assistant", "content": "still thinking"},
]
)
snapshot = [dict(m) for m in messages]
result = await _apply(guardrail, messages)
out = result["structured_messages"]
assert out[3]["content"] == DROPPED_RESULT_TEXT
assert out[3]["tool_call_id"] == "call_1"
assert out[3]["role"] == "tool"
assert out[5]["content"] == TOOL_OUTPUT_LONG
assert out[2] == messages[2]
assert out[4] == messages[4]
assert out[6]["content"] == "still thinking"
assert messages == snapshot
@pytest.mark.asyncio
async def test_last_exchange_and_protected_rows_never_evaluated():
handler = _make_handler({"e0": 0.05})
guardrail = _make_guardrail(handler)
messages = _messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), *_exchange("call_2", TOOL_OUTPUT_LONG)])
result = await _apply(guardrail, messages)
payload = handler.post.call_args.kwargs["json"]
assert list(payload["questions"]) == ["e0"]
assert list(payload["state"]["tool_exchanges"]) == ["e0"]
assert payload["state"]["task"] == USER_TEXT
assert payload["state"]["system"] == SYSTEM_TEXT
out = result["structured_messages"]
assert out[3]["content"] == DROPPED_RESULT_TEXT
assert out[5]["content"] == TOOL_OUTPUT_LONG
@pytest.mark.asyncio
async def test_short_exchange_not_sent():
handler = _make_handler({"e0": 0.9})
guardrail = _make_guardrail(handler)
messages = _messages(
tail=[
*_exchange("call_1", TOOL_OUTPUT_SHORT),
*_exchange("call_2", TOOL_OUTPUT_LONG),
{"role": "assistant", "content": "done"},
]
)
result = await _apply(guardrail, messages)
payload = handler.post.call_args.kwargs["json"]
assert list(payload["questions"]) == ["e0"]
exchange = payload["state"]["tool_exchanges"]["e0"]
assert exchange["result"] == TOOL_OUTPUT_LONG
assert result is not None
@pytest.mark.asyncio
async def test_request_body_shape_and_truncation():
handler = _make_handler({"e0": 0.9})
guardrail = _make_guardrail(handler, max_result_chars_in_state=50)
messages = _messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "done"}])
await _apply(guardrail, messages)
kwargs = handler.post.call_args.kwargs
assert kwargs["url"].endswith("/v1/systemone")
assert kwargs["url"].startswith(FAKE_API_BASE)
assert kwargs["headers"]["Authorization"] == f"Bearer {FAKE_API_KEY}"
assert kwargs["headers"]["Content-Type"] == "application/json"
payload = kwargs["json"]
assert payload["model"] == "jev-latest"
assert list(payload["questions"]) == ["e0"]
assert payload["questions"]["e0"]["type"] == "noul"
assert "e0" in payload["questions"]["e0"]["instructions"]
assert payload["state"]["task"] == USER_TEXT
exchange = payload["state"]["tool_exchanges"]["e0"]
assert len(exchange["result"]) == 50
assert exchange["result"].startswith(TOOL_OUTPUT_LONG[:10])
assert exchange["result"].endswith(TOOL_OUTPUT_LONG[-11:])
assert list(exchange["tool_calls"]) == [{"name": "web_search", "arguments": '{"query": "ev"}'}]
@pytest.mark.asyncio
async def test_no_candidates_returns_identity_and_skips_http():
handler = _make_handler({})
guardrail = _make_guardrail(handler)
inputs = _inputs(_messages(tail=[{"role": "assistant", "content": "plain answer"}]))
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
assert result is inputs
handler.post.assert_not_called()
@pytest.mark.asyncio
async def test_all_above_threshold_returns_identity():
handler = _make_handler({"e0": 0.9})
guardrail = _make_guardrail(handler)
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
assert result is inputs
@pytest.mark.asyncio
async def test_fail_open_returns_inputs_on_exception():
handler = MagicMock()
handler.post = AsyncMock(side_effect=Exception("connection refused"))
guardrail = _make_guardrail(handler, unreachable_fallback="fail_open")
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
assert result is inputs
@pytest.mark.asyncio
async def test_fail_closed_raises_http_exception():
handler = MagicMock()
handler.post = AsyncMock(side_effect=Exception("connection refused"))
guardrail = _make_guardrail(handler, unreachable_fallback="fail_closed")
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
assert exc_info.value.status_code == 502
@pytest.mark.asyncio
async def test_fail_open_on_non_2xx():
handler = _make_handler({"e0": 0.9}, status=500)
guardrail = _make_guardrail(handler)
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
assert result is inputs
@pytest.mark.asyncio
async def test_response_input_type_passthrough():
handler = _make_handler({"e0": 0.05})
guardrail = _make_guardrail(handler)
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG)]))
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="response", logging_obj=None)
assert result is inputs
handler.post.assert_not_called()
def test_initialize_guardrail_applies_optional_params_and_registry_keys():
from litellm.types.guardrails import LitellmParams
litellm_params = LitellmParams(
guardrail="typesafe",
mode="pre_call",
api_key=FAKE_API_KEY,
api_base=FAKE_API_BASE,
optional_params={
"relevance_threshold": 0.5,
"min_chars_to_evaluate": 10,
"max_result_chars_in_state": 100,
},
)
callback = initialize_guardrail(litellm_params, {"guardrail_name": "jev-compaction"})
assert isinstance(callback, TypeSafeGuardrail)
assert callback.relevance_threshold == 0.5
assert callback.min_chars_to_evaluate == 10
assert callback.max_result_chars_in_state == 100
assert callback.unreachable_fallback == "fail_open"
assert guardrail_initializer_registry[SupportedGuardrailIntegrations.TYPESAFE.value] is initialize_guardrail
assert guardrail_class_registry[SupportedGuardrailIntegrations.TYPESAFE.value] is TypeSafeGuardrail
def test_missing_api_key_raises(monkeypatch):
monkeypatch.delenv("TYPESAFE_API_KEY", raising=False)
with pytest.raises(ValueError, match="requires an API key"):
TypeSafeGuardrail(api_key=None)
def test_get_config_model_and_ui_name():
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
TypeSafeGuardrailConfigModel,
)
assert TypeSafeGuardrail.get_config_model() is TypeSafeGuardrailConfigModel
assert TypeSafeGuardrailConfigModel.ui_friendly_name() == "TypeSafe (Jev) Compaction"
@pytest.mark.asyncio
async def test_non_list_and_non_dict_messages_return_identity():
guardrail = _make_guardrail()
not_a_list = GenericGuardrailAPIInputs(structured_messages={"role": "user"})
assert (
await guardrail.apply_guardrail(inputs=not_a_list, request_data={}, input_type="request", logging_obj=None)
is not_a_list
)
with_bad_row = _inputs(_messages(tail=[["not", "a", "dict"]]))
assert (
await guardrail.apply_guardrail(inputs=with_bad_row, request_data={}, input_type="request", logging_obj=None)
is with_bad_row
)
def test_odd_tool_call_shapes_yield_no_entries():
from litellm.proxy.guardrails.guardrail_hooks.typesafe.typesafe import _tool_call_entries
assert _tool_call_entries({"tool_calls": "not-a-list"}) == ()
assert _tool_call_entries({"tool_calls": None}) == ()
assert list(_tool_call_entries({"tool_calls": [42]})) == []
entries = _tool_call_entries({"tool_calls": [{"function": {"name": "web_search", "arguments": "{}"}}]})
assert list(entries) == [{"name": "web_search", "arguments": "{}"}]
@pytest.mark.asyncio
async def test_short_max_chars_uses_prefix_slice():
handler = _make_handler({"e0": 0.9})
guardrail = _make_guardrail(handler, max_result_chars_in_state=5)
await _apply(
guardrail, _messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}])
)
result = handler.post.call_args.kwargs["json"]["state"]["tool_exchanges"]["e0"]["result"]
assert result == TOOL_OUTPUT_LONG[:5]
@pytest.mark.asyncio
async def test_unreadable_json_body_fails_open():
handler = MagicMock()
response = MagicMock()
response.status_code = 200
response.text = "not json"
response.json.side_effect = ValueError("no json")
handler.post = AsyncMock(return_value=response)
guardrail = _make_guardrail(handler)
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
assert result is inputs
@pytest.mark.asyncio
async def test_malformed_answers_shape_fails_open():
handler = MagicMock()
response = MagicMock()
response.status_code = 200
response.text = '{"answers": "oops"}'
response.json.return_value = {"answers": "oops"}
handler.post = AsyncMock(return_value=response)
guardrail = _make_guardrail(handler)
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
assert result is inputs
@pytest.mark.asyncio
async def test_http_status_error_includes_status_and_undecodable_body():
import httpx
response = MagicMock()
response.status_code = 503
type(response).text = PropertyMock(side_effect=httpx.DecodingError("bad codec"))
handler = MagicMock()
handler.post = AsyncMock(side_effect=httpx.HTTPStatusError("unavailable", request=MagicMock(), response=response))
guardrail = _make_guardrail(handler)
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
assert result is inputs
@pytest.mark.asyncio
async def test_cancelled_jev_call_propagates():
import asyncio
handler = MagicMock()
handler.post = AsyncMock(side_effect=asyncio.CancelledError())
guardrail = _make_guardrail(handler)
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
with pytest.raises(asyncio.CancelledError):
await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
def test_optional_params_defaults_and_event_hook_coercion():
from litellm.proxy.guardrails.guardrail_hooks.typesafe import _coerce_event_hook, _optional_params
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
assert _coerce_event_hook("pre_call") is GuardrailEventHooks.pre_call
assert _coerce_event_hook(["pre_call", "post_call"]) == [
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
]
litellm_params = LitellmParams(guardrail="typesafe", mode="pre_call", api_key=FAKE_API_KEY)
params = _optional_params(litellm_params)
assert params.relevance_threshold is None
def test_typesafe_initializer_discoverable_via_hook_registries():
from litellm.proxy.guardrails.guardrail_registry import get_guardrail_initializer_from_hooks
initializers = get_guardrail_initializer_from_hooks()
assert initializers["typesafe"] is initialize_guardrail

View file

@ -422,7 +422,7 @@ def test_apply_patch_ops_invalid_entitlements_value_raises_400():
patch_ops = SCIMPatchOp(
Operations=[
SCIMPatchOperation(
op="replace", path="entitlements", value=[{"display": "no value"}]
op="replace", path="entitlements", value=[42]
)
]
)
@ -433,6 +433,22 @@ def test_apply_patch_ops_invalid_entitlements_value_raises_400():
assert exc_info.value.status_code == 400
def test_apply_patch_ops_replace_entitlements_without_value_member_is_stored_as_sent():
patch_ops = SCIMPatchOp(
Operations=[
SCIMPatchOperation(
op="replace", path="entitlements", value=[{"groups": ["S0506MKA55L"]}]
)
]
)
update_data, _ = _apply_patch_ops(
existing_user=_user_with_metadata({}), patch_ops=patch_ops
)
assert update_data["metadata"]["scim_entitlements"] == [{"groups": ["S0506MKA55L"]}]
def test_apply_patch_ops_add_without_value_raises_400_naming_value_member():
patch_ops = SCIMPatchOp(
Operations=[SCIMPatchOperation(op="add", path="entitlements")]

View file

@ -1,3 +1,4 @@
import json
import logging
import time
from collections.abc import Callable, Mapping, Sequence
@ -1303,6 +1304,75 @@ async def test_update_user_success(mocker):
assert call_args[1]["data"]["teams"] == ["new-team"]
@pytest.mark.asyncio
async def test_update_user_put_with_valueless_entitlements_deactivates_user(scim_test_client, mocker):
existing_user = mocker.MagicMock()
existing_user.teams = []
existing_user.metadata = {"scim_active": True}
updated_user = {
"user_id": "suspend-me",
"user_email": "suspend@example.com",
"user_alias": None,
"teams": [],
"metadata": "{}",
}
response_scim_user = SCIMUser(
schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
id="suspend-me",
userName="suspend-me",
active=False,
)
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
AsyncMock(return_value=mock_prisma_client),
)
mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable
"litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists",
AsyncMock(return_value=existing_user),
)
mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable
"litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes",
AsyncMock(),
)
set_keys_blocked_mock = mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable
"litellm.proxy.management_endpoints.scim.scim_v2._set_user_keys_blocked",
AsyncMock(return_value=1),
)
mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable
"litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user",
AsyncMock(return_value=response_scim_user),
)
async with scim_test_client as client:
response = await client.put(
"/scim/v2/Users/suspend-me",
json={
"schemas": ["urn:ietf:params:scim:schemas:core:2.0:User"],
"userName": "suspend-me",
"emails": [{"value": "suspend@example.com", "primary": True}],
"entitlements": [{"groups": ["S0506MKA55L", "S0506MKA56M"]}],
"roles": [{"display": "Viewer"}],
"active": False,
},
)
assert response.status_code == 200, response.text
assert response.json()["active"] is False
written_metadata = json.loads(mock_prisma_client.db.litellm_usertable.update.call_args.kwargs["data"]["metadata"])
assert written_metadata["scim_active"] is False
assert written_metadata["scim_entitlements"] == [{"groups": ["S0506MKA55L", "S0506MKA56M"]}]
assert written_metadata["scim_roles"] == [{"display": "Viewer"}]
set_keys_blocked_mock.assert_awaited_once_with(user_id="suspend-me", blocked=True)
@pytest.mark.asyncio
@pytest.mark.parametrize("groups", [None, []], ids=["groups-omitted", "groups-empty"])
async def test_update_user_without_groups_preserves_memberships_and_role(mocker, monkeypatch, groups):

View file

@ -1,13 +1,19 @@
import re
from collections.abc import Sequence
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import psycopg
import pytest
from psycopg.rows import dict_row
from pytest_postgresql import factories
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT
from litellm.proxy.management_endpoints.common_daily_activity import (
_adjust_dates_for_timezone,
_build_aggregated_sql_query,
@ -169,6 +175,7 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
"endpoint": "/v1/chat/completions",
"api_key": None,
"group_level": 62,
"distinct_api_keys": None,
"spend": 15.0,
"prompt_tokens": 150,
"completion_tokens": 75,
@ -181,31 +188,7 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
"endpoint": "/v1/embeddings",
"api_key": None,
"group_level": 62,
"spend": 3.0,
"prompt_tokens": 30,
"completion_tokens": 0,
"api_requests": 1,
"successful_requests": 1,
},
# (date, endpoint, api_key) — populates the per-key sub-bucket
{
**base,
"date": "2024-01-01",
"endpoint": "/v1/chat/completions",
"api_key": "key-1",
"group_level": 30,
"spend": 15.0,
"prompt_tokens": 150,
"completion_tokens": 75,
"api_requests": 2,
"successful_requests": 2,
},
{
**base,
"date": "2024-01-01",
"endpoint": "/v1/embeddings",
"api_key": "key-2",
"group_level": 30,
"distinct_api_keys": None,
"spend": 3.0,
"prompt_tokens": 30,
"completion_tokens": 0,
@ -219,6 +202,7 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
"endpoint": None,
"api_key": None,
"group_level": 63,
"distinct_api_keys": None,
"spend": 18.0,
"prompt_tokens": 180,
"completion_tokens": 75,
@ -232,12 +216,40 @@ async def test_get_daily_activity_aggregated_with_endpoint_breakdown():
"endpoint": None,
"api_key": None,
"group_level": 127,
"distinct_api_keys": None,
"spend": 18.0,
"prompt_tokens": 180,
"completion_tokens": 75,
"api_requests": 3,
"successful_requests": 3,
},
# (date, endpoint, api_key) — populates the per-key sub-bucket
{
**base,
"date": "2024-01-01",
"endpoint": "/v1/chat/completions",
"api_key": "key-1",
"group_level": 30,
"distinct_api_keys": 2,
"spend": 15.0,
"prompt_tokens": 150,
"completion_tokens": 75,
"api_requests": 2,
"successful_requests": 2,
},
{
**base,
"date": "2024-01-01",
"endpoint": "/v1/embeddings",
"api_key": "key-2",
"group_level": 30,
"distinct_api_keys": 2,
"spend": 3.0,
"prompt_tokens": 30,
"completion_tokens": 0,
"api_requests": 1,
"successful_requests": 1,
},
]
mock_prisma.db.query_raw = AsyncMock(return_value=mock_rows)
@ -474,9 +486,7 @@ async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash(
return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")]
)
mock_prisma.db.query_raw = AsyncMock(
return_value=[
{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"}
]
return_value=[{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"}]
)
result = await get_api_key_metadata(
@ -835,6 +845,7 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys():
"endpoint": "/v1/chat/completions",
"api_key": None,
"group_level": 62,
"distinct_api_keys": None,
"spend": 10.0,
"prompt_tokens": 100,
"completion_tokens": 50,
@ -847,6 +858,7 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys():
"endpoint": "/v1/chat/completions",
"api_key": "deleted-key-hash",
"group_level": 30,
"distinct_api_keys": 1,
"spend": 10.0,
"prompt_tokens": 100,
"completion_tokens": 50,
@ -1230,42 +1242,11 @@ class TestBuildAggregatedSqlQuery:
"user-1",
"bedrock/global.anthropic.claude-opus-4-8",
"sk-test",
PTU_SENTINEL_API_KEY,
]
assert "model = $4" in sql
assert "api_key = $5" in sql
def test_model_group_rollups_fall_back_to_model_name(self):
"""Aggregated model_groups rollups must fall back to model for group-less rows.
The (date, model_group) grouping level cannot recover the model column
after the fact (it is rolled up), so the fallback has to happen in SQL;
without it, group-less rows silently vanish from the model_groups
breakdown that the usage UI now renders by default. Group-less rows are
stored as empty strings, not NULL (spend_tracking_utils defaults
model_group to ""), so a plain COALESCE is not enough: the fallback must
be NULLIF-wrapped to catch both
"""
sql, _ = _build_aggregated_sql_query(
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=None,
start_date="2026-07-01",
end_date="2026-07-01",
model=None,
api_key=None,
)
normalized = " ".join(sql.split())
fallback = "COALESCE(NULLIF(model_group, ''), model)"
assert f"{fallback} AS model_group" in normalized
assert (
f"GROUPING(date, api_key, model, {fallback}, "
"custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level" in normalized
)
assert f"(date, {fallback}), (date, {fallback}, api_key)," in normalized
assert "(date, model_group)" not in normalized
assert "COALESCE(model_group, model)" not in normalized
class TestAggregatedEmptyEntityFilter:
_BUILDERS: Final = (_build_aggregated_sql_query, _build_entity_rollup_sql_query)
@ -1285,7 +1266,8 @@ class TestAggregatedEmptyEntityFilter:
normalized = " ".join(sql.split())
assert "IN ()" not in normalized
assert '"team_id" IN' not in normalized
assert params == ["2026-08-01", "2026-08-19"]
sentinel_params = [PTU_SENTINEL_API_KEY] if build is _build_aggregated_sql_query else []
assert params == ["2026-08-01", "2026-08-19", *sentinel_params]
@pytest.mark.parametrize("build", _BUILDERS)
def test_empty_entity_list_matches_nothing_rather_than_everything(self, build):
@ -1316,7 +1298,8 @@ class TestAggregatedEmptyEntityFilter:
normalized = " ".join(sql.split())
assert '"team_id" IN ($3, $4)' in normalized
assert "FALSE" not in normalized
assert params == ["2026-08-01", "2026-08-19", "team-alpha", "team-beta"]
sentinel_params = [PTU_SENTINEL_API_KEY] if build is _build_aggregated_sql_query else []
assert params == ["2026-08-01", "2026-08-19", "team-alpha", "team-beta", *sentinel_params]
@pytest.mark.asyncio
@ -1341,6 +1324,7 @@ async def test_get_daily_activity_aggregated_empty_result_set():
"mcp_namespaced_tool_name": None,
"endpoint": None,
"group_level": 127,
"distinct_api_keys": None,
"spend": None,
"prompt_tokens": None,
"completion_tokens": None,
@ -1385,6 +1369,305 @@ async def test_get_daily_activity_aggregated_empty_result_set():
assert result.metadata.total_compression_saved_tokens == 0
_aggregated_postgresql_proc: Final = factories.postgresql_proc()
_aggregated_postgresql: Final = factories.postgresql("_aggregated_postgresql_proc")
_DAILY_USER_SPEND_DDL: Final = """
CREATE TABLE "LiteLLM_DailyUserSpend" (
id TEXT PRIMARY KEY,
user_id TEXT,
date TEXT NOT NULL,
api_key TEXT NOT NULL,
model TEXT,
model_group TEXT,
custom_llm_provider TEXT,
mcp_namespaced_tool_name TEXT,
endpoint TEXT,
prompt_tokens BIGINT DEFAULT 0,
completion_tokens BIGINT DEFAULT 0,
cache_read_input_tokens BIGINT DEFAULT 0,
cache_creation_input_tokens BIGINT DEFAULT 0,
compression_saved_tokens BIGINT DEFAULT 0,
compression_savings_spend DOUBLE PRECISION DEFAULT 0,
prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0,
gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0,
autorouter_savings_spend DOUBLE PRECISION DEFAULT 0,
spend DOUBLE PRECISION DEFAULT 0,
api_requests BIGINT DEFAULT 0,
successful_requests BIGINT DEFAULT 0,
failed_requests BIGINT DEFAULT 0,
total_response_time_ms BIGINT DEFAULT 0,
timed_requests BIGINT DEFAULT 0
)
"""
def _seed_daily_user_spend(conn: psycopg.Connection, rows: Sequence[tuple[object, ...]]) -> None:
with conn.cursor() as cur:
cur.execute(_DAILY_USER_SPEND_DDL)
cur.executemany(
"""
INSERT INTO "LiteLLM_DailyUserSpend"
(id, user_id, date, api_key, model, model_group, custom_llm_provider,
endpoint, prompt_tokens, spend, api_requests, successful_requests)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
""",
rows,
)
conn.commit()
def _psycopg_query_raw(conn: psycopg.Connection, row_counts: list[int]):
"""Run the proxy's $N-parameterized SQL through psycopg, recording each result size."""
async def query_raw(sql: str, *params: str) -> list[dict[str, object]]:
converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql)
with conn.cursor(row_factory=dict_row) as cur:
cur.execute(
converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query
{f"p{i}": v for i, v in enumerate(params, start=1)},
)
rows: Final = cur.fetchall()
row_counts.append(len(rows))
return rows
return query_raw
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_bounds_api_key_rollups(
_aggregated_postgresql: psycopg.Connection,
):
"""Run the GROUPING SETS statement against real Postgres with more keys than the cap.
key-004 and key-005 tie on spend exactly at the USAGE_TOP_API_KEYS_LIMIT
cutoff; the api_key tiebreaker must keep key-004 and drop key-005. The PTU
sentinel outspends every key but must not take a slot. Excluded keys and the
sentinel still count toward the totals and the model rollup, which come from
the key-free arm.
"""
n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 5
key_rows: Final = [
(
f"row-{i:03d}",
f"user-{i:03d}",
"2026-06-01",
f"key-{i:03d}",
"gpt-5",
"",
"openai",
"/v1/chat/completions",
10,
6.0 if i == 4 else float(i + 1),
1,
1,
)
for i in range(n_keys)
]
sentinel_row: Final = (
"row-ptu",
None,
"2026-06-01",
PTU_SENTINEL_API_KEY,
"gpt-5",
"",
"azure",
None,
0,
1000.0,
0,
0,
)
_seed_daily_user_spend(_aggregated_postgresql, [*key_rows, sentinel_row])
key_spend: Final = sum(6.0 if i == 4 else float(i + 1) for i in range(n_keys))
row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts)
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
result = await get_daily_activity_aggregated(
prisma_client=mock_prisma,
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=None,
entity_metadata_field=None,
start_date="2026-06-01",
end_date="2026-06-01",
model=None,
api_key=None,
)
# Key-free arm: (), (date), (date, model), (date, model_group), two providers,
# one mcp NULL bucket, endpoint plus its NULL bucket = 9 rows regardless of key count.
# Per-key arm: six per-key grouping sets, each capped at the limit.
assert row_counts == [9 + 6 * USAGE_TOP_API_KEYS_LIMIT]
assert result.metadata.total_spend == pytest.approx(key_spend + 1000.0)
assert result.metadata.total_api_requests == n_keys
assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT
assert result.metadata.total_api_keys == n_keys
expected_top: Final = {f"key-{i:03d}" for i in range(6, n_keys)} | {"key-004"}
day: Final = result.results[0]
assert day.metrics.spend == pytest.approx(key_spend + 1000.0)
assert set(day.breakdown.api_keys) == expected_top
assert day.breakdown.api_keys["key-004"].metrics.spend == 6.0
assert "key-005" not in day.breakdown.api_keys
assert PTU_SENTINEL_API_KEY not in day.breakdown.api_keys
assert day.breakdown.models["gpt-5"].metrics.spend == pytest.approx(key_spend + 1000.0)
assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == expected_top
assert day.breakdown.providers["openai"].metrics.spend == pytest.approx(key_spend)
assert set(day.breakdown.providers["openai"].api_key_breakdown) == expected_top
assert day.breakdown.endpoints["/v1/chat/completions"].metrics.api_requests == n_keys
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_both_arms(
_aggregated_postgresql: psycopg.Connection,
):
"""An explicit api_key filter must scope the key-free totals and the per-key
rollups to that key alone, so the two arms never disagree."""
rows: Final = [
(
f"row-{i}",
f"user-{i}",
"2026-06-01",
f"key-{i}",
"gpt-5",
"",
"openai",
"/v1/chat/completions",
10,
float(i + 1),
1,
1,
)
for i in range(3)
]
_seed_daily_user_spend(_aggregated_postgresql, rows)
row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts)
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
result = await get_daily_activity_aggregated(
prisma_client=mock_prisma,
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=None,
entity_metadata_field=None,
start_date="2026-06-01",
end_date="2026-06-01",
model=None,
api_key="key-1",
)
assert result.metadata.total_spend == 2.0
assert result.metadata.total_api_keys == 1
day: Final = result.results[0]
assert set(day.breakdown.api_keys) == {"key-1"}
assert day.breakdown.api_keys["key-1"].metrics.spend == 2.0
assert day.breakdown.models["gpt-5"].metrics.spend == 2.0
assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == {"key-1"}
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_reports_exact_limit_key_count_as_complete(
_aggregated_postgresql: psycopg.Connection,
):
"""With exactly USAGE_TOP_API_KEYS_LIMIT keys nothing is dropped, and the
response must say so: total_api_keys equals the limit rather than exceeding it."""
rows: Final = [
(
f"row-{i:03d}",
f"user-{i:03d}",
"2026-06-01",
f"key-{i:03d}",
"gpt-5",
"",
"openai",
"/v1/chat/completions",
10,
float(i + 1),
1,
1,
)
for i in range(USAGE_TOP_API_KEYS_LIMIT)
]
_seed_daily_user_spend(_aggregated_postgresql, rows)
row_counts: Final[list[int]] = [] # mutable-ok: out-param for the query_raw shim
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, row_counts)
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
result = await get_daily_activity_aggregated(
prisma_client=mock_prisma,
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=None,
entity_metadata_field=None,
start_date="2026-06-01",
end_date="2026-06-01",
model=None,
api_key=None,
)
assert result.metadata.total_api_keys == USAGE_TOP_API_KEYS_LIMIT
assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT
assert len(result.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT
@pytest.mark.asyncio
async def test_get_daily_activity_aggregated_model_group_rollups_fall_back_to_model_name(
_aggregated_postgresql: psycopg.Connection,
):
"""Rows stored with an empty or NULL model_group must land in the model_groups
breakdown under their model name instead of vanishing from the usage UI."""
rows: Final = [
("row-0", "user-0", "2026-06-01", "key-0", "gpt-5", "gpt-5-eu", "openai", "/v1/chat/completions", 10, 7.0, 1, 1),
("row-1", "user-1", "2026-06-01", "key-1", "gpt-5", "", "openai", "/v1/chat/completions", 10, 3.0, 1, 1),
("row-2", "user-2", "2026-06-01", "key-2", "claude-x", None, "anthropic", "/v1/messages", 10, 2.0, 1, 1),
]
_seed_daily_user_spend(_aggregated_postgresql, rows)
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, [])
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
result = await get_daily_activity_aggregated(
prisma_client=mock_prisma,
table_name="litellm_dailyuserspend",
entity_id_field="user_id",
entity_id=None,
entity_metadata_field=None,
start_date="2026-06-01",
end_date="2026-06-01",
model=None,
api_key=None,
)
breakdown: Final = result.results[0].breakdown
assert set(breakdown.model_groups) == {"gpt-5-eu", "gpt-5", "claude-x"}
assert breakdown.model_groups["gpt-5-eu"].metrics.spend == 7.0
assert breakdown.model_groups["gpt-5"].metrics.spend == 3.0
assert breakdown.model_groups["claude-x"].metrics.spend == 2.0
assert set(breakdown.model_groups["gpt-5"].api_key_breakdown) == {"key-1"}
assert set(breakdown.models) == {"gpt-5", "claude-x"}
assert breakdown.models["gpt-5"].metrics.spend == 10.0
def _no_spend_record():
"""A rollup row for a key with no spend, where SUM() returns NULL (None)."""
return SimpleNamespace(
@ -2170,7 +2453,7 @@ def test_entity_rollup_sql_query_and_api_key_list_filter():
api_key=[],
)
assert "FALSE" in empty_sql
assert empty_params == ["2024-01-01", "2024-01-31"]
assert empty_params == ["2024-01-01", "2024-01-31", PTU_SENTINEL_API_KEY]
@pytest.mark.asyncio
@ -2204,10 +2487,10 @@ async def test_get_daily_activity_aggregated_with_entity_breakdown():
"successful_requests": 0,
}
main_rows = [
{**base, "date": None, "group_level": 127, "spend": 18.0},
{**base, "date": "2024-01-01", "group_level": 63, "spend": 18.0},
{**base, "date": "2024-01-01", "model": "gpt-4o", "group_level": 47, "spend": 18.0},
{**base, "date": "2024-01-01", "api_key": "key-1", "group_level": 31, "spend": 12.0},
{**base, "date": None, "group_level": 127, "distinct_api_keys": None, "spend": 18.0},
{**base, "date": "2024-01-01", "group_level": 63, "distinct_api_keys": None, "spend": 18.0},
{**base, "date": "2024-01-01", "model": "gpt-4o", "group_level": 47, "distinct_api_keys": None, "spend": 18.0},
{**base, "date": "2024-01-01", "api_key": "key-1", "group_level": 31, "distinct_api_keys": 1, "spend": 12.0},
]
entity_base = {
key: value

View file

@ -18,6 +18,7 @@ import inspect
from litellm.proxy._types import (
GenerateKeyRequest,
KeyManagementRoutes,
NewUserRequest,
LiteLLM_BudgetTable,
LiteLLM_ObjectPermissionBase,
@ -3297,7 +3298,7 @@ async def test_validate_key_team_change_with_member_permissions():
# Verify the permission check was called with correct parameters
mock_has_perms.assert_called_once_with(
team_member_object=mock_member_object,
team_member_role=mock_member_object.role,
team_table=mock_team,
route=KeyManagementRoutes.KEY_UPDATE.value,
)
@ -19937,3 +19938,130 @@ async def test_bulk_update_team_keys_runs_custom_key_policy_per_key(monkeypatch)
assert [policy_request.operation for policy_request in received] == ["update", "update"]
assert [policy_request.effective_key.max_budget for policy_request in received] == [50.0, 50.0]
assert [policy_request.effective_key.team_id for policy_request in received] == ["team-abc", "team-abc"]
class TestServiceAccountKeyGenerationCheck:
"""Service account keys (user_id=None, team_id set, metadata.service_account_id)
may only create keys for their own team."""
def _service_account_token(self, team_id: str) -> UserAPIKeyAuth:
return UserAPIKeyAuth(
api_key="sk-sa",
user_id=None,
team_id=team_id,
metadata={"service_account_id": "sa-1"},
)
def test_other_team_denied(self):
data = GenerateKeyRequest(team_id="team-b")
with pytest.raises(HTTPException) as exc_info:
key_generation_check(
team_table=None,
user_api_key_dict=self._service_account_token(team_id="team-a"),
data=data,
route=KeyManagementRoutes.KEY_GENERATE,
)
assert exc_info.value.status_code == 403
def test_personal_key_denied(self):
"""team_id=None would mint a personal key; service accounts may only
create keys for their own team."""
data = GenerateKeyRequest()
with pytest.raises(HTTPException) as exc_info:
key_generation_check(
team_table=None,
user_api_key_dict=self._service_account_token(team_id="team-a"),
data=data,
route=KeyManagementRoutes.KEY_GENERATE,
)
assert exc_info.value.status_code == 403
def test_own_team_with_permission_allowed(self):
team_table = LiteLLM_TeamTableCachedObj(
team_id="team-a",
members_with_roles=[],
team_member_permissions=["/key/generate"],
)
data = GenerateKeyRequest(team_id="team-a")
assert (
key_generation_check(
team_table=team_table,
user_api_key_dict=self._service_account_token(team_id="team-a"),
data=data,
route=KeyManagementRoutes.KEY_GENERATE,
)
is True
)
def test_own_team_without_permission_denied(self):
team_table = LiteLLM_TeamTableCachedObj(
team_id="team-a",
members_with_roles=[],
team_member_permissions=["/key/info"],
)
data = GenerateKeyRequest(team_id="team-a")
with pytest.raises(ProxyException) as exc_info:
key_generation_check(
team_table=team_table,
user_api_key_dict=self._service_account_token(team_id="team-a"),
data=data,
route=KeyManagementRoutes.KEY_GENERATE,
)
assert str(exc_info.value.code) == "401"
def _stub_service_account_generation(monkeypatch):
"""Stub the DB lookups generate_service_account_key_fn needs so the test
exercises only the service_account_id stamping and user_id clearing."""
from litellm.proxy import proxy_server
from litellm.proxy.management_endpoints import key_management_endpoints as kme
mock_helper = AsyncMock(return_value=MagicMock())
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(kme, "validate_team_id_used_in_service_account_request", AsyncMock())
monkeypatch.setattr(kme, "_common_key_generation_helper", mock_helper)
return mock_helper
@pytest.mark.asyncio
async def test_generate_service_account_key_stamps_service_account_id(monkeypatch):
"""generate_service_account_key_fn must stamp metadata.service_account_id
(key_alias fallback) so the key is identifiable as a service account by
is_team_service_account and check_if_token_is_service_account."""
from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_service_account_key_fn,
)
mock_helper = _stub_service_account_generation(monkeypatch)
data = GenerateKeyRequest(team_id="team-a", key_alias="sa-alias")
await generate_service_account_key_fn(
data=data,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"),
litellm_changed_by=None,
)
assert data.metadata is not None
assert data.metadata["service_account_id"] == "sa-alias"
assert data.user_id is None
mock_helper.assert_awaited_once()
@pytest.mark.asyncio
async def test_generate_service_account_key_generates_uuid_when_no_alias(monkeypatch):
"""Without key_alias, service_account_id falls back to a generated uuid."""
from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_service_account_key_fn,
)
_stub_service_account_generation(monkeypatch)
data = GenerateKeyRequest(team_id="team-a")
await generate_service_account_key_fn(
data=data,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"),
litellm_changed_by=None,
)
assert data.metadata is not None
assert data.metadata["service_account_id"]

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