scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge branch 'litellm_internal_staging' into litellm_/admin-ui-e2e-models-endpoints-bf273a

Browse source
This commit is contained in:
yuneng-jiang 2026-09-07 11:35:19 -07:00 • committed by GitHub
commit bb55565a82
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
140 changed files with 8854 additions and 1366 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/workflows/auto-close-duplicates.yml vendored
View file

@ -22,7 +22,7 @@ on:
permissions: {}
jobs:
test:
sweep-tests:
if: github.event_name == 'pull_request'
runs-on: ubuntu-latest
timeout-minutes: 5

129
.github/workflows/report-rust-release-wheel.yml vendored
View file

@ -1,129 +0,0 @@
name: Report LiteLLM Rust release wheel
on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consumes no PR artifacts or outputs
workflow_run:
workflows:
- LiteLLM Rust
types:
- completed
permissions: {}
concurrency:
group: ${{ github.workflow }}-${{ github.event.workflow_run.pull_requests[0].number || github.event.workflow_run.id }}
cancel-in-progress: false
jobs:
report-release-wheel:
name: report release wheel
if: >-
github.event.workflow_run.event == 'pull_request' &&
github.event.workflow_run.path == '.github/workflows/test-rust.yml' &&
github.event.workflow_run.head_repository.full_name == github.repository &&
github.event.workflow_run.pull_requests[0].number != null
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
pull-requests: write
steps:
- name: Link release wheel report on PR
uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
env:
COMMENT_MARKER: "<!-- litellm-release-wheel-size -->"
with:
script: |
const marker = process.env.COMMENT_MARKER;
const workflowRun = context.payload.workflow_run;
const allowedConclusions = new Set([
"action_required",
"cancelled",
"failure",
"neutral",
"skipped",
"stale",
"startup_failure",
"success",
"timed_out",
]);
if (
!allowedConclusions.has(workflowRun.conclusion) ||
workflowRun.event !== "pull_request" ||
workflowRun.path !== ".github/workflows/test-rust.yml" ||
workflowRun.head_repository?.full_name !==
`${context.repo.owner}/${context.repo.repo}` ||
workflowRun.pull_requests?.length !== 1
) {
throw new Error("unexpected source workflow");
}
const pullRequest = workflowRun.pull_requests[0];
const pullRequestNumber = pullRequest.number;
const headSha = workflowRun.head_sha;
const runId = workflowRun.id;
if (
!Number.isSafeInteger(pullRequestNumber) ||
pullRequestNumber <= 0 ||
!Number.isSafeInteger(runId) ||
runId <= 0 ||
!/^[0-9a-f]{40}$/.test(headSha) ||
pullRequest.head?.sha !== headSha
) {
throw new Error("invalid source workflow metadata");
}
const runUrl =
`${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` +
`/actions/runs/${runId}`;
const result =
workflowRun.conclusion === "success"
? "successfully"
: `with \`${workflowRun.conclusion}\``;
const body = [
marker,
"## LiteLLM Rust workflow",
"",
`Workflow completed ${result} for \`${headSha}\``,
"",
`[View workflow run](${runUrl})`,
].join("\n");
const comments = await github.paginate(github.rest.issues.listComments, {
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: pullRequestNumber,
per_page: 100,
});
const existing = comments.find(
(comment) =>
comment.user?.login === "github-actions[bot]" &&
comment.body?.startsWith(marker),
);
const currentPullRequest = (
await github.rest.pulls.get({
owner: context.repo.owner,
repo: context.repo.repo,
pull_number: pullRequestNumber,
})
).data;
if (
currentPullRequest.state !== "open" ||
currentPullRequest.head.repo?.full_name !==
`${context.repo.owner}/${context.repo.repo}` ||
currentPullRequest.head.sha !== headSha
) {
core.info("source workflow no longer matches the current pull request head");
return;
}
if (existing) {
await github.rest.issues.updateComment({
owner: context.repo.owner,
repo: context.repo.repo,
comment_id: existing.id,
body,
});
} else {
await github.rest.issues.createComment({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: pullRequestNumber,
body,
});
}

6
.github/workflows/test-code-quality.yml vendored
View file

@ -74,6 +74,12 @@ jobs:
- name: check_workflow_startup_safety
run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py
- name: check_workflow_job_name_collisions
run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_job_name_collisions.py
- name: test_workflow_job_name_collisions
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.py
- name: test_e2e_changed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py

111
.github/workflows/test-rust.yml vendored
View file

@ -7,6 +7,7 @@ on:
- ".cargo/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
@ -22,6 +23,7 @@ on:
- ".cargo/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
@ -34,102 +36,89 @@ concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
env:
CARGO_TERM_COLOR: always
jobs:
rust-checks:
name: rustfmt, clippy, test
rust-lint:
runs-on: ubuntu-latest
timeout-minutes: 10
defaults:
run:
working-directory: litellm-rust
env:
CARGO_TERM_COLOR: always
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Rust
run: rustup toolchain install
- run: rustup toolchain install --no-self-update
- name: Cache Cargo registry and target
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
- run: cargo fmt --check
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ hashFiles('rust-toolchain.toml', 'litellm-rust/Cargo.lock') }}
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-
${{ runner.os }}-cargo-${{ github.job }}-
- name: Check Rust formatting
run: cargo fmt --check
- run: cargo clippy --workspace --all-targets --locked -- -D warnings
- name: Run Clippy
run: cargo clippy --workspace --all-targets --locked -- -D warnings
- run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
- name: Run Clippy with Bedrock auth
run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
- run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings
- name: Run Clippy with all gateway features
run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings
- name: Run Rust tests
run: cargo test --workspace --locked
- name: Run core tests with Bedrock auth
run: cargo test -p litellm-core --features bedrock-auth --locked
# Not --all-features: python-config links libpython, which this job does not install.
- name: Run gateway tests with the server feature
run: cargo test -p litellm-ai-gateway --features server --locked
release-wheel:
name: release wheel
rust-test:
runs-on: ubuntu-latest
timeout-minutes: 20
permissions:
contents: read
env:
CARGO_TERM_COLOR: always
timeout-minutes: 30
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
- uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Set up Rust
run: rustup toolchain install
- run: rustup toolchain install --no-self-update
- name: Build release wheel
run: uv build --wheel --out-dir dist
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-${{ github.job }}-
- name: Build panic contract wheel
run: >-
- run: cargo test --workspace --locked
working-directory: litellm-rust
- run: cargo test -p litellm-core --features bedrock-auth --locked
working-directory: litellm-rust
- run: cargo test -p litellm-ai-gateway --features server --locked
working-directory: litellm-rust
- run: uv build --wheel --out-dir dist
- run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
- run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- run: >-
uv build --wheel --out-dir panic-dist
--config-setting "maturin.build-args=--features panic-test,extension-module"
- name: Smoke-test native panic unwinding
run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl
- name: Verify stripped native extension
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
- name: Test native route wheel
run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl

4
CLAUDE.md
View file

@ -37,7 +37,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
@ -52,7 +52,7 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a
Python max line length is 120, not 88
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on `litellm_internal_staging` in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. If your branch already carries a budget edit, drop it before opening the PR
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice

4
basedpyright-code-budget.json
View file

@ -84,7 +84,7 @@
"limit": 56
},
"reportPrivateUsage": {
"limit": 1808
"limit": 1804
},
"reportRedeclaration": {
"limit": 8
@ -135,7 +135,7 @@
"limit": 21
},
"reportUnusedFunction": {
"limit": 138
"limit": 136
},
"reportUnusedImport": {
"limit": 542

2
litellm-rust/Cargo.lock generated
View file

@ -1415,6 +1415,8 @@ dependencies = [
"litellm-config",
"litellm-core",
"reqwest",
"rustls 0.23.42",
"rustls-native-certs",
"serde",
"serde_json",
"sha2 0.10.9",

2
litellm-rust/Cargo.toml
View file

@ -28,6 +28,8 @@ pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
rustls-native-certs = "0.8"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
sha2 = "0.10"

4
litellm-rust/crates/ai-gateway/Cargo.toml
View file

@ -20,6 +20,10 @@ litellm-config.workspace = true
# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
# Python proxy callbacks API.
reqwest.workspace = true
# rustls and its root store are direct dependencies so `io::tls` can build the
# one TLS config the outbound dials use; see that module for why it has to.
rustls.workspace = true
rustls-native-certs.workspace = true
# `sync` powers the bounded mpsc channel the realtime logger drains.
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] }
tokio-tungstenite.workspace = true

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

@ -3,3 +3,4 @@ pub mod ocr;
pub mod realtime;
pub mod realtime_pool;
pub mod responses_ws;
pub(crate) mod tls;

33
litellm-rust/crates/ai-gateway/src/io/realtime.rs
View file

@ -23,10 +23,12 @@ use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
use crate::io::tls::connect_upstream;
/// Environment variable holding the OpenAI API key (last-resort fallback).
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
@ -84,7 +86,7 @@ pub(crate) async fn dial_upstream(
.map_err(|err| Error::Auth(err.to_string()))?,
);
let (upstream, _response) = connect_async(request)
let (upstream, _response) = connect_upstream(request)
.await
.map_err(|err| Error::Network(err.to_string()))?;
Ok(upstream)
@ -284,6 +286,33 @@ mod tests {
serde_json::from_str(raw).expect("valid event json")
}
/// The realtime dial has to reach a `wss://` upstream without a process-wide
/// crypto provider installed, which is what dialing through `io::tls` buys.
#[tokio::test]
async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
let result = dial_upstream(
"gpt-realtime",
"sk-test",
Some(&format!("wss://127.0.0.1:{port}")),
)
.await;
assert!(matches!(result, Err(Error::Network(_))));
}
#[test]
fn resolve_api_key_prefers_param_then_blank_falls_through() {
assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test");

35
litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
View file

@ -14,7 +14,9 @@ use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use crate::io::tls::connect_upstream;
use crate::constants::{
DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
@ -49,14 +51,14 @@ impl ResponsesWebSocketConnection {
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
request.headers_mut().insert(header_name, header_value);
}
let connect = connect_async(request);
let connect = connect_upstream(request);
let result = match timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
Error::Network("Responses WebSocket connection timed out".to_string())
})?,
None => connect.await,
};
let (socket, _) = result.map_err(|error| match error {
let (socket, _) = result.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
@ -138,13 +140,13 @@ async fn dial_upstream(
);
let result = tokio::time::timeout(
Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS),
connect_async(request),
connect_upstream(request),
)
.await
.map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?;
result
.map(|(socket, _)| socket)
.map_err(|error| match error {
.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
@ -324,6 +326,29 @@ mod tests {
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
/// The Responses dial has to reach a `wss://` upstream without a process-wide
/// crypto provider installed, which is what dialing through `io::tls` buys.
#[tokio::test]
async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
let result =
dial_upstream("gpt-5", "sk-test", Some(&format!("wss://127.0.0.1:{port}"))).await;
assert!(matches!(result, Err(Error::Network(_))));
}
async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("local address");

80
litellm-rust/crates/ai-gateway/src/io/tls.rs Normal file
View file

@ -0,0 +1,80 @@
//! Outbound WebSocket dials over a TLS config this crate builds once and owns.
//!
//! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth`
//! enables `rustls/aws-lc-rs`, so the bare `ClientConfig::builder()` that
//! `tokio-tungstenite` uses when handed no connector panics rather than guess
//! between them. Naming ring on a connector of our own settles that for these
//! dials without touching the process-wide default, and building the config
//! once keeps the platform trust store, which `tokio-tungstenite` would
//! otherwise re-read on every dial, off the dial path.
use std::io;
use std::sync::{Arc, OnceLock};
use rustls::{ClientConfig, RootCertStore};
use tokio::net::TcpStream;
use tokio_tungstenite::tungstenite::Error;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::error::TlsError;
use tokio_tungstenite::tungstenite::handshake::client::Response;
use tokio_tungstenite::{
Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
};
static TLS_CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
fn build_config() -> Result<ClientConfig, Box<Error>> {
let native = rustls_native_certs::load_native_certs();
let roots = {
let mut store = RootCertStore::empty();
let (added, _ignored) = store.add_parsable_certificates(native.certs);
if added == 0 {
return Err(Box::new(Error::Io(io::Error::other(format!(
"no usable native root certificates: {:?}",
native.errors
)))));
}
store
};
ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
.with_safe_default_protocol_versions()
.map(|builder| builder.with_root_certificates(roots).with_no_client_auth())
.map_err(|error| Box::new(Error::Tls(TlsError::Rustls(error))))
}
fn tls_config() -> Result<Arc<ClientConfig>, Box<Error>> {
if let Some(config) = TLS_CONFIG.get() {
return Ok(Arc::clone(config));
}
let built = Arc::new(build_config()?);
Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built)))
}
pub(crate) async fn connect_upstream<R>(
request: R,
) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), Box<Error>>
where
R: IntoClientRequest + Unpin,
{
let request = request.into_client_request().map_err(Box::new)?;
let connector = match request.uri().scheme_str() {
Some("wss") => Some(Connector::Rustls(tls_config()?)),
_ => None,
};
connect_async_tls_with_config(request, None, false, connector)
.await
.map_err(Box::new)
}
#[cfg(test)]
mod tests {
use super::build_config;
#[test]
fn builds_a_usable_config_with_both_provider_features_enabled() {
let config = build_config().expect("a client config");
assert!(!config.crypto_provider().cipher_suites.is_empty());
}
}

12
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
View file

@ -265,6 +265,12 @@ impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLi
Box::pin(async move { Ok(request) })
}
#[tracing::instrument(
name = "success_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_success_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
@ -288,6 +294,12 @@ impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLi
})
}
#[tracing::instrument(
name = "failure_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,

48
litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs Normal file
View file

@ -0,0 +1,48 @@
//! Guards the wiring, not just the helper: a `wss://` dial through the public
//! API has to resolve its own crypto provider, in a test binary where nothing
//! has installed a process-wide one, and has to leave it uninstalled.
use std::collections::HashMap;
use std::time::Duration;
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection;
use tokio::net::TcpListener;
async fn dead_tls_server() -> u16 {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
port
}
#[tokio::test]
async fn dialing_wss_returns_an_error_instead_of_panicking() {
let port = dead_tls_server().await;
let result = ResponsesWebSocketConnection::connect_url(
&format!("wss://127.0.0.1:{port}/"),
&HashMap::new(),
Some(Duration::from_secs(10)),
)
.await;
assert!(
result.is_err(),
"a plain TCP server cannot finish a TLS handshake"
);
assert!(
rustls::crypto::CryptoProvider::get_default().is_none(),
"the dial settles its provider on its own connector, not process-wide"
);
}

52
litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
View file

@ -11,9 +11,13 @@ use litellm_ai_gateway::integrations::custom_logger::{
use litellm_ai_gateway::integrations::types::RequestMetadata;
use litellm_ai_gateway::ocr::{OcrRequest, ocr};
use litellm_core::error::Error;
#[cfg(feature = "trace-parity")]
use litellm_core::observability::FunctionTrace;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
#[cfg(feature = "trace-parity")]
use tracing::instrument::WithSubscriber;
async fn read_http_headers(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
@ -320,14 +324,17 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
GuardrailEventHook::PreCall,
GuardrailEventHook::DuringCall,
]));
let response = ocr(OcrRequest {
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
@ -339,9 +346,10 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
..Default::default()
},
litellm_call_id: Some("ocr-call-1"),
})
.await
.expect("ocr request succeeds");
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let response = call.await.expect("ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
assert_eq!(
@ -359,6 +367,16 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
error_kind: None,
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["success_callback"]
);
let request = server.await.expect("server task completes");
assert!(request.contains(r#""guarded_pre":true"#), "{request}");
@ -388,14 +406,17 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
});
let logger = Arc::new(RecordingOcrLogger::default());
let err = ocr(OcrRequest {
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
@ -404,9 +425,10 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("ocr-call-2"),
})
.await
.expect_err("provider error propagates");
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let err = call.await.expect_err("provider error propagates");
assert!(matches!(err, Error::Http { status: 500, .. }));
server.await.expect("server task completes");
@ -421,6 +443,16 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
error_kind: Some("HttpError".to_string()),
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["failure_callback"]
);
}
#[tokio::test]

28
litellm-rust/crates/python-bridge/src/function_trace.rs
View file

@ -1,3 +1,4 @@
use std::fmt::Display;
use std::future::Future;
use litellm_core::observability::{FunctionTrace, FunctionTraceEvent};
@ -6,17 +7,32 @@ use tracing::instrument::WithSubscriber;
#[derive(Serialize)]
pub(crate) struct TracedResponse<T> {
response: T,
#[serde(skip_serializing_if = "Option::is_none")]
response: Option<T>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
trace: Vec<FunctionTraceEvent>,
}
pub(crate) async fn capture<T, E>(
future: impl Future<Output = Result<T, E>>,
) -> Result<TracedResponse<T>, E> {
) -> Result<TracedResponse<T>, E>
where
E: Display,
{
let trace = FunctionTrace::default();
let response = future.with_subscriber(trace.dispatcher()).await?;
Ok(TracedResponse {
response,
trace: trace.events(),
let result = future.with_subscriber(trace.dispatcher()).await;
let events = trace.events();
Ok(match result {
Ok(response) => TracedResponse {
response: Some(response),
error: None,
trace: events,
},
Err(error) => TracedResponse {
response: None,
error: Some(error.to_string()),
trace: events,
},
})
}

9
litellm-rust/crates/python-bridge/src/routes/definition.rs
View file

@ -486,10 +486,11 @@ asyncio.run(exercise())
let code = CString::new(
r#"
result = routes.echo("traced")
assert result == {
"response": "traced",
"trace": [{"function": "execute_echo", "depth": 0}],
}
assert result["response"] == "traced", result
assert [event["function"] for event in result["trace"]] == ["execute_echo"], result
failure = routes.echo("error")
assert failure["error"] == "invalid request: synthetic error", failure
assert [event["function"] for event in failure["trace"]] == ["execute_echo"], failure
"#,
)
.expect("Python source should not contain null bytes");

21
litellm/batches/batch_utils.py
View file

@ -25,6 +25,27 @@ class BatchCostUsageResult:
failed_requests: int
_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"})
_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"})
def batch_cost_is_final(batch: Batch) -> bool:
"""Whether this retrieve of the batch is the one to account its cost from.
A batch still in flight has nothing to price, and a "completed" batch can report
no output_file_id for a moment before the output populates; pricing either records
$0 under the batch's single spend row and pins it there. Final means a completed
batch whose output file has arrived or whose counts prove no line succeeded, or
any other terminal status (failed, cancelled, expired).
"""
if batch.status not in _TERMINAL_BATCH_STATUSES:
return False
if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None:
return True
request_counts: Final = batch.request_counts
return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0
async def calculate_batch_cost_and_usage(
file_content_dictionary: list[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],

40
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -5,8 +5,11 @@ Handler for transforming /chat/completions api requests to litellm.responses req
import json
import os
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args
from openai.types.chat import ChatCompletion
from openai.types.responses import Response
from openai.types.responses.custom_tool_param import CustomToolParam
from openai.types.responses.response_input_param import (
FunctionCallOutput,
@ -33,7 +36,7 @@ from litellm.responses.sse_output_recovery import (
record_output_item_chunk,
record_output_text_chunk,
)
from litellm.responses.utils import normalize_responses_api_stream_options
from litellm.responses.utils import ResponsesAPIRequestUtils, normalize_responses_api_stream_options
from litellm.types.llms.openai import (
REASONING_EFFORT,
ChatCompletionAnnotation,
@ -43,6 +46,7 @@ from litellm.types.llms.openai import (
ChatCompletionToolParamFunctionChunk,
Reasoning,
ResponsesAPIOptionalRequestParams,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
@ -54,7 +58,7 @@ if TYPE_CHECKING:
)
from pydantic import BaseModel
from litellm import LiteLLMLoggingObj, ModelResponse
from litellm import LiteLLMLoggingObj
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.types.llms.openai import (
ALL_RESPONSES_API_TOOL_PARAMS,
@ -69,6 +73,28 @@ if TYPE_CHECKING:
from litellm.types.utils import Choices
_CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage"))
_RESPONSES_API_ONLY_FIELDS: Final = frozenset((*Response.model_fields, *ResponsesAPIResponse.model_fields)) - frozenset(
ChatCompletion.model_fields
)
def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]:
return MappingProxyType(
{
key: value
for key, value in (response_fields.items() if response_fields else ())
if value is not None and key not in _CHAT_COMPLETION_FIELDS and key not in _RESPONSES_API_ONLY_FIELDS
}
)
def _upstream_response_id(response_id: str | None) -> str | None:
if response_id is None:
return None
return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id)
class _ReasoningSummaryText(TypedDict):
type: str
text: str
@ -904,6 +930,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_response.usage),
)
model_response.id = _upstream_response_id(raw_response.id) or raw_response.id
for key, value in _provider_metadata(raw_response.model_extra).items():
setattr(model_response, key, value)
# Preserve hidden params from the ResponsesAPIResponse, especially the headers
# which contain important provider information like x-request-id
raw_response_hidden_params: Final = getattr(raw_response, "_hidden_params", {})
@ -1359,14 +1389,16 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
if event_type == "response.created":
# Initial response creation event
verbose_logger.debug("Chat provider: response.created -> %s", parsed_chunk)
created_response: Final = parsed_chunk.get("response")
return ModelResponseStream(
id=_upstream_response_id(created_response.get("id")) if created_response else None,
choices=[
StreamingChoices(
index=0,
delta=Delta(content=""),
finish_reason=None,
)
]
],
)
elif event_type == "response.output_item.added":
# New output item added
@ -1534,6 +1566,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
from litellm.responses.utils import ResponseAPILoggingUtils
usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_data.get("usage"))
provider_metadata: Final = _provider_metadata(response_data)
return ModelResponseStream(
choices=[
StreamingChoices(
@ -1546,6 +1579,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
)
],
usage=usage,
provider_specific_fields=dict(provider_metadata) or None, # mutable-ok: field is typed dict
)
else:
pass

10
litellm/constants.py
View file

@ -1901,6 +1901,16 @@ HTTP_FRAMING_HEADERS: Final[frozenset[str]] = frozenset(
}
)
PROVIDER_REQUEST_ID_HEADERS: Final[tuple[str, ...]] = (
"x-amzn-requestid",
"x-request-id",
"request-id",
"x-ms-request-id",
"apim-request-id",
"x-goog-request-id",
"cf-ray",
)
# Browser-facing security headers that a malicious or misconfigured upstream
# provider must not be able to set on the proxy's own response.
BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(

9
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -21,13 +21,10 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset(
}
)
# The per-deployment Rust opt-in.
RUST_KWARG_KEY: Final = "rust"
# Keys `completion()` forwards from its own kwargs into `get_litellm_params`,
# which are otherwise invisible to it because that call site passes explicit
# named arguments rather than `**kwargs`.
FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({RUST_KWARG_KEY})
FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS
# Pre-define optional kwargs keys as frozenset for O(1) lookups
# These are extracted from kwargs only if present, avoiding unnecessary .get() calls
@ -58,10 +55,6 @@ OPTIONAL_KWARGS_KEYS: Final = (
"itpm",
"otpm",
"use_xai_oauth",
# The per-deployment Rust opt-in. `all_litellm_params` keeps it out
# of the provider body; this keeps it *in* litellm_params, which is
# where the chat completions handlers read it from.
RUST_KWARG_KEY,
}
)
| AWS_CREDENTIAL_KWARGS_KEYS

49
litellm/litellm_core_utils/litellm_logging.py
View file

@ -36,12 +36,13 @@ from litellm._logging import (
verbose_logger,
)
from litellm._uuid import uuid
from litellm.batches.batch_utils import _handle_completed_batch
from litellm.batches.batch_utils import _handle_completed_batch, batch_cost_is_final
from litellm.caching.caching import DualCache, InMemoryCache
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.constants import (
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
PROVIDER_REQUEST_ID_HEADERS,
SENTRY_DENYLIST,
SENTRY_PII_DENYLIST,
)
@ -255,6 +256,30 @@ _in_memory_loggers: Final[list[CustomLogger]] = []
_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggingMetadata.__annotations__.keys())
def _get_provider_request_id(original_exception: Exception) -> str | None:
try:
error_response: Final = getattr(original_exception, "response", None)
header_sources: Final = (
_get_response_headers(original_exception),
getattr(error_response, "headers", None),
getattr(original_exception, "litellm_response_headers", None),
)
return next(
(
str(value)
for expected_header_name in PROVIDER_REQUEST_ID_HEADERS
for headers in header_sources
if isinstance(headers, Mapping)
for header_name, value in headers.items()
if isinstance(header_name, str) and header_name.lower() == expected_header_name and value
),
None,
)
except Exception:
return None
### GLOBAL VARIABLES ###
# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
@ -2899,13 +2924,6 @@ class Logging(LiteLLMLoggingBaseClass):
): # polling job will query these frequently, don't spam db logs
return
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
# check if file id is a unified file id
is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(result.id)
batch_cost: Final = kwargs.get("batch_cost", None)
batch_usage = kwargs.get("batch_usage", None)
batch_models = kwargs.get("batch_models", None)
@ -2913,9 +2931,7 @@ class Logging(LiteLLMLoggingBaseClass):
batch_failed_requests: Final = kwargs.get("batch_failed_requests", None)
has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models))
should_compute_batch_data: Final = (
not is_base64_unified_file_id or not has_explicit_batch_data and result.status == "completed"
)
should_compute_batch_data: Final = not has_explicit_batch_data and batch_cost_is_final(result)
if has_explicit_batch_data:
result._hidden_params["response_cost"] = batch_cost
result._hidden_params["batch_models"] = batch_models
@ -3918,11 +3934,12 @@ class Logging(LiteLLMLoggingBaseClass):
LiteLLMResponsesTransformationHandler,
)
served_id: Final = _provider_response_id(result)
try:
return LiteLLMResponsesTransformationHandler().transform_response(
translated: Final = LiteLLMResponsesTransformationHandler().transform_response(
model=self.model,
raw_response=result,
model_response=litellm.ModelResponse(id=_provider_response_id(result)),
model_response=litellm.ModelResponse(id=served_id),
logging_obj=self,
request_data={},
messages=[],
@ -3930,6 +3947,8 @@ class Logging(LiteLLMLoggingBaseClass):
litellm_params={},
encoding=litellm.encoding,
)
translated.id = served_id or translated.id
return translated
except Exception as e:
verbose_logger.debug(
"Responses API -> ModelResponse translation failed for "
@ -3937,7 +3956,7 @@ class Logging(LiteLLMLoggingBaseClass):
"usage-only ModelResponse to keep the spend_logs row.",
str(e),
)
model_response: Final = litellm.ModelResponse(id=_provider_response_id(result))
model_response: Final = litellm.ModelResponse(id=served_id)
model_response.model = self.model
usage: Final = getattr(result, "usage", None)
if usage is not None and ResponseAPILoggingUtils._is_response_api_usage(usage):
@ -5670,6 +5689,7 @@ class StandardLoggingPayloadSetup:
rate_limit_category: Final = validate_rate_limit_category(getattr(original_exception, "category", None))
rate_limit_type: Final = validate_rate_limit_type(getattr(original_exception, "rate_limit_type", None))
budget_error: Final = original_exception if isinstance(original_exception, BudgetExceededError) else None
provider_request_id: Final = _get_provider_request_id(original_exception) if original_exception else None
return StandardLoggingPayloadErrorInformation(
error_code=error_status,
@ -5677,6 +5697,7 @@ class StandardLoggingPayloadSetup:
llm_provider=_llm_provider_in_exception,
traceback=_redact_string(traceback_info),
error_message=_redact_string(error_message),
error_provider_request_id=provider_request_id,
error_rate_limit_category=rate_limit_category,
error_rate_limit_type=rate_limit_type,
error_budget_entity_type=budget_error.entity_type if budget_error else None,

9
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -26,6 +26,7 @@ from litellm.types.utils import (
PromptTokensDetailsWrapper,
ServiceTier,
Usage,
text_tokens_without_nested_reasoning,
)
from litellm.utils import get_model_info
@ -860,7 +861,7 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
)
or 0
)
text_tokens: Final = (
reported_text_tokens: Final = (
cast(
int | None,
getattr(usage.completion_tokens_details, "text_tokens", None),
@ -882,6 +883,12 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
or 0
)
video_tokens: Final = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0))
text_tokens: Final = text_tokens_without_nested_reasoning(
completion_tokens=usage.completion_tokens,
text_tokens=reported_text_tokens,
reasoning_tokens=reasoning_tokens,
other_modality_tokens=audio_tokens + image_tokens + video_tokens,
)
return CompletionTokensDetailsResult(
audio_tokens=audio_tokens,

202
litellm/litellm_core_utils/prompt_templates/image_handling.py
View file

@ -2,7 +2,11 @@
Helper functions to handle images passed in messages
"""
import asyncio
import base64
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from httpx import Response
@ -11,9 +15,11 @@ import litellm
from litellm import verbose_logger
from litellm.caching.caching import InMemoryCache
from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get
from litellm.types.llms.openai import AllMessageValues
MAX_IMGS_IN_MEMORY: Final = 10
MAX_CONCURRENT_REMOTE_MEDIA_FETCHES: Final = 20
in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY)
@ -72,6 +78,14 @@ def _process_image_response(response: Response, url: str) -> str:
return result
def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError":
verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict)
return litellm.ImageFetchError(
"Error: Unable to fetch image from URL. The proxy could not resolve this host or its URL policy rejected it; "
f"an admin can check the proxy log and `user_url_allowed_hosts` in general_settings. url={url}"
)
async def async_convert_url_to_base64(url: str) -> str:
if url.startswith("data:") and ";base64," in url:
return url
@ -93,6 +107,8 @@ async def async_convert_url_to_base64(url: str) -> str:
return _process_image_response(response, url)
except litellm.ImageFetchError:
raise
except SSRFError as e:
raise _rejected_image_fetch(url, e) from e
except Exception:
pass
raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}")
@ -119,8 +135,192 @@ def convert_url_to_base64(url: str) -> str:
return _process_image_response(response, url)
except litellm.ImageFetchError:
raise
except SSRFError as e:
raise _rejected_image_fetch(url, e) from e
except Exception as e:
verbose_logger.exception(e)
raise litellm.ImageFetchError(
f"Error: Unable to fetch image from URL after 3 attempts. url={url}",
)
_REMOTE_URL_PREFIXES: Final = ("http://", "https://")
@dataclass(frozen=True, slots=True)
class _RemoteImage:
part: Mapping[str, object]
image_url: Mapping[str, object] | None
url: str
@dataclass(frozen=True, slots=True)
class _RemoteFile:
part: Mapping[str, object]
file: Mapping[str, object]
url: str
def _as_mapping(value: object) -> Mapping[str, object] | None:
return value if isinstance(value, Mapping) else None # pyright: ignore[reportUnknownVariableType] # fields are parsed one by one
def _remote_url(candidate: object) -> str | None:
return candidate if isinstance(candidate, str) and candidate.startswith(_REMOTE_URL_PREFIXES) else None
_ANTHROPIC_MEDIA_BLOCK_TYPES: Final = frozenset({"document", "image"})
@dataclass(frozen=True, slots=True)
class _RemoteSource:
part: Mapping[str, object]
source: Mapping[str, object]
url: str
@dataclass(frozen=True, slots=True)
class RemoteMedia:
url: str
fields: Mapping[str, object]
_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
def inline_every_remote_url(_media: RemoteMedia) -> bool:
return True
def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None:
if fields.get("type") != "image_url":
return None
image_url: Final = fields.get("image_url")
image_url_fields: Final = _as_mapping(image_url)
url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url)
return _RemoteImage(fields, image_url_fields, url) if url is not None else None
def _parse_remote_file(fields: Mapping[str, object]) -> _RemoteFile | None:
file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None
url: Final = _remote_url(file.get("file_id")) if file is not None else None
return _RemoteFile(fields, file, url) if file is not None and url is not None else None
def _parse_remote_source(fields: Mapping[str, object]) -> _RemoteSource | None:
source: Final = _as_mapping(fields.get("source")) if fields.get("type") in _ANTHROPIC_MEDIA_BLOCK_TYPES else None
url: Final = _remote_url(source.get("url")) if source is not None and source.get("type") == "url" else None
return _RemoteSource(fields, source, url) if source is not None and url is not None else None
def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSource | None:
fields: Final = _as_mapping(part)
if fields is None:
return None
return _parse_remote_image(fields) or _parse_remote_file(fields) or _parse_remote_source(fields)
def _remote_media(remote: _RemoteImage | _RemoteFile | _RemoteSource) -> RemoteMedia:
match remote:
case _RemoteImage(_, image_url, url):
return RemoteMedia(url, image_url if image_url is not None else _NO_FIELDS)
case _RemoteFile(_, file, url):
return RemoteMedia(url, file)
case _RemoteSource(_, source, url):
return RemoteMedia(url, source)
_PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"})
def _inferred_format(file: Mapping[str, object], url: str) -> Mapping[str, str]:
return _PDF_FORMAT if "format" not in file and url.lower().endswith(".pdf") else MappingProxyType({})
def _inlined_image_url(image_url: Mapping[str, object] | None, data_url: str) -> Mapping[str, object] | str:
return {**image_url, "url": data_url} if image_url is not None else data_url # mutable-ok: json-serialized part
def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mapping[str, object]:
kept: Final = {k: v for k, v in file.items() if k != "file_id"} # mutable-ok: json-serialized message part
return {**kept, **_inferred_format(file, url), "file_data": data_url} # mutable-ok: json-serialized part
def _base64_source(url: str, data_url: str) -> Mapping[str, str]:
fetched_media_type, data = data_url.removeprefix("data:").split(";base64,", 1)
media_type: Final = "application/pdf" if url.lower().endswith(".pdf") else fetched_media_type
return {"type": "base64", "media_type": media_type, "data": data} # mutable-ok: json-serialized message part
def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) -> Mapping[str, object]:
match remote:
case _RemoteImage(part, image_url, _):
return {**part, "image_url": _inlined_image_url(image_url, data_url)} # mutable-ok: json-serialized part
case _RemoteFile(part, file, url):
return {**part, "file": _inlined_file(file, url, data_url)} # mutable-ok: json-serialized message part
case _RemoteSource(part, _, url):
return {**part, "source": _base64_source(url, data_url)} # mutable-ok: json-serialized message part
def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]:
content: Final = message.get("content")
return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one
def _inline_part(part: object, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]) -> object:
remote: Final = _parse_remote_part(part)
if remote is None or not should_inline(_remote_media(remote)):
return part
data_url: Final = data_urls.get(remote.url)
return _inline(remote, data_url) if data_url is not None else part
def _inline_message(
message: AllMessageValues, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]
) -> AllMessageValues:
parts: Final = _content_parts(message)
if not parts:
return message
inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks
_inline_part(part, data_urls, should_inline) for part in parts
]
inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message
return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined
async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str:
async with in_flight:
return await async_convert_url_to_base64(url)
async def _fetch_data_urls(remote_urls: tuple[str, ...]) -> tuple[str, ...]:
in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES)
fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight)) for url in remote_urls)
try:
return tuple(await asyncio.gather(*fetches))
except BaseException:
for fetch in fetches:
fetch.cancel()
await asyncio.gather(*fetches, return_exceptions=True)
raise
async def async_inline_remote_media(
messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues]
should_inline: Callable[[RemoteMedia], bool] = inline_every_remote_url,
) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues]
remote_urls: Final = tuple(
dict.fromkeys(
remote.url
for message in messages
for part in _content_parts(message)
if (remote := _parse_remote_part(part)) is not None and should_inline(_remote_media(remote))
)
)
if not remote_urls:
return messages
data_urls: Final = await _fetch_data_urls(remote_urls)
inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True)))
return [ # mutable-ok: transform_request takes a list
_inline_message(message, inlined, should_inline) for message in messages
]

3
litellm/litellm_core_utils/url_utils.py
View file

@ -19,6 +19,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config):
check but still resolve DNS and still rewrite HTTP to the resolved IP.
"""
import asyncio
import socket
from ipaddress import ip_address, ip_network
from typing import Any, Final, Protocol
@ -471,7 +472,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response
kwargs.pop("follow_redirects", None)
headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})}
for _ in range(_MAX_REDIRECTS):
validated_url, original_host = validate_url(url)
validated_url, original_host = await asyncio.to_thread(validate_url, url)
response = await fetcher.get(
validated_url,
headers={**headers_view["headers"], "Host": original_host},

49
litellm/llms/azure_ai/chat/transformation.py
View file

@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
from litellm.llms.azure.common_utils import BaseAzureLLM
from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config
from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error
from litellm.llms.openai.openai import OpenAIConfig
from litellm.llms.xai.chat.transformation import XAIChatConfig
@ -42,12 +43,37 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = (
)
class AzureAIGPT5Config(OpenAIGPT5Config):
@classmethod
def _model_map_lookup_name(cls, model: str) -> str:
"""Normalise a Foundry routing name to its cost-map key, when the map has one.
A Foundry deployment and its OpenAI-hosted namesake are different products with
different capabilities, so ``azure_ai/<model>`` is the entry to read whenever the map
carries it. Most gpt-5-family names have no ``azure_ai/`` row, though, and prefixing
those anyway costs them every flag: ``get_llm_provider`` re-resolves an ``azure_ai/``
name to the azure provider when a global AZURE_AI_API_BASE points at an
openai.azure.com host, ``azure/<model>`` is not a key either, so the lookup lands
nowhere and every effort answer degrades to False. A missing key defers to the base
resolver instead.
"""
prefixed: Final = model if model.startswith("azure_ai/") else f"azure_ai/{model}"
return prefixed if prefixed in litellm.model_cost else super()._model_map_lookup_name(model)
azureAIGPT5Config: Final = AzureAIGPT5Config()
class AzureAIStudioConfig(OpenAIConfig):
def get_supported_openai_params(self, model: str) -> list:
model_supports_tool_choice = True # azure ai supports this by default
if not supports_tool_choice(model=f"azure_ai/{model}"):
model_supports_tool_choice = False
supported_params = super().get_supported_openai_params(model)
supported_params = (
azureAIGPT5Config.get_supported_openai_params(model)
if azureAIGPT5Config.is_model_gpt_5_model(model)
else super().get_supported_openai_params(model)
)
if not model_supports_tool_choice:
filtered_supported_params: Final = []
for param in supported_params:
@ -61,6 +87,27 @@ class AzureAIStudioConfig(OpenAIConfig):
return supported_params
def map_openai_params(
self,
non_default_params: dict[str, object], # mutable-ok: OpenAIConfig.map_openai_params signature
optional_params: dict[str, object], # mutable-ok: OpenAIConfig.map_openai_params signature
model: str,
drop_params: bool,
) -> dict[str, object]: # mutable-ok: OpenAIConfig.map_openai_params signature
if not azureAIGPT5Config.is_model_gpt_5_model(model):
return super().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=drop_params,
)
return azureAIGPT5Config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=drop_params,
)
def _supports_stop_reason(self, model: str) -> bool:
"""
Check if the model supports stop tokens.

4
litellm/llms/base_llm/chat/transformation.py
View file

@ -411,6 +411,10 @@ class BaseConfig(ABC):
def has_custom_stream_wrapper(self) -> bool:
return False
@property
def uses_async_transform_request(self) -> bool:
return False
@property
def supports_stream_param_in_request_body(self) -> bool:
"""

19
litellm/llms/base_llm/image_edit/transformation.py
View file

@ -1,5 +1,6 @@
import types
from abc import ABC, abstractmethod
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any
import httpx
@ -102,6 +103,24 @@ class BaseImageEditConfig(ABC):
) -> tuple[dict, RequestFiles]:
pass
async def async_transform_image_edit_request(
self,
model: str,
prompt: str | None,
image: FileTypes | None,
image_edit_optional_request_params: Mapping[str, object],
litellm_params: GenericLiteLLMParams,
headers: Mapping[str, str],
) -> tuple[dict, RequestFiles]:
return self.transform_image_edit_request(
model=model,
prompt=prompt,
image=image,
image_edit_optional_request_params=dict(image_edit_optional_request_params),
litellm_params=litellm_params,
headers=dict(headers),
)
def finalize_image_edit_request_data(self, data: dict, resolved_request_url: str) -> dict:
"""
Last pass on the request dict after ``transform_image_edit_request``, using the

61
litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -8,7 +8,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_anthropic_image_obj,
)
from litellm.litellm_core_utils.prompt_templates.image_handling import (
async_convert_url_to_base64,
async_inline_remote_media,
convert_url_to_base64,
)
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
@ -172,6 +172,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
return _anthropic_request
@property
def uses_async_transform_request(self) -> bool:
return True
async def async_transform_request(
self,
model: str,
@ -180,26 +184,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
litellm_params: dict,
headers: dict,
) -> dict:
_anthropic_request: Final = self._build_bedrock_anthropic_request_base(
return self.transform_request(
model=model,
messages=messages,
messages=await async_inline_remote_media(messages),
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
await self._async_convert_document_url_sources_to_base64(_anthropic_request)
beta_list: Final = self._compute_bedrock_invoke_beta_headers(
model=model,
messages=messages,
optional_params=optional_params,
headers=headers,
)
if beta_list:
_anthropic_request["anthropic_beta"] = beta_list
return _anthropic_request
def _build_bedrock_anthropic_request_base(
self,
model: str,
@ -321,45 +313,6 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
"data": image_chunk["data"],
}
async def _async_convert_document_url_sources_to_base64(self, anthropic_request: dict) -> None:
"""
Async version of document URL conversion for async completion paths.
"""
messages: Final = anthropic_request.get("messages")
if not isinstance(messages, list):
return
for message in messages:
if not isinstance(message, dict):
continue
content = message.get("content")
if not isinstance(content, list):
continue
for block in content:
if not isinstance(block, dict) or block.get("type") != "document":
continue
source = block.get("source")
if not isinstance(source, dict) or source.get("type") != "url":
continue
source_url = source.get("url")
if not isinstance(source_url, str):
continue
inferred_format: str | None = None
if source_url.lower().endswith(".pdf"):
inferred_format = "application/pdf"
base64_url = await async_convert_url_to_base64(url=source_url)
image_chunk = convert_to_anthropic_image_obj(
openai_image_url=base64_url,
format=inferred_format,
)
block["source"] = {
"type": "base64",
"media_type": image_chunk["media_type"],
"data": image_chunk["data"],
}
def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict:
"""
Convert tool search entries to the format supported by the Bedrock Invoke API.

15
litellm/llms/bedrock/chat/mantle/transformation.py
View file

@ -10,6 +10,7 @@ at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth.
from collections.abc import AsyncIterator, Iterator
from typing import TYPE_CHECKING, Any, Final
from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeConfig,
)
@ -110,21 +111,13 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig):
litellm_params: dict,
headers: dict,
) -> dict:
model_id: Final = model.replace("mantle/", "", 1)
request: Final = self._build_bedrock_anthropic_request_base(
model=model_id,
messages=messages,
return self.transform_request(
model=model,
messages=await async_inline_remote_media(messages),
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
await self._async_convert_document_url_sources_to_base64(request)
return self._restore_mantle_body_fields(
request=request,
model_id=model_id,
optional_params=optional_params,
)
@staticmethod
def _restore_mantle_body_fields(request: dict, model_id: str, optional_params: dict) -> dict:

102
litellm/llms/black_forest_labs/image_edit/transformation.py
View file

@ -9,6 +9,7 @@ API Reference: https://docs.bfl.ai/
import base64
import time
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
import httpx
@ -16,7 +17,7 @@ from httpx._types import RequestFiles
import litellm
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.litellm_core_utils.url_utils import safe_get
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
@ -37,6 +38,22 @@ else:
LiteLLMLoggingObj = Any
_BFL_REQUEST_PARAMS: Final = (
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
"aspect_ratio",
"steps",
"guidance",
"grow_mask",
"top",
"bottom",
"left",
"right",
)
class BlackForestLabsImageEditConfig(BaseImageEditConfig):
"""
Configuration for Black Forest Labs image editing.
@ -85,34 +102,10 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
BFL-specific params are passed through directly.
"""
optional_params: Final[dict[str, object]] = {}
# Pass through BFL-specific params
bfl_params: Final = [
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
# Kontext-specific
"aspect_ratio",
# Fill/Inpaint-specific
"steps",
"guidance",
"grow_mask",
# Expand-specific
"top",
"bottom",
"left",
"right",
]
# Convert TypedDict to regular dict for access
params_dict: Final = dict(image_edit_optional_params)
for param in bfl_params:
if param in params_dict:
value = params_dict[param]
if value is not None:
optional_params[param] = value
params: Final[Mapping[str, object]] = image_edit_optional_params
for param in _BFL_REQUEST_PARAMS:
if (value := params.get(param)) is not None:
optional_params[param] = value
# Set default output format
if "output_format" not in optional_params:
@ -251,23 +244,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
"input_image": b64_image,
}
# Add optional params (only BFL-recognized parameters)
bfl_request_params: Final = [
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
"aspect_ratio",
"steps",
"guidance",
"grow_mask",
"top",
"bottom",
"left",
"right",
]
for key, value in image_edit_optional_request_params.items():
if key in bfl_request_params and value is not None:
if key in _BFL_REQUEST_PARAMS and value is not None:
request_body[key] = value
# Handle mask if provided (for inpainting)
@ -277,7 +255,39 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8")
# BFL uses JSON, not multipart - return empty files
return request_body, []
return request_body, ()
async def async_transform_image_edit_request(
self,
model: str,
prompt: str | None,
image: FileTypes | None,
image_edit_optional_request_params: Mapping[str, object],
litellm_params: GenericLiteLLMParams,
headers: Mapping[str, str],
) -> tuple[dict, RequestFiles]:
downloaded_image: Final = await self._fetch_remote_image(image)
downloaded_mask: Final = await self._fetch_remote_image(image_edit_optional_request_params.get("mask"))
return self.transform_image_edit_request(
model=model,
prompt=prompt,
image=image if downloaded_image is None else downloaded_image,
image_edit_optional_request_params=(
dict(image_edit_optional_request_params)
if downloaded_mask is None
else {**image_edit_optional_request_params, "mask": downloaded_mask}
),
litellm_params=litellm_params,
headers=dict(headers),
)
async def _fetch_remote_image(self, image: object) -> bytes | None:
candidate: Final = image[0] if isinstance(image, list) and image else image
if not isinstance(candidate, str) or not candidate.startswith(("http://", "https://")):
return None
response: Final = await async_safe_get(litellm.module_level_aclient, candidate, timeout=60.0)
response.raise_for_status()
return response.content
def transform_image_edit_response(
self,

207
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -116,6 +116,7 @@ from litellm.types.llms.anthropic_skills import (
Skill,
)
from litellm.types.llms.openai import (
AllMessageValues,
CreateBatchRequest,
CreateFileRequest,
FileContentRequest,
@ -163,13 +164,10 @@ from litellm.utils import (
def _rust_responses_websocket_enabled(
custom_llm_provider: str | None,
litellm_params: GenericLiteLLMParams,
) -> bool:
from litellm.rust_bridge.configuration import rust_enabled
raw_request_override: Final = litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
return custom_llm_provider == "openai" and rust_enabled(request_override=request_override)
return custom_llm_provider == "openai" and rust_enabled()
from .http_handler import get_shared_realtime_ssl_context
@ -488,7 +486,7 @@ class BaseLLMHTTPHandler:
def completion(
self,
model: str,
messages: list,
messages: list[AllMessageValues],
api_base: str | None,
custom_llm_provider: str,
model_response: ModelResponse,
@ -507,7 +505,7 @@ class BaseLLMHTTPHandler:
shared_session: Optional["ClientSession"] = None,
):
json_mode: Final[bool] = optional_params.pop("json_mode", False)
extra_body: Final[dict | None] = optional_params.pop("extra_body", None)
extra_body: Final[Mapping[str, object] | None] = optional_params.pop("extra_body", None)
provider_config = provider_config or ProviderConfigManager.get_provider_chat_config(
model=model, provider=litellm.LlmProviders(custom_llm_provider)
@ -522,14 +520,17 @@ class BaseLLMHTTPHandler:
)
# get config from model, custom llm provider
headers = provider_config.validate_environment(
api_key=api_key,
headers=headers or {},
model=model,
messages=messages,
optional_params=optional_params,
api_base=api_base,
litellm_params=litellm_params,
request_headers: Final = cast( # cast-ok: validate_environment is declared as a bare dict
"dict[str, object]",
provider_config.validate_environment(
api_key=api_key,
headers=headers or {},
model=model,
messages=messages,
optional_params=optional_params,
api_base=api_base,
litellm_params=litellm_params,
),
)
api_base = provider_config.get_complete_url(
@ -541,93 +542,117 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
)
data: dict[str, object] = provider_config.transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
if extra_body is not None:
data = {**data, **extra_body}
headers, signed_json_body = provider_config.sign_request(
headers=headers,
optional_params={
**optional_params,
**_aws_signing_overrides(optional_params, litellm_params),
},
request_data=data,
api_base=api_base,
api_key=api_key,
stream=stream,
fake_stream=fake_stream,
model=model,
)
## LOGGING
logging_obj.pre_call(
input=messages,
api_key=api_key,
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": headers,
},
)
# Check if stream was converted for WebSearch interception
# This is set by the async_pre_request_hook in WebSearchInterceptionLogger
if litellm_params.get("_websearch_interception_converted_stream", False):
logging_obj.model_call_details["websearch_interception_converted_stream"] = True
if acompletion is True:
if stream is True:
data = self._add_stream_param_to_request_body(
data=data,
provider_config=provider_config,
def sign_and_log(
transformed: dict[str, object], # mutable-ok: async_completion takes dict
) -> tuple[dict[str, object], dict[str, object], bytes | None]: # mutable-ok: async_completion takes dict
data: Final = {**transformed, **extra_body} if extra_body is not None else transformed
signed: Final = cast( # cast-ok: sign_request is declared as a bare dict
"tuple[dict[str, object], bytes | None]",
provider_config.sign_request(
headers=request_headers,
optional_params={
**optional_params,
**_aws_signing_overrides(optional_params, litellm_params),
},
request_data=data,
api_base=api_base,
api_key=api_key,
stream=stream,
fake_stream=fake_stream,
)
model=model,
),
)
logging_obj.pre_call(
input=messages,
api_key=api_key,
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": signed[0],
},
)
if litellm_params.get("_websearch_interception_converted_stream", False):
logging_obj.model_call_details["websearch_interception_converted_stream"] = True
return data, signed[0], signed[1]
def dispatch_async(
data: dict[str, object], # mutable-ok: async_completion takes dict
signed_headers: dict[str, object], # mutable-ok: async_completion takes dict
signed_json_body: bytes | None,
):
async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None
if stream is True:
return self.acompletion_stream_function(
model=model,
messages=messages,
api_base=api_base,
headers=headers,
headers=signed_headers,
custom_llm_provider=custom_llm_provider,
provider_config=provider_config,
timeout=timeout,
logging_obj=logging_obj,
data=data,
data=self._add_stream_param_to_request_body(
data=data,
provider_config=provider_config,
fake_stream=fake_stream,
),
fake_stream=fake_stream,
client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
client=async_client,
litellm_params=litellm_params,
json_mode=json_mode,
optional_params=optional_params,
signed_json_body=signed_json_body,
)
return self.async_completion(
custom_llm_provider=custom_llm_provider,
provider_config=provider_config,
api_base=api_base,
headers=signed_headers,
data=data,
timeout=timeout,
model=model,
model_response=model_response,
logging_obj=logging_obj,
api_key=api_key,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
client=async_client,
json_mode=json_mode,
signed_json_body=signed_json_body,
shared_session=shared_session,
)
else:
return self.async_completion(
custom_llm_provider=custom_llm_provider,
provider_config=provider_config,
api_base=api_base,
headers=headers,
data=data,
timeout=timeout,
model=model,
model_response=model_response,
logging_obj=logging_obj,
api_key=api_key,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
json_mode=json_mode,
signed_json_body=signed_json_body,
shared_session=shared_session,
if acompletion is True and provider_config.uses_async_transform_request:
async def transform_then_dispatch():
transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict
"dict[str, object]",
await provider_config.async_transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=request_headers,
),
)
return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed))
return transform_then_dispatch()
data, signed_headers, signed_json_body = sign_and_log(
provider_config.transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=request_headers,
)
)
if acompletion is True:
return dispatch_async(data, signed_headers, signed_json_body)
if stream is True:
data = self._add_stream_param_to_request_body(
@ -641,7 +666,7 @@ class BaseLLMHTTPHandler:
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
api_base=api_base,
headers=headers,
headers=signed_headers,
data=data,
signed_json_body=signed_json_body,
messages=messages,
@ -651,7 +676,7 @@ class BaseLLMHTTPHandler:
completion_stream, headers = self.make_sync_call(
provider_config=provider_config,
api_base=api_base,
headers=headers,
headers=signed_headers,
data=data,
signed_json_body=signed_json_body,
original_data=data,
@ -684,7 +709,7 @@ class BaseLLMHTTPHandler:
sync_httpx_client=sync_httpx_client,
provider_config=provider_config,
api_base=api_base,
headers=headers,
headers=signed_headers,
data=data,
signed_json_body=signed_json_body,
timeout=timeout,
@ -2403,9 +2428,7 @@ class BaseLLMHTTPHandler:
return None
from litellm.rust_bridge.configuration import rust_enabled
raw_request_override: Final = litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
if not rust_enabled(request_override=request_override):
if not rust_enabled():
return None
if has_agentic_hook:
return None
@ -6514,7 +6537,7 @@ class BaseLLMHTTPHandler:
@asynccontextmanager
async def _backend_connection():
if _rust_responses_websocket_enabled(custom_llm_provider, litellm_params):
if _rust_responses_websocket_enabled(custom_llm_provider):
from litellm.rust_bridge import responses_websocket as rust_responses_websocket
rust_backend: Final = await rust_responses_websocket.connect(
@ -6759,7 +6782,7 @@ class BaseLLMHTTPHandler:
litellm_params=dict(litellm_params),
)
data, files = image_edit_provider_config.transform_image_edit_request(
data, files = await image_edit_provider_config.async_transform_image_edit_request(
model=model,
image=image,
prompt=prompt,

14
litellm/llms/gemini/chat/transformation.py
View file

@ -12,7 +12,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionFileObject
from litellm.types.llms.vertex_ai import ContentType, PartType
from litellm.utils import supports_reasoning
from ...vertex_ai.gemini.transformation import _gemini_convert_messages_with_history
from ...vertex_ai.gemini.transformation import GEMINI_FILES_API_URI_PREFIX, _gemini_convert_messages_with_history
from ...vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
@ -127,7 +127,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
if element.get("type") == "image_url":
img_element = cast(ChatCompletionImageObject, element) # cast-ok: runtime type tag checked
_image_url, format, detail = _image_url_fields(img_element)
if _image_url and "https://" in _image_url:
if (
_image_url
and "https://" in _image_url
and not _image_url.startswith(GEMINI_FILES_API_URI_PREFIX)
):
image_obj = convert_to_anthropic_image_obj(_image_url, format=format)
converted_image_url = convert_generic_image_chunk_to_openai_image_obj(image_obj)
if detail is not None:
@ -147,7 +151,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
llm_provider="gemini",
)
file_id = _file_field.get("file_id")
if file_id and ("http://" in file_id or "https://" in file_id):
if (
file_id
and ("http://" in file_id or "https://" in file_id)
and not file_id.startswith(GEMINI_FILES_API_URI_PREFIX)
):
# Convert HTTP/HTTPS file URL to base64 data
try:
base64_data = convert_url_to_base64(file_id)

210
litellm/llms/mistral/audio_speech/transformation.py Normal file
View file

@ -0,0 +1,210 @@
"""
Support for Mistral Voxtral text-to-speech via ``/v1/audio/speech``.
API reference: https://docs.mistral.ai/api/#tag/audio/operation/audio_speech_v1_audio_speech_post
"""
import base64
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import httpx
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.text_to_speech.transformation import (
BaseTextToSpeechConfig,
TextToSpeechRequestData,
)
from litellm.secret_managers.main import get_secret_str
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.llms.openai import HttpxBinaryResponseContent
class MistralTextToSpeechException(BaseLLMException):
pass
class MistralTextToSpeechConfig(BaseTextToSpeechConfig):
TTS_BASE_URL: Final[str] = "https://api.mistral.ai/v1"
AUDIO_CONTENT_TYPES: Final[MappingProxyType[str, str]] = MappingProxyType(
{
"mp3": "audio/mpeg",
"wav": "audio/wav",
"pcm": "audio/pcm",
"flac": "audio/flac",
"opus": "audio/ogg",
}
)
DROPPED_RESPONSE_HEADERS: Final[frozenset[str]] = frozenset(
{"content-encoding", "transfer-encoding", "content-length", "content-type"}
)
OPENAI_VOICE_ALIASES: Final[MappingProxyType[str, str]] = MappingProxyType(
{
"alloy": "en_paul_neutral",
"echo": "gb_oliver_neutral",
"fable": "en_paul_cheerful",
"onyx": "en_paul_confident",
"nova": "gb_jane_sarcasm",
"shimmer": "gb_jane_sarcasm",
}
)
def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class contract returns a plain list
return ["voice", "response_format"] # mutable-ok: base class contract returns a plain list
def _map_openai_voice(self, voice_id: str) -> str:
return self.OPENAI_VOICE_ALIASES.get(voice_id.lower(), voice_id)
def _resolve_voice_id(self, voice: object) -> str | None:
if isinstance(voice, str) and voice.strip():
return self._map_openai_voice(voice.strip())
if isinstance(voice, Mapping):
candidates: Final = (voice.get(key) for key in ("voice_id", "id", "name"))
resolved: Final = next(
(candidate.strip() for candidate in candidates if isinstance(candidate, str) and candidate.strip()),
None,
)
return self._map_openai_voice(resolved) if resolved else None
return None
def map_openai_params(
self,
model: str,
optional_params: Mapping[str, object],
voice: object = None,
drop_params: bool = False,
kwargs: Mapping[str, object] | None = None,
) -> tuple[str | None, dict]: # mutable-ok: base class contract returns a plain dict
response_format: Final = optional_params.get("response_format")
ref_audio: Final = kwargs.get("ref_audio") if kwargs else None
voice_id_kwarg: Final = kwargs.get("voice_id") if kwargs else None
mapped_voice: Final = self._resolve_voice_id(voice) or self._resolve_voice_id(voice_id_kwarg)
mapped_params: Final = { # mutable-ok: base class contract returns a plain dict
key: value
for key, value in (("response_format", response_format), ("ref_audio", ref_audio))
if isinstance(value, str)
}
return mapped_voice, mapped_params
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
api_key: str | None = None,
api_base: str | None = None,
) -> dict: # mutable-ok: base class contract returns a plain dict
resolved_key: Final = api_key or get_secret_str("MISTRAL_API_KEY")
if resolved_key is None:
raise MistralTextToSpeechException(
status_code=401,
message="Mistral API key is required. Set MISTRAL_API_KEY or pass api_key.",
)
return { # mutable-ok: base class contract returns a plain dict
**headers,
"Authorization": f"Bearer {resolved_key}",
"Content-Type": "application/json",
}
def get_complete_url(
self,
model: str,
api_base: str | None,
litellm_params: Mapping[str, object],
) -> str:
configured_base: Final = (api_base or self.TTS_BASE_URL).rstrip("/")
versioned_base: Final = configured_base if configured_base.endswith("/v1") else f"{configured_base}/v1"
return f"{versioned_base}/audio/speech"
def transform_text_to_speech_request(
self,
model: str,
input: str,
voice: str | None,
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
headers: Mapping[str, str],
) -> TextToSpeechRequestData:
response_format: Final = optional_params.get("response_format")
ref_audio: Final = optional_params.get("ref_audio")
request_data: Final[TextToSpeechRequestData] = {
"dict_body": {
"model": model,
"input": input,
**({"voice_id": voice} if voice else {}),
**({"response_format": response_format} if isinstance(response_format, str) else {}),
**({"ref_audio": ref_audio} if isinstance(ref_audio, str) else {}),
},
"headers": {"Content-Type": "application/json"},
}
return request_data
def _requested_content_type(self, request: httpx.Request) -> str:
request_body: Final = json.loads(request.content or b"{}")
requested_format: Final = request_body.get("response_format")
if not isinstance(requested_format, str):
return "audio/mpeg"
return self.AUDIO_CONTENT_TYPES.get(requested_format, "audio/mpeg")
def transform_text_to_speech_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> "HttpxBinaryResponseContent":
from litellm.types.llms.openai import HttpxBinaryResponseContent
try:
response_json: Final = raw_response.json()
except (json.JSONDecodeError, ValueError):
raise MistralTextToSpeechException(
status_code=raw_response.status_code,
message=f"Non-JSON response from Mistral speech API: {raw_response.text[:500]}",
headers=raw_response.headers,
)
audio_b64: Final = response_json.get("audio_data")
if not isinstance(audio_b64, str) or not audio_b64:
raise MistralTextToSpeechException(
status_code=500,
message=f"No audio_data in Mistral speech response. Response keys: {tuple(response_json.keys())}",
headers=raw_response.headers,
)
try:
audio_bytes: Final = base64.b64decode(audio_b64, validate=True)
except ValueError:
raise MistralTextToSpeechException(
status_code=500,
message="Invalid base64 audio_data in Mistral speech response.",
headers=raw_response.headers,
)
retained_headers: Final = tuple(
(key, value)
for key, value in raw_response.headers.items()
if key.lower() not in self.DROPPED_RESPONSE_HEADERS
)
response_headers: Final = retained_headers + (
("content-length", str(len(audio_bytes))),
("content-type", self._requested_content_type(raw_response.request)),
)
binary_response: Final = httpx.Response(
status_code=200,
headers=response_headers,
content=audio_bytes,
request=raw_response.request,
)
return HttpxBinaryResponseContent(binary_response)
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict | httpx.Headers, # mutable-ok: BaseLLMException takes a plain dict or httpx.Headers
) -> BaseLLMException:
return MistralTextToSpeechException(
message=error_message,
status_code=status_code,
headers=headers,
)

52
litellm/llms/oci/chat/cohere.py
View file

@ -290,19 +290,14 @@ def handle_cohere_stream_chunk(
) -> ModelResponseStream:
"""Parse a single Cohere SSE chunk into a LiteLLM ModelResponseStream.
``prior_tool_calls_emitted`` lets the caller signal whether tool calls
were already emitted in earlier chunks of the same stream. When set, the
terminal consolidation chunk's tool calls are suppressed (they would
duplicate prior deltas); otherwise they are passed through so a stream
that delivers tool calls only on the terminal chunk doesn't silently
drop them.
``prior_text_emitted`` plays the analogous role for the ``text`` field:
when set, the terminal consolidation chunk's ``text`` is suppressed
(it would re-emit the full assembled response on top of prior deltas);
when unset (e.g. a degenerate stream that delivers the entire response
in a single SSE event carrying both ``chatHistory`` and ``finishReason``),
the text is passed through so the response content isn't silently lost.
OCI Cohere streams the answer as single-token ``text`` deltas, then restates
the whole assembled ``text`` on every chunk that carries ``toolCalls`` or
``chatHistory`` (the tool-calls event and the terminal event). Once the
caller reports that earlier chunks already emitted text
(``prior_text_emitted``), those restatements are dropped so the client does
not see the answer twice; a stream whose only text lives on such a chunk
keeps it. ``prior_tool_calls_emitted`` plays the same role for the tool
calls the terminal ``chatHistory`` chunk repeats.
"""
try:
typed_chunk: Final = CohereStreamChunk.model_validate(dict_chunk)
@ -315,33 +310,10 @@ def handle_cohere_stream_chunk(
if typed_chunk.index is None:
typed_chunk.index = 0
# OCI Cohere's terminal SSE event re-sends the full assembled response in
# `text` alongside a populated `chatHistory` and a non-null `finishReason`.
# Emitting that text would concatenate the whole response onto the
# already-streamed deltas. We require both signals to be present so that a
# future API change which adds `chatHistory` to intermediate chunks (or a
# rare early-populated case) doesn't silently drop legitimate token deltas.
is_terminal_consolidation: Final = typed_chunk.chatHistory is not None and typed_chunk.finishReason is not None
# On non-terminal text-free chunks (e.g. tool-call-only or keep-alive
# chunks) emit ``content=None`` rather than ``content=""`` so downstream
# stream-mergers that distinguish "no text in this delta" from "an
# explicitly empty text delta" behave correctly.
#
# We only suppress the terminal chunk's ``text`` when the caller has
# confirmed that text deltas were already emitted earlier — otherwise
# (e.g. a degenerate stream that delivers the whole response in a
# single SSE event), passing it through is the only chance to surface it.
text: Final[str | None] = None if (is_terminal_consolidation and prior_text_emitted) else typed_chunk.text
# Tool calls on the terminal consolidation chunk (whether from
# `typed_chunk.toolCalls` or from `chatHistory`) typically restate what
# was already streamed in intermediate chunks. Re-emitting them would
# mint fresh `uuid4` IDs and cause downstream consumers to execute each
# tool call twice. We only suppress when the caller has confirmed that
# tool calls were already emitted earlier — otherwise (e.g. a short
# response that delivers tool calls exclusively on the terminal chunk),
# passing them through is the only chance to surface them.
cohere_tool_calls = None if (is_terminal_consolidation and prior_tool_calls_emitted) else typed_chunk.toolCalls
restates_text: Final = typed_chunk.chatHistory is not None or typed_chunk.toolCalls is not None
restates_tool_calls: Final = typed_chunk.chatHistory is not None
text: Final[str | None] = None if (restates_text and prior_text_emitted) else typed_chunk.text
cohere_tool_calls: Final = None if (restates_tool_calls and prior_tool_calls_emitted) else typed_chunk.toolCalls
tool_calls: list[dict[str, object]] | None = None
if cohere_tool_calls:

16
litellm/llms/snowflake/chat/transformation.py
View file

@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
create_anthropic_image_param,
select_anthropic_content_block_type_for_file,
)
from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload
@ -421,6 +422,21 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body)
return self._transform_request_openai(model, messages, optional_params, stream, extra_body)
@property
def uses_async_transform_request(self) -> bool:
return True
async def async_transform_request(
self,
model: str,
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
optional_params: dict[str, object], # mutable-ok: BaseConfig signature
litellm_params: dict[str, object], # mutable-ok: BaseConfig signature
headers: dict[str, object], # mutable-ok: BaseConfig signature
) -> dict[str, object]: # mutable-ok: BaseConfig signature
inlined_messages: Final = await async_inline_remote_media(messages) if _is_claude_model(model) else messages
return self.transform_request(model, inlined_messages, optional_params, litellm_params, headers)
def _transform_request_openai(
self,
model: str,

32
litellm/llms/vertex_ai/gemini/transformation.py
View file

@ -7,6 +7,7 @@ Why separate file? Make it easy to see how transformation works
import json
import os
import re
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from urllib.parse import quote
@ -27,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_gemini_tool_call_result,
response_schema_prompt,
)
from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, async_inline_remote_media
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
from litellm.types.files import (
@ -68,6 +70,7 @@ _GCS_METADATA_VERTEX_BASE: Any | None = None
# Shared sync client for GCS JSON API metadata reads so proxy/SSL settings
# from litellm's HTTP stack apply (see Greptile review on PR #27278).
_GCS_METADATA_HTTP_HANDLER: HTTPHandler | None = None
GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/"
_GEMINI_MIME_TYPE_ALIASES: Final[dict[str, str]] = {
"image/jpg": "image/jpeg",
}
@ -556,7 +559,7 @@ def _process_gemini_media(
file_data = FileDataType(mime_type=mime_type, file_uri=image_url)
part: PartType = {"file_data": file_data}
return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata)
elif image_url.startswith("https://generativelanguage.googleapis.com/v1beta/files/"):
elif image_url.startswith(GEMINI_FILES_API_URI_PREFIX):
# Gemini Files API URIs — the file is already uploaded to Google's
# servers; pass the URI through as file_data without fetching it.
# These URLs return 403 when accessed directly, so we must not try
@ -1307,6 +1310,23 @@ def sync_transform_request_body(
)
def _explicit_mime_type(fields: Mapping[str, object]) -> str | None:
hint: Final = fields.get("format") or fields.get("mime_type") or fields.get("content_type")
return hint if isinstance(hint, str) else None
def _ai_studio_inlines(media: RemoteMedia) -> bool:
return not media.url.startswith(GEMINI_FILES_API_URI_PREFIX)
def _vertex_inlines(media: RemoteMedia) -> bool:
if media.url.startswith(GEMINI_FILES_API_URI_PREFIX):
return False
return media.url.startswith("http://") or (
_explicit_mime_type(media.fields) is None and _get_image_mime_type_from_url(media.url) is None
)
async def async_transform_request_body(
gemini_api_key: str | None,
messages: list[AllMessageValues],
@ -1348,13 +1368,17 @@ async def async_transform_request_body(
vertex_auth_header=vertex_auth_header,
)
if _openai_messages_may_need_sync_gcs_metadata_fetch(messages):
inlined_messages: Final = await async_inline_remote_media(
messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines
)
if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages):
# _transform_request_body may issue a sync httpx.get (up to 5s timeout)
# via _get_gcs_object_content_type to fetch GCS object metadata. Run the
# whole sync transformation on a worker thread so it does not block the
# async event loop.
return await asyncify(_transform_request_body)(
messages=messages,
messages=inlined_messages,
model=model,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
@ -1363,7 +1387,7 @@ async def async_transform_request_body(
)
return _transform_request_body(
messages=messages,
messages=inlined_messages,
model=model,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,

36
litellm/main.py
View file

@ -4474,7 +4474,7 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR
acompletion: Final = ctx.acompletion
api_base: Final = ctx.api_base
api_key: Final = ctx.api_key
client = _dispatch_client_http(ctx)
injected_client: Final = _dispatch_client_http(ctx)
custom_llm_provider: Final = ctx.custom_llm_provider
headers: Final = ctx.headers
litellm_params: Final = ctx.litellm_params
@ -4486,11 +4486,11 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR
shared_session: Final = ctx.shared_session
stream: Final = ctx.stream
timeout: Final = ctx.timeout
client: Final = (
injected_client if injected_client is not None else (HTTPHandler(timeout=timeout) if stream is False else None)
) # Keep this here, otherwise, the httpx.client closes and streaming is impossible
try:
client = (
HTTPHandler(timeout=timeout) if stream is False else None
) # Keep this here, otherwise, the httpx.client closes and streaming is impossible
response: Final = base_llm_http_handler.completion(
model=model,
messages=messages,
@ -8389,6 +8389,34 @@ def speech(
client=client,
_is_async=aspeech or False,
)
elif custom_llm_provider == "mistral":
from litellm.llms.mistral.audio_speech.transformation import (
MistralTextToSpeechConfig,
)
mistral_tts_config: Final = text_to_speech_provider_config or MistralTextToSpeechConfig()
if api_base is not None:
litellm_params_dict["api_base"] = api_base
if api_key is not None:
litellm_params_dict["api_key"] = api_key
mistral_voice: Final[str | None] = voice if isinstance(voice, str) else None
response = base_llm_http_handler.text_to_speech_handler(
model=model,
input=input,
voice=mistral_voice,
text_to_speech_provider_config=mistral_tts_config,
text_to_speech_optional_params=optional_params,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=client,
_is_async=aspeech or False,
)
elif custom_llm_provider == "aws_polly":
from litellm.llms.aws_polly.text_to_speech.transformation import (
AWSPollyTextToSpeechConfig,

57
litellm/model_prices_and_context_window_backup.json
View file

@ -3485,6 +3485,55 @@
"supports_response_schema": true,
"supports_tool_choice": true
},
"azure_ai/gpt-6-astra": {
"cache_creation_input_token_cost": 1.25e-05,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-05,
"cache_read_input_token_cost": 1e-06,
"cache_read_input_token_cost_above_272k_tokens": 2e-06,
"input_cost_per_token": 1e-05,
"input_cost_per_token_above_272k_tokens": 2e-05,
"litellm_provider": "azure_ai",
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 5e-05,
"output_cost_per_token_above_272k_tokens": 7.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"source": "https://ai.azure.com/catalog/models/gpt-6-astra",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": false,
"supports_minimal_reasoning_effort": false,
"supports_native_streaming": true,
"supports_none_reasoning_effort": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_cache_breakpoint": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true,
"supports_xhigh_reasoning_effort": true
},
"azure_ai/gpt-5.5": {
"deprecation_date": "2027-10-26",
"cache_read_input_token_cost": 5e-07,
@ -7189,7 +7238,7 @@
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": true,
"supports_max_reasoning_effort": false,
"supports_minimal_reasoning_effort": false,
"supports_native_streaming": true,
"supports_none_reasoning_effort": true,
@ -7455,7 +7504,7 @@
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": true,
"supports_max_reasoning_effort": false,
"supports_minimal_reasoning_effort": false,
"supports_native_streaming": true,
"supports_none_reasoning_effort": true,
@ -34898,9 +34947,9 @@
"supports_audio_input": true
},
"mistral/voxtral-mini-tts-latest": {
"input_cost_per_character": 1.6e-05,
"litellm_provider": "mistral",
"mode": "audio_speech",
"output_cost_per_character": 1.6e-05,
"source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03",
"supported_endpoints": [
"/v1/audio/speech"
@ -56467,9 +56516,9 @@
"supports_audio_input": true
},
"mistral/voxtral-mini-tts-2603": {
"input_cost_per_character": 1.6e-05,
"litellm_provider": "mistral",
"mode": "audio_speech",
"output_cost_per_character": 1.6e-05,
"source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03",
"supported_endpoints": [
"/v1/audio/speech"

85
litellm/ocr/main.py
View file

@ -29,6 +29,8 @@ from litellm.llms.base_llm.ocr.transformation import (
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import ocr as rust_ocr_bridge
from litellm.rust_bridge.bindings import native_exception_types
from litellm.rust_bridge.configuration import rust_enabled
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager, client
@ -196,12 +198,6 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool:
raw_request_override: Final = prepared_request.litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override)
def _rust_bridge_optional_params(
prepared_request: _PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
@ -286,6 +282,33 @@ def _prepare_rust_ocr_call(
)
def _map_rust_ocr_error(
error: Exception,
prepared_request: _PreparedOCRRequest,
exception_types: tuple[type[BaseException], type[BaseException]] | None,
) -> Exception:
if exception_types is None:
return error
_, upstream_error = exception_types
if not isinstance(error, upstream_error):
return error
error_args: Final = cast( # cast-ok: BaseException.args is typed with Any in the standard library stubs
tuple[object, ...], error.args
)
status_value: Final = error_args[0] if error_args else 0
message_value: Final = error_args[1] if len(error_args) > 1 else str(error)
status: Final = status_value if isinstance(status_value, int) else 0
message: Final = message_value if isinstance(message_value, str) else str(message_value)
error_factory: Final = cast( # cast-ok: the legacy provider interface leaves callable parameters untyped
Callable[..., Exception], prepared_request.provider_config.get_error_class
)
return error_factory(
error_message=message,
status_code=status or 500,
headers={}, # mutable-ok: provider error factories require a concrete header dict
)
def _run_rust_ocr(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
@ -296,16 +319,19 @@ def _run_rust_ocr(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
)
rust_response: Final = rust_ocr_bridge.ocr(
model=prepared_request.model,
document=prepared_request.document,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout=prepared_request.effective_timeout,
)
try:
rust_response: Final = rust_ocr_bridge.ocr(
model=prepared_request.model,
document=prepared_request.document,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout=prepared_request.effective_timeout,
)
except Exception as error:
raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error
if rust_response is None:
return None
return OCRResponse.model_validate(rust_response)
@ -321,16 +347,19 @@ async def _run_rust_aocr(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
)
rust_response: Final = await rust_ocr_bridge.aocr(
model=prepared_request.model,
document=prepared_request.document,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout=prepared_request.effective_timeout,
)
try:
rust_response: Final = await rust_ocr_bridge.aocr(
model=prepared_request.model,
document=prepared_request.document,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout=prepared_request.effective_timeout,
)
except Exception as error:
raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error
if rust_response is None:
return None
return OCRResponse.model_validate(rust_response)
@ -430,7 +459,7 @@ async def aocr(
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
if _rust_ocr_supported(prepared) and rust_enabled():
from litellm.secret_managers.main import get_secret_str
rust_response: Final = await _run_rust_aocr(
@ -702,7 +731,7 @@ def ocr(
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
if _rust_ocr_supported(prepared) and rust_enabled():
from litellm.secret_managers.main import get_secret_str
rust_response: Final = _run_rust_ocr(

136
litellm/proxy/db/db_spend_update_writer.py
View file

@ -14,6 +14,7 @@ import time
import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload
import litellm
@ -84,6 +85,25 @@ else:
RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value})
def _is_batch_cost_row(payload: SpendLogsPayload) -> bool:
return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success"
_BATCH_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"})
def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]:
"""Reduce a batch's cost row to what tells the retrieves apart when logging is off.
A proxy run with spend logs disabled still needs one row per batch to charge it once,
so the row is written either way, but it carries no request of its own: no metadata,
no requester IP, no key, model, or token counts (LIT-7048).
"""
if disable_spend_logs is False:
return payload
return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS})
class _SpendBatch(Protocol):
litellm_usertable: BatchTable
litellm_verificationtoken: BatchTable
@ -215,7 +235,12 @@ class DBSpendUpdateWriter:
start_time: datetime | None,
end_time: datetime | None,
response_cost: float | None,
) -> None:
) -> bool:
"""Record the request's spend, answering whether its cost still needs charging.
False only for a batch retrieve whose cost row another retrieve already wrote,
so the caller leaves the key, team, and user counters alone (LIT-7048).
"""
from litellm.proxy.proxy_server import (
disable_spend_logs,
litellm_proxy_budget_name,
@ -232,7 +257,7 @@ class DBSpendUpdateWriter:
team_id,
)
if ProxyUpdateSpend.disable_spend_updates() is True:
return
return True
if token is not None and isinstance(token, str) and token.startswith("sk-"):
hashed_token = hash_token(token=token)
else:
@ -262,11 +287,12 @@ class DBSpendUpdateWriter:
if team_id is not None and team_id != "":
payload["team_id"] = team_id
if not await self._record_spend_log(
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
):
return False
if disable_spend_logs is False:
await self._insert_spend_log_to_db(
payload=payload,
prisma_client=prisma_client,
)
await self._enqueue_tool_usage_transaction(
payload=payload,
completion_response=completion_response,
@ -306,6 +332,7 @@ class DBSpendUpdateWriter:
)
verbose_proxy_logger.debug("Runs spend update on all tables")
return True
except Exception:
spend_log_error(
"Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue "
@ -318,7 +345,102 @@ class DBSpendUpdateWriter:
org_id,
end_user_id,
)
return
return True
async def _record_spend_log(
self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool
) -> bool:
if prisma_client is not None and _is_batch_cost_row(payload):
return await self._claim_batch_cost_spend_log(
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
)
if disable_spend_logs is False:
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client)
return True
async def _claim_batch_cost_spend_log(
self, payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool
) -> bool:
"""Write the batch's cost row now, or learn that another retrieve already did.
Every retrieve of one batch shares this row, so the insert that lands first owns
the charge and every later one finds the row and charges nothing (LIT-7048). Only
a row that recorded a charge counts: a failed retrieve, a request whose client
picked the batch id as its call id, and the $0 row an older proxy left behind
while the batch was still running all leave the charge to be made.
"""
from litellm.repositories.table_repositories import SpendLogsRepository
request_id: Final = payload["request_id"]
row: Final = _batch_cost_row_to_write(payload, disable_spend_logs)
spend_logs: Final = SpendLogsRepository(prisma_client).table
try:
claimed: Final = await spend_logs.create_many(
data=[prisma_client.jsonify_object(row)], # mutable-ok: prisma create_many takes a list
skip_duplicates=True,
)
if claimed == 1:
return True
existing: Final = await spend_logs.find_unique(
where={"request_id": request_id} # mutable-ok: prisma where clause
)
except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log
verbose_proxy_logger.warning(
"Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e
)
await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client)
return True
if existing is None or existing.call_type != CallTypes.aretrieve_batch.value or existing.status != "success":
verbose_proxy_logger.warning(
"Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own",
request_id,
getattr(existing, "call_type", None),
)
return True
if existing.spend > 0:
verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id)
return False
return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row)
async def _take_over_uncharged_batch_cost_row(
self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object]
) -> bool:
"""Take the batch's cost row over from the poll that left it charging nothing.
A pre-upgrade proxy wrote that row every time it polled the batch while it was still
running, so the charge is still to be made and the row still has to end up carrying
it. The row stops matching the moment it carries a charge, so it is one retrieve that
takes it over and charges, and every later one reads the charge and charges nothing.
"""
from litellm.repositories.table_repositories import SpendLogsRepository
request_id: Final = payload["request_id"]
if payload["spend"] <= 0:
verbose_proxy_logger.debug(
"Cost tracking skipped: this batch costs nothing and spend row %s says so", request_id
)
return False
try:
taken_over: Final = await SpendLogsRepository(prisma_client).table.update_many(
data=prisma_client.jsonify_object(
MappingProxyType({field: value for field, value in row.items() if field != "request_id"})
),
where={ # mutable-ok: prisma where clause
"request_id": request_id,
"call_type": CallTypes.aretrieve_batch.value,
"status": "success",
"spend": 0.0,
},
)
except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; the next retrieve takes the row over
verbose_proxy_logger.warning(
"Could not take over spend row %s, leaving this batch's cost to the next retrieve: %s", request_id, e
)
return False
if taken_over == 0:
verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id)
return False
return True
async def _enqueue_tool_usage_transaction(
self,

26
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, cast
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.batch_utils import batch_cost_is_final
from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
@ -37,6 +38,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
from litellm.proxy.utils import ProxyUpdateSpend
from litellm.types.utils import (
CallTypes,
LiteLLMBatch,
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
)
@ -248,6 +250,18 @@ class _ProxyDBLogger(CustomLogger):
)
_write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata)
if (
isinstance(completion_response, LiteLLMBatch)
and kwargs.get("call_type") == CallTypes.aretrieve_batch.value
and not batch_cost_is_final(completion_response)
):
verbose_proxy_logger.debug(
"Cost tracking deferred for batch %s still in status %s",
completion_response.id,
completion_response.status,
)
await _release_budget_reservation(budget_reservation=budget_reservation)
return
user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None))
team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None))
org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None))
@ -285,7 +299,7 @@ class _ProxyDBLogger(CustomLogger):
call_type=call_type,
):
## UPDATE DATABASE
await _update_database_and_spend_counters(
charged: Final = await _update_database_and_spend_counters(
proxy_logging_obj=proxy_logging_obj,
increment_spend_counters=increment_spend_counters,
user_api_key=user_api_key,
@ -302,6 +316,8 @@ class _ProxyDBLogger(CustomLogger):
request_tags=tags,
model_access_groups=model_access_groups,
)
if not charged:
return
# update cache (fire-and-forget for backward compat:
# cached object fields, soft budget alerts, etc.)
@ -578,9 +594,9 @@ async def _update_database_and_spend_counters(
budget_reservation: dict | None,
request_tags: list[str] | None = None,
model_access_groups: Sequence[str] | None = None,
) -> None:
) -> bool:
try:
await proxy_logging_obj.db_spend_update_writer.update_database(
charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database(
token=user_api_key,
response_cost=response_cost,
user_id=user_id,
@ -605,6 +621,9 @@ async def _update_database_and_spend_counters(
"Failed to invalidate budget reservation counters after release failed"
)
raise
if not charged:
await _release_budget_reservation(budget_reservation=budget_reservation)
return False
try:
await increment_spend_counters(
@ -630,6 +649,7 @@ async def _update_database_and_spend_counters(
finally:
budget_reservation["finalized"] = True
raise
return True
async def _release_budget_reservation(budget_reservation: dict | None) -> None:

40
litellm/proxy/litellm_pre_call_utils.py
View file

@ -2882,6 +2882,28 @@ def _add_guardrails_from_policies_in_metadata(
)
def add_guardrails_from_auth_metadata(
user_api_key_dict: UserAPIKeyAuth,
data: dict, # mutable-ok: writes guardrails into the live request dict, same contract as the helpers it wraps
metadata_variable_name: str,
) -> None:
"""Resolve key, team, and project guardrails, direct and via policies, onto the request metadata."""
_add_guardrails_from_key_or_team_metadata(
key_metadata=user_api_key_dict.metadata,
team_metadata=user_api_key_dict.team_metadata,
project_metadata=user_api_key_dict.project_metadata,
data=data,
metadata_variable_name=metadata_variable_name,
)
_add_guardrails_from_policies_in_metadata(
key_metadata=user_api_key_dict.metadata,
team_metadata=user_api_key_dict.team_metadata,
project_metadata=user_api_key_dict.project_metadata,
data=data,
metadata_variable_name=metadata_variable_name,
)
async def move_guardrails_to_metadata(
data: dict,
_metadata_variable_name: str,
@ -2914,22 +2936,8 @@ async def move_guardrails_to_metadata(
data.pop("policies", None)
return
# Check key/team/project-level guardrails
_add_guardrails_from_key_or_team_metadata(
key_metadata=user_api_key_dict.metadata,
team_metadata=user_api_key_dict.team_metadata,
project_metadata=project_metadata,
data=data,
metadata_variable_name=_metadata_variable_name,
)
#########################################################################################
# Add guardrails from policies attached to key/team/project metadata
#########################################################################################
_add_guardrails_from_policies_in_metadata(
key_metadata=user_api_key_dict.metadata,
team_metadata=user_api_key_dict.team_metadata,
project_metadata=project_metadata,
add_guardrails_from_auth_metadata(
user_api_key_dict=user_api_key_dict,
data=data,
metadata_variable_name=_metadata_variable_name,
)

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

@ -2840,6 +2840,11 @@ async def update_key_fn(
"""
Update an existing API key's parameters.
The body is a merge patch: a field left out keeps its stored value, and on the key's own columns
an explicit null clears it. The metadata-backed fields below are the exception, merging into the
stored metadata instead: passing one as null leaves it unchanged, while `metadata` itself
replaces the stored metadata wholesale.
Parameters:
- key: Optional[str] - The key to update. Either key or key_alias must be provided.
- key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases)

8
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -15,6 +15,7 @@ from typing import (
runtime_checkable,
)
from litellm.batches.batch_utils import batch_cost_is_final
from litellm.proxy._types import ProxyException
from litellm.repositories.table_repositories import (
ManagedFileRepository,
@ -1357,12 +1358,7 @@ def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool:
enumerated the batch and none succeeded. A zero or unknown total means counts
are unreported, so stay eligible and let the next poller pass revisit it. (#37713)
"""
if response.output_file_id is not None:
return True
request_counts = response.request_counts
if request_counts is None:
return False
return request_counts.total > 0 and request_counts.completed == 0
return batch_cost_is_final(response)
async def update_batch_in_database(

10
litellm/proxy/utils.py
View file

@ -152,7 +152,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
from litellm.proxy.hooks.sensitive_data_routing import (
_PROXY_SensitiveDataRoutingHandler,
)
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
from litellm.repositories.budget_repository import BudgetRepository
@ -924,7 +924,13 @@ class ProxyLogging:
"incoming_bearer_token": kwargs.get("incoming_bearer_token"),
"metadata": {"headers": kwargs.get("headers") or {}},
}
user_api_key_auth: Final = kwargs.get("user_api_key_auth")
if isinstance(user_api_key_auth, UserAPIKeyAuth):
add_guardrails_from_auth_metadata(
user_api_key_dict=user_api_key_auth,
data=synthetic_data,
metadata_variable_name="metadata",
)
return synthetic_data
def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> MCPPreCallResponseObject | None:

3
litellm/repositories/model_repository.py
View file

@ -4,7 +4,6 @@ Model repository for database operations on LiteLLM_ProxyModelTable.
import json
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol
from litellm.models.model import LiteLLM_ProxyModelTable
@ -109,7 +108,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
async def find_all_except(self, model_id: str) -> Sequence[LiteLLM_ProxyModelTable]:
"""Find every model except the row currently being updated."""
records: Final = await self.table.find_many(
where=MappingProxyType({"model_id": MappingProxyType({"not": model_id})})
where={"model_id": {"not": model_id}} # mutable-ok: Prisma requires plain dicts for query serialization
)
return tuple(self._to_model_list(records))

25
litellm/responses/utils.py
View file

@ -25,9 +25,15 @@ from litellm.types.utils import (
PromptTokensDetailsWrapper,
SpecialEnums,
Usage,
text_tokens_without_nested_reasoning,
)
def _output_token_detail(details: object, field: str) -> int | None:
value: Final = getattr(details, field, None)
return value if isinstance(value, int) else None
def _is_object_sequence(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a list is a Sequence of anything
return isinstance(value, list)
@ -1137,11 +1143,22 @@ class ResponseAPILoggingUtils:
response_api_usage, "output_tokens_details", None
)
if output_tokens_details:
reasoning_tokens: Final = _output_token_detail(output_tokens_details, "reasoning_tokens")
image_tokens: Final = _output_token_detail(output_tokens_details, "image_tokens")
audio_tokens: Final = _output_token_detail(output_tokens_details, "audio_tokens")
reported_text_tokens: Final = _output_token_detail(output_tokens_details, "text_tokens")
completion_tokens_details = CompletionTokensDetailsWrapper(
reasoning_tokens=getattr(output_tokens_details, "reasoning_tokens", None),
image_tokens=getattr(output_tokens_details, "image_tokens", None),
text_tokens=getattr(output_tokens_details, "text_tokens", None),
audio_tokens=getattr(output_tokens_details, "audio_tokens", None),
reasoning_tokens=reasoning_tokens,
image_tokens=image_tokens,
text_tokens=None
if reported_text_tokens is None
else text_tokens_without_nested_reasoning(
completion_tokens=completion_tokens,
text_tokens=reported_text_tokens,
reasoning_tokens=reasoning_tokens or 0,
other_modality_tokens=(audio_tokens or 0) + (image_tokens or 0),
),
audio_tokens=audio_tokens,
)
extra_usage_fields: Final = {

69
litellm/router.py
View file

@ -403,6 +403,10 @@ _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else ()
def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]:
"""
Realtime client-secret requests carry the model inside ``session`` as well, and the caller's copy of it still
@ -4449,7 +4453,7 @@ class Router:
self.fail_calls[model_name] += 1
raise e
async def aspeech(self, model: str, input: str, voice: str, **kwargs):
async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs):
"""
Example Usage:
@ -4501,7 +4505,7 @@ class Router:
)
raise e
async def _aspeech(self, model: str, input: str, voice: str, **kwargs):
async def _aspeech(self, model: str, input: str, voice: str | None = None, **kwargs):
model_name: Final = model
try:
verbose_router_logger.debug("Inside _aspeech()- model: %s; kwargs: %s", model, kwargs)
@ -4525,7 +4529,7 @@ class Router:
**{
**data,
"input": input,
"voice": voice,
"voice": data.get("voice") if voice is None else voice,
"client": model_client,
**kwargs,
}
@ -7458,6 +7462,21 @@ class Router:
Context_Policy_Fallbacks={content_policy_fallbacks}",
)
@staticmethod
def _deployment_ids_to_skip_on_retry(exception: Exception, already_skipped: object) -> tuple[str, ...]:
failed_deployment_id: Final[str | None] = getattr(exception, "failed_deployment_id", None)
status_code: Final = getattr(exception, "status_code", None)
if not failed_deployment_id or not isinstance(status_code, int):
return ()
if litellm._should_retry(status_code): # pyright: ignore[reportPrivateUsage] # as in should_retry_this_error
return ()
already_skipped_ids: Final = _as_retry_skipped_deployment_ids(already_skipped)
skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id))))
verbose_router_logger.debug(
"Retry skips deployments that already answered %s to this request: %s", status_code, skipped
)
return skipped
@tracer.wrap()
async def async_function_with_retries(self, *args, **kwargs):
verbose_router_logger.debug("Inside async function with retries.")
@ -7553,6 +7572,12 @@ class Router:
## LOGGING
if num_retries > 0:
kwargs = self.log_retry(kwargs=kwargs, e=original_exception)
first_skipped_ids: Final = self._deployment_ids_to_skip_on_retry(
exception=original_exception,
already_skipped=kwargs.get("_retry_skipped_deployment_ids"),
)
if first_skipped_ids:
kwargs["_retry_skipped_deployment_ids"] = first_skipped_ids # rebind-ok: the next attempt reads it
else:
raise
@ -7622,6 +7647,12 @@ class Router:
except Exception:
raise e
skipped_ids = self._deployment_ids_to_skip_on_retry(
exception=e,
already_skipped=kwargs.get("_retry_skipped_deployment_ids"),
)
if skipped_ids:
kwargs["_retry_skipped_deployment_ids"] = skipped_ids # rebind-ok: the next attempt reads it
_timeout = self._time_to_sleep_before_retry(
e=e,
remaining_retries=remaining_retries,
@ -12454,7 +12485,7 @@ class Router:
## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2)
_target_order: Final = (request_kwargs or {}).pop("_target_order", None)
healthy_deployments = litellm.utils._get_order_filtered_deployments(
healthy_deployments = litellm.utils.get_order_filtered_deployments(
cast(list[dict], healthy_deployments), target_order=_target_order
)
@ -12462,11 +12493,24 @@ class Router:
## this request via weighted-failover. Always honored, regardless of the
## router-level flag, so a stale exclusion key on kwargs cannot escape.
_excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None)
healthy_deployments = litellm.utils._get_excluded_filtered_deployments(
healthy_deployments = litellm.utils.get_excluded_filtered_deployments(
cast(list[dict], healthy_deployments),
excluded_deployment_ids=_excluded_deployment_ids,
)
## RETRY SKIP ## -> drop deployments that already refused this request with a
## non-retryable status, unless that leaves nothing, so the caller still gets
## the provider's own error instead of a no-deployments error.
_retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids(
request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None
)
healthy_deployments = (
litellm.utils.get_excluded_filtered_deployments(
healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids
)
or healthy_deployments
)
if len(healthy_deployments) == 0:
exception: Final = await async_raise_no_deployment_exception(
litellm_router_instance=self,
@ -13359,7 +13403,7 @@ class Router:
## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2)
_target_order: Final = (request_kwargs or {}).pop("_target_order", None)
healthy_deployments = litellm.utils._get_order_filtered_deployments(
healthy_deployments = litellm.utils.get_order_filtered_deployments(
healthy_deployments, target_order=_target_order
)
@ -13367,11 +13411,22 @@ class Router:
## this request via weighted-failover. See async counterpart in
## async_get_healthy_deployments for details.
_excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None)
healthy_deployments = litellm.utils._get_excluded_filtered_deployments(
healthy_deployments = litellm.utils.get_excluded_filtered_deployments(
healthy_deployments,
excluded_deployment_ids=_excluded_deployment_ids,
)
## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments.
_retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids(
request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None
)
healthy_deployments = (
litellm.utils.get_excluded_filtered_deployments(
healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids
)
or healthy_deployments
)
if len(healthy_deployments) == 0:
model_ids = self.get_model_ids(model_name=model)
_cooldown_time = self.cooldown_cache.get_min_cooldown(

5
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -2835,8 +2835,9 @@ class ComplexityRouter(CustomLogger):
where the prompt never arrives as messages.
Probed on a COPY of request_kwargs because the owner pops routing bookkeeping off the
dict it is handed (`_target_order`, `_excluded_deployment_ids`), and this is a
speculative question about a model that may never be picked.
dict it is handed (`_target_order`, `_excluded_deployment_ids`,
`_retry_skipped_deployment_ids`), and this is a speculative question about a model
that may never be picked.
Every way the owner says "nothing here can serve this" is a negative verdict: no healthy
deployment for the group at all (BadRequestError, which ContextWindowExceededError

4
litellm/router_utils/reasoning_effort_capability.py
View file

@ -10,8 +10,8 @@ opt-in. none is opt-out everywhere except the azure gpt-5 family, whose config r
UnsupportedParamsError without an explicit true.
xhigh is gated on the request path by the openai and azure gpt-5 configs. max is not gated there at
all: every entry carrying supports_max_reasoning_effort is Claude-family, and
anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort
all: outside the gpt-6-astra rows every entry carrying supports_max_reasoning_effort is Claude-family,
and anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort
path maps any level to a thinking budget. Making max opt-in is a deliberate trade, then, since an
explicit flag is the only signal that the tier is a real one rather than litellm rounding the level
to a budget, and a missing flag costs advisory metadata rather than a rejected request.

3
litellm/rust_bridge/chat_completions.py
View file

@ -247,8 +247,7 @@ def rust_chat_completions_accepts(
return False
if stream:
return False
request_override: Final = litellm_params.get("rust") if litellm_params is not None else None
if not rust_enabled(request_override=request_override if isinstance(request_override, bool) else None):
if not rust_enabled():
return False
if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params):
verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path")

33
litellm/rust_bridge/configuration.py
View file

@ -1,13 +1,11 @@
from __future__ import annotations
import os
import warnings
from typing import Final
DEFAULT_RUST_ENABLED: Final = False
_TRUE_ENV_VALUES: Final = frozenset({"1", "true", "yes", "on"})
_GLOBAL_ENV_NAME: Final = "LITELLM_RUST"
_LEGACY_OCR_ENV_NAME: Final = "LITELLM_USE_RUST_OCR"
class _RustConfiguration:
@ -26,49 +24,24 @@ def _parse_env_bool(value: str | None) -> bool | None:
def resolve_rust_enabled(
*,
request_override: bool | None,
process_override: bool | None,
environment_override: bool | None,
legacy_environment_override: bool | None = None,
release_default: bool = DEFAULT_RUST_ENABLED,
) -> bool:
if request_override is not None:
return request_override
if process_override is not None:
return process_override
if environment_override is not None:
return environment_override
if legacy_environment_override is not None:
return legacy_environment_override
return release_default
def rust_enabled(*, request_override: bool | None = None) -> bool:
if request_override is not None:
return request_override
process_override: Final = _CONFIGURATION.override
if process_override is not None:
return process_override
global_override: Final = _parse_env_bool(os.getenv(_GLOBAL_ENV_NAME))
legacy_override: Final = None if global_override is not None else _parse_env_bool(os.getenv(_LEGACY_OCR_ENV_NAME))
if legacy_override is not None:
warnings.warn(
f"{_LEGACY_OCR_ENV_NAME} is deprecated; use {_GLOBAL_ENV_NAME} instead",
DeprecationWarning,
stacklevel=2,
)
def rust_enabled() -> bool:
return resolve_rust_enabled(
request_override=None,
process_override=None,
environment_override=global_override,
legacy_environment_override=legacy_override,
process_override=_CONFIGURATION.override,
environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)),
)
def rust_ocr_enabled(*, request_override: bool | None = None) -> bool:
return rust_enabled(request_override=request_override)
def reset_rust_configuration() -> None:
_CONFIGURATION.override = None

46
litellm/rust_bridge/ocr.py
View file

@ -7,12 +7,9 @@ from typing import Final, Protocol, cast # noqa: TID251 # native extension exp
import httpx
from litellm.rust_bridge import configuration as _configuration
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds
rust_ocr_enabled = _configuration.rust_ocr_enabled
rust = _configuration.rust
class RustOcr(Protocol):
def __call__(
@ -44,49 +41,24 @@ class RustAocr(Protocol):
raise NotImplementedError
class _Unset:
pass
def _as_ocr(value: object) -> RustOcr | None:
return cast(RustOcr, value) if callable(value) else None
_UNSET: Final[_Unset] = _Unset()
def _as_aocr(value: object) -> RustAocr | None:
return cast(RustAocr, value) if callable(value) else None
_rust_ocr_impl: RustOcr | None = None
_rust_aocr_impl: RustAocr | None = None
def set_rust_ocr(
*,
ocr: RustOcr | None | _Unset = _UNSET,
aocr: RustAocr | None | _Unset = _UNSET,
) -> None:
global _rust_ocr_impl, _rust_aocr_impl
if not isinstance(ocr, _Unset):
_rust_ocr_impl = ocr
if not isinstance(aocr, _Unset):
_rust_aocr_impl = aocr
_OCR: Final = NativeBinding("ocr", validate=_as_ocr)
_AOCR: Final = NativeBinding("aocr", validate=_as_aocr)
def load_rust_ocr() -> RustOcr | None:
if _rust_ocr_impl is not None:
return _rust_ocr_impl
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
return cast(RustOcr, native_bridge.ocr)
return _OCR.load()
def load_rust_aocr() -> RustAocr | None:
if _rust_aocr_impl is not None:
return _rust_aocr_impl
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
return cast(RustAocr, getattr(native_bridge, "aocr", None))
return _AOCR.load()
def ocr(

1
litellm/rust_bridge/transcription.py
View file

@ -56,7 +56,6 @@ _STATE: Final = _RustTranscriptionState()
def configure_rust_transcription(
enabled: bool = True,
*,
transcription: RustTranscription | None | _Unset = _UNSET,
atranscription: RustAtranscription | None | _Unset = _UNSET,

1
litellm/types/router.py
View file

@ -307,7 +307,6 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
"""
custom_llm_provider: str | None = None
rust: bool | None = None
tpm: int | None = None
rpm: int | None = None
itpm: int | None = None

12
litellm/types/utils.py
View file

@ -1625,6 +1625,17 @@ class Choices(SafeAttributeModel, OpenAIObject):
setattr(self, key, value)
def text_tokens_without_nested_reasoning(
completion_tokens: int,
text_tokens: int,
reasoning_tokens: int,
other_modality_tokens: int,
) -> int:
reported_total: Final = text_tokens + reasoning_tokens + other_modality_tokens
nested_reasoning_tokens: Final = min(reasoning_tokens, text_tokens, max(reported_total - completion_tokens, 0))
return text_tokens - nested_reasoning_tokens
class CompletionTokensDetailsWrapper(CompletionTokensDetails): # wrapper for older openai versions
text_tokens: int | None = None
"""Text tokens generated by the model."""
@ -3053,6 +3064,7 @@ class StandardLoggingPayloadErrorInformation(TypedDict, total=False):
llm_provider: str | None
traceback: str | None
error_message: str | None
error_provider_request_id: ReadOnly[str | None]
# error_rate_limit_category:
# For 429 / rate-limit errors, the source of the rate limit. One of the
# string values defined by `litellm.exceptions.RateLimitErrorCategory`

20
litellm/utils.py
View file

@ -4889,7 +4889,7 @@ def _get_deployment_order(deployment: dict | Any) -> int | None:
return order
def _get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list:
def get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list:
if target_order is not None:
return [d for d in healthy_deployments if _get_deployment_order(d) == target_order]
@ -4908,7 +4908,7 @@ def _get_order_filtered_deployments(healthy_deployments: list[dict], target_orde
return healthy_deployments
def _get_excluded_filtered_deployments(
def get_excluded_filtered_deployments(
healthy_deployments: list[dict],
excluded_deployment_ids: Iterable[str] | None = None,
) -> list:
@ -4919,10 +4919,12 @@ def _get_excluded_filtered_deployments(
across the remaining deployments in the same model group after one of them
has failed.
If the filter would leave no deployments, an empty list is returned so the
caller raises its usual no-deployments error and the weighted-failover
helper falls through to the cross-group fallback path. Returning the
original unfiltered list here would re-include the just-failed deployment.
If the filter would leave no deployments, an empty list is returned and the
caller decides what that means. Weighted failover lets it raise the usual
no-deployments error and fall through to the cross-group fallback path; the
retry skip in `async_get_healthy_deployments` deliberately falls back to the
unfiltered list, so a request every deployment refused still comes back with
the provider's own error rather than a no-deployments one.
"""
if not excluded_deployment_ids:
return healthy_deployments
@ -9451,6 +9453,12 @@ class ProviderConfigManager:
)
return MinimaxTextToSpeechConfig()
elif litellm.LlmProviders.MISTRAL == provider:
from litellm.llms.mistral.audio_speech.transformation import (
MistralTextToSpeechConfig,
)
return MistralTextToSpeechConfig()
elif litellm.LlmProviders.AWS_POLLY == provider:
from litellm.llms.aws_polly.text_to_speech.transformation import (
AWSPollyTextToSpeechConfig,

57
model_prices_and_context_window.json
View file

@ -3485,6 +3485,55 @@
"supports_response_schema": true,
"supports_tool_choice": true
},
"azure_ai/gpt-6-astra": {
"cache_creation_input_token_cost": 1.25e-05,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-05,
"cache_read_input_token_cost": 1e-06,
"cache_read_input_token_cost_above_272k_tokens": 2e-06,
"input_cost_per_token": 1e-05,
"input_cost_per_token_above_272k_tokens": 2e-05,
"litellm_provider": "azure_ai",
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 5e-05,
"output_cost_per_token_above_272k_tokens": 7.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"source": "https://ai.azure.com/catalog/models/gpt-6-astra",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": false,
"supports_minimal_reasoning_effort": false,
"supports_native_streaming": true,
"supports_none_reasoning_effort": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_cache_breakpoint": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true,
"supports_xhigh_reasoning_effort": true
},
"azure_ai/gpt-5.5": {
"deprecation_date": "2027-10-26",
"cache_read_input_token_cost": 5e-07,
@ -7189,7 +7238,7 @@
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": true,
"supports_max_reasoning_effort": false,
"supports_minimal_reasoning_effort": false,
"supports_native_streaming": true,
"supports_none_reasoning_effort": true,
@ -7455,7 +7504,7 @@
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": true,
"supports_max_reasoning_effort": false,
"supports_minimal_reasoning_effort": false,
"supports_native_streaming": true,
"supports_none_reasoning_effort": true,
@ -34898,9 +34947,9 @@
"supports_audio_input": true
},
"mistral/voxtral-mini-tts-latest": {
"input_cost_per_character": 1.6e-05,
"litellm_provider": "mistral",
"mode": "audio_speech",
"output_cost_per_character": 1.6e-05,
"source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03",
"supported_endpoints": [
"/v1/audio/speech"
@ -56467,9 +56516,9 @@
"supports_audio_input": true
},
"mistral/voxtral-mini-tts-2603": {
"input_cost_per_character": 1.6e-05,
"litellm_provider": "mistral",
"mode": "audio_speech",
"output_cost_per_character": 1.6e-05,
"source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03",
"supported_endpoints": [
"/v1/audio/speech"

23
scripts/pre_commit_lint.sh
View file

@ -12,6 +12,10 @@
# - litellm/ Python -> `make lint` (test-linting.yml's lint job)
# - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step)
# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests)
# - tests/ Python, ruff-tests.toml, test-quality-budget.json, scripts/check_test_quality.py,
# scripts/test_quality_gate.py
# -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's
# test-tree ruff and test-quality budget steps)
# - dashboard -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint)
# - proxy/types -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift (check-ui-api-types.yml)
#
@ -88,15 +92,14 @@ existing_files() {
litellm_py_pattern='^litellm/.*\.py$'
e2e_py_pattern='^tests/e2e/.*\.py$'
test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|test-quality-budget\.json|scripts/(check_test_quality|test_quality_gate)\.py)$'
spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$'
ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$'
ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$'
# CI's lint job (test-linting.yml) only inspects litellm/, so a tests-only or
# scripts-only commit can't turn it red; scope the trigger there to skip the slow
# make lint when it couldn't catch anything.
litellm_py_files=$(scope_match "$litellm_py_pattern")
e2e_py_files=$(scope_match "$e2e_py_pattern")
test_tree_files=$(scope_match "$test_tree_pattern")
# ruff format (and CI's format step) skip enterprise; the rest of make lint covers it.
fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files)
# check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types
@ -136,6 +139,7 @@ if [ -n "$staged" ]; then
}
warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files"
warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files"
warn_skipped "test-tree lint (ruff-tests.toml + test-quality budget)" "$test_tree_pattern" "$test_tree_files"
warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_changed"
warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files"
fi
@ -288,6 +292,15 @@ if [ -n "$spec_files" ]; then
set +m
fi
if [ -n "$test_tree_files" ] && [ -z "$litellm_py_files" ]; then
echo "check: linting the test tree (ruff check --config ruff-tests.toml tests)"
uv run --no-sync ruff check --config ruff-tests.toml tests \
|| { echo "✗ Test-tree ruff failed. Fix the errors above, then re-run make check." >&2; status=1; }
echo "check: checking the test-quality budget (make lint-test-quality)"
make lint-test-quality \
|| { echo "✗ Test-quality budget failed. Fix the errors above, then re-run make check." >&2; status=1; }
fi
if [ -n "${python_pid:-}" ]; then
wait "$python_pid" || status=1
cat "$python_log"; rm -f "$python_log"
@ -313,10 +326,12 @@ summary_item() {
echo "check: summary"
summary_item "Python lint (make lint)" "$litellm_py_files" "no litellm/ Python files in scope"
summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_files" "no tests/e2e Python files in scope"
summary_item "test-tree lint (ruff-tests.toml + test-quality budget)" "$test_tree_files" \
"no tests/ Python files or test-tree lint inputs in scope"
summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope"
summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope"
if [ -z "$litellm_py_files$e2e_py_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then
if [ -z "$litellm_py_files$e2e_py_files$test_tree_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then
echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2
printf '%s\n' "$scope" | sed 's/^/ /' >&2
echo " A pass here is a no-op, not a lint verdict." >&2

86
scripts/test_quality_gate.py
View file

@ -10,17 +10,12 @@ base.
Every rule is seeded at exactly its count on the day the gate landed, so the
suite's existing debt is grandfathered and any net-new violation trips the gate
immediately. ``--update`` ratchets a limit down by the violations this branch
fixed relative to its branch point (the merge-base), so the ceilings only ever
fall. Base counts are measured with the *current* checker, so a rule introduced
on this branch is counted at the base too and ratchets like every other one.
Only ever falling is not the same as always falling, so the gate enforces the
second half: a branch that clears violations and leaves the ceiling above its
new count fails, naming the rules and telling the author to run
``make lint-budget-update``. Without that, a removed violation could come back
later under a ceiling nobody lowered. Drift already in the base is never
blamed, so this fires only on the branch that did the clearing.
immediately. ``--update`` ratchets a limit down by the violations fixed relative
to ``--base``, so the ceilings only ever fall. Base counts are measured with the
*current* checker, so a rule introduced on this branch is counted at the base too
and ratchets like every other one. The ratchet runs as a scheduled automation
against litellm_internal_staging, not on PR branches, so concurrent PRs never
race to edit the same limit.
The deliberate difference from its sibling: this gate has no headroom anywhere.
Type discipline seeded LIT010/LIT011 at 1.5x to leave room for an in-flight
@ -34,13 +29,14 @@ import argparse
import json
import re
import shutil
import signal
import subprocess
import sys
import tempfile
from collections import Counter
from collections.abc import Mapping, Sequence
from pathlib import Path
from types import MappingProxyType
from types import FrameType, MappingProxyType
from typing import Final, NamedTuple
REPO_ROOT: Final = Path(__file__).resolve().parent.parent
@ -48,6 +44,7 @@ CHECKER: Final = REPO_ROOT / "scripts" / "check_test_quality.py"
BUDGET_PATH: Final = REPO_ROOT / "test-quality-budget.json"
TARGET: Final = "tests"
DEFAULT_BASE: Final = "origin/litellm_internal_staging"
TERMINATION_SIGNALS: Final = (signal.SIGTERM, signal.SIGHUP)
_HUNK: Final = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@", re.MULTILINE)
_FILE_HEADER: Final = re.compile(r"^\+\+\+ b/(.+)$", re.MULTILINE)
@ -116,22 +113,33 @@ def count_by_rule(violations: Sequence[Violation]) -> Mapping[str, int]:
return MappingProxyType(dict(Counter(v.code for v in violations)))
def base_counts(ref: str) -> Mapping[str, int]:
def _exit_on_termination(signum: int, _frame: FrameType | None) -> None:
raise SystemExit(128 + signum)
def _install_termination_handlers() -> None:
for termination in TERMINATION_SIGNALS:
if signal.getsignal(termination) == signal.SIG_DFL:
signal.signal(termination, _exit_on_termination)
def base_counts(ref: str, repo_root: Path = REPO_ROOT, checker: Path = CHECKER) -> Mapping[str, int]:
"""Rule counts at `ref`, measured with the *current* rule logic rather than
whatever the checker looked like at that commit."""
_install_termination_handlers()
parent: Final = Path(tempfile.mkdtemp(prefix="tq_base_"))
worktree: Final = parent / "wt"
try:
_run(["git", "worktree", "add", "--detach", str(worktree), ref])
_run(["git", "worktree", "add", "--detach", str(worktree), ref], cwd=repo_root)
(worktree / "scripts").mkdir(parents=True, exist_ok=True)
checker: Final = worktree / "scripts" / "check_test_quality.py"
shutil.copy(CHECKER, checker)
return count_by_rule(_check(worktree, checker))
base_checker: Final = worktree / "scripts" / "check_test_quality.py"
shutil.copy(checker, base_checker)
return count_by_rule(_check(worktree, base_checker))
finally:
# Teardown must never raise, or it masks the real error when the body failed.
subprocess.run(
["git", "worktree", "remove", "--force", str(worktree)],
cwd=REPO_ROOT, capture_output=True, text=True,
cwd=repo_root, capture_output=True, text=True,
)
shutil.rmtree(parent, ignore_errors=True)
@ -144,21 +152,6 @@ def over_ceiling(head: Mapping[str, int], budget: Mapping[str, Mapping[str, int]
)
def unratcheted(
head: Mapping[str, int],
base: Mapping[str, int],
budget: Mapping[str, Mapping[str, int]],
) -> tuple[Breach, ...]:
"""Rules this branch cleared without lowering the ceiling behind them. Requires
both `head < base`, so drift already in the base is never blamed on this change,
and `head < limit`, so a ceiling already at the count is left alone."""
return tuple(sorted(
Breach(rule, head.get(rule, 0), spec["limit"], head.get(rule, 0) - base.get(rule, 0))
for rule, spec in budget.items()
if head.get(rule, 0) < base.get(rule, 0) and head.get(rule, 0) < spec["limit"]
))
def evaluate(
head: Mapping[str, int],
base: Mapping[str, int],
@ -198,38 +191,15 @@ def introduced(
return tuple(v for v in violations if v.line in changed.get(v.file, frozenset()))
def touches_measured_tree(base_point: str) -> bool:
"""Whether this branch changed anything that can move a count. A branch that
touches neither the test tree nor the checker cannot have cleared a violation,
so the base scan is skipped and the gate stays cheap on the common change."""
changed: Final = _run(
["git", "diff", "--name-only", base_point, "--", TARGET, str(CHECKER.relative_to(REPO_ROOT))]
)
return bool(changed.strip())
def cmd_check(base: str) -> None:
budget: Final = json.loads(BUDGET_PATH.read_text())
head: Final = head_violations()
head_counts: Final = count_by_rule(head)
base_point: Final = resolve_base_point(base)
if not over_ceiling(head_counts, budget) and not touches_measured_tree(base_point):
if not over_ceiling(head_counts, budget):
print(f"OK: every TQ rule is within its test-suite ceiling (base {base})")
return
base_point: Final = resolve_base_point(base)
base_at_point: Final = base_counts(base_point)
stale: Final = unratcheted(head_counts, base_at_point, budget)
if stale:
print(f"FAIL: TQ-rule limits were left above the count this branch reached (base {base}):")
for breach in stale:
print(
f" {breach.rule}: this branch cleared {-breach.added} down to {breach.total}, "
f"but the limit is still {breach.cap}"
)
print(
"Run `make lint-budget-update` and commit the lowered limits, so the "
"violations you cleared cannot come back under a ceiling nobody moved."
)
raise SystemExit(1)
breaches: Final = evaluate(head_counts, base_at_point, budget)
if not breaches:
print(f"OK: every TQ rule is within its test-suite ceiling (base {base})")

521
tests/code_coverage_tests/check_workflow_job_name_collisions.py Normal file
View file

@ -0,0 +1,521 @@
"""Catch workflow jobs that publish check runs under the same name.
A ruleset's required status check names a check run and GitHub matches it by that
name alone. When two jobs publish the same name the required context stops
mapping to the job that proves it: the commit carries two check runs under one
name and nothing says which one the ruleset required. Both being green hides the
clash completely, so the context quietly stops meaning what the ruleset intended.
One job lands in the same place when its `name:` holds no matrix value, since
every combination it runs then reports under that one name.
`.github/workflows/auto-close-duplicates.yml` shipped a job id `test` while
`.github/workflows/test-mcp.yml` already published the required `test` context,
and commit ed5761daef4ae17152446d182c860630c38b7268 carried both check runs.
This invariant has to be enforced here because CI cannot enforce it on itself.
A job publishes its `name:` when it sets one, and otherwise its job id plus the
values of the combination it runs, the way GitHub writes `build (3.12)`. A name
carrying `${{ ... }}` publishes one check run per combination the matrix
produces: `exclude` rows drop combinations before `include` rows fold into the
survivors, and each `include` row's values stay together rather than crossing
with the other rows', so two shard lists that overlap collide even though their
templates read differently. Each expression is evaluated per combination over the
pieces a job name can hold: string literals, `matrix.<key>`, `format()`, `==` and
`!=`, and the `<cond> && <a> || <b>` idiom, which is how the shards reach their
real `<shard> / Run tests` names rather than staying opaque.
Whatever the sweep cannot work out is left out of the comparison and reported
instead of guessed, because a guess that lands wrong fails a workflow GitHub
would have published perfectly well. A name still holding an expression once the
combination is filled in is usually one GitHub resolves per job, so it is one of
those: guessing that two jobs sharing such a template clash would fail workflows
over a context this sweep cannot read. The exception is a name whose leftover
expressions all read a `github.` property other than `github.job`, which one run
fills in the same way for every job in it, so those are compared against the
other jobs of their own workflow and stay out of the comparison across files,
where two workflows can run on different events. A matrix that is itself an
expression or that lists values which are not scalars, an `include` or `exclude`
row shaped the same way, a whole `strategy:` that comes from an expression, and a
call this sweep cannot follow, go in the same bucket. The cost is that a real clash hiding behind
one of them goes unseen, which leaves a merge no worse off than before this check
existed, where the opposite direction would block work that was fine.
A job calling a local reusable workflow publishes one check run per job of the
callee, named `<caller> / <callee>` and chained through however many levels of
local calls it takes, which is why a caller's name never collides with a plain
job that happens to match it. A file under `.github/workflows/` that does not
read as one workflow at all is reported rather than skipped, since skipping it
silently would hide every job it holds.
"""
import itertools
import operator
import re
import sys
from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
import yaml
from pydantic import BaseModel, Field, ValidationError
REPO_ROOT: Final = Path(__file__).resolve().parent.parent.parent
WORKFLOWS_DIR: Final = REPO_ROOT / ".github" / "workflows"
EXPRESSION: Final = re.compile(r"\$\{\{(?P<body>.*?)\}\}", re.DOTALL)
MATRIX_REF: Final = re.compile(r"^matrix\.(?P<key>[\w-]+)$")
LITERAL: Final = re.compile(r"^'(?P<text>[^']*)'$")
FORMAT_CALL: Final = re.compile(r"^format\((?P<args>.*)\)$", re.DOTALL)
COMPARISON: Final = re.compile(r"^(?P<left>.+?)\s*(?P<operator>==|!=)\s*(?P<right>.+)$", re.DOTALL)
RUN_WIDE: Final = re.compile(r"^github\.(?!job\b)[\w.]+$")
GITHUB_PLACEHOLDER: Final = re.compile(r"\{\{|\}\}|\{\d+\}")
NO_MATRIX: Final[Mapping[str, str]] = MappingProxyType({})
NO_CALLERS: Final[frozenset[str]] = frozenset()
SCALAR: Final = (str, int, float)
MATRIX_DIRECTIVES: Final = frozenset({"include", "exclude"})
LOCAL_CALL_PREFIX: Final = "./"
@dataclass(frozen=True, slots=True)
class Unreadable:
reason: str
@dataclass(frozen=True, slots=True)
class Opaque:
reason: str
@dataclass(frozen=True, slots=True)
class Names:
"""The check-run names a job publishes, beside the reasons the rest of them stay unknown."""
known: tuple[str, ...] = ()
unknown: tuple[str, ...] = ()
local: tuple[str, ...] = ()
class Job(BaseModel):
name: object = None
uses: str | None = None
strategy: object = Field(default_factory=dict)
class Workflow(BaseModel):
jobs: Mapping[str, Job] = Field(default_factory=dict)
def scalar_text(value: object) -> str:
"""A YAML scalar the way GitHub renders it, so `true` never reaches a name as `True`."""
return str(value).lower() if isinstance(value, bool) else str(value)
def parse(source: str) -> tuple[Workflow, object] | Unreadable:
"""The workflow plus its raw `on:` value, or why the file does not read as one."""
try:
parsed: Final = yaml.safe_load(source)
except yaml.YAMLError:
return Unreadable("it does not read as one YAML document")
if not isinstance(parsed, dict):
return Unreadable("its top level is not a mapping of workflow keys")
try:
return Workflow.model_validate(parsed), parsed.get(True, parsed.get("on"))
except ValidationError as error:
return Unreadable(f"{error.error_count()} of its job definitions have a shape GitHub would reject")
def events(raw_on: object) -> frozenset[str]:
if isinstance(raw_on, Mapping):
return frozenset(str(key) for key in raw_on)
if isinstance(raw_on, str):
return frozenset({raw_on})
if isinstance(raw_on, Sequence):
return frozenset(str(event) for event in raw_on)
return frozenset()
def publishes_check_runs(raw_on: object) -> bool:
"""A `workflow_call`-only workflow posts its check runs through callers, never itself."""
return events(raw_on) != frozenset({"workflow_call"})
def scalar_list(value: object) -> tuple[str, ...] | Opaque:
"""One matrix key's values, or why the combinations it produces cannot be worked out."""
if not isinstance(value, Sequence) or isinstance(value, str):
return Opaque("a matrix key holds something other than a list of values")
if any(not isinstance(item, SCALAR) for item in value):
return Opaque("a matrix key lists values that are not plain scalars")
return tuple(scalar_text(item) for item in value)
def listed_values(matrix: Mapping[str, object]) -> tuple[tuple[str, tuple[str, ...]], ...] | Opaque:
listed: Final = tuple(
(str(key), scalar_list(values)) for key, values in matrix.items() if str(key) not in MATRIX_DIRECTIVES
)
opaque: Final = next((values for _, values in listed if isinstance(values, Opaque)), None)
if opaque is not None:
return opaque
return tuple((key, values) for key, values in listed if not isinstance(values, Opaque))
def directive_rows(matrix: Mapping[str, object], directive: str) -> tuple[Mapping[str, str], ...] | Opaque:
"""One `include` or `exclude` row, or why the combinations they shape cannot be worked out."""
rows: Final = matrix.get(directive)
if rows is None:
return ()
if not isinstance(rows, Sequence) or isinstance(rows, str):
return Opaque(f"a matrix `{directive}` is itself an expression rather than a list of rows")
mappings: Final = tuple(row for row in rows if isinstance(row, Mapping))
if len(mappings) != len(rows):
return Opaque(f"a matrix `{directive}` row is not a mapping of values")
if any(not isinstance(value, SCALAR) for row in mappings for value in row.values()):
return Opaque(f"a matrix `{directive}` row holds a value that is not a plain scalar")
return tuple(MappingProxyType({str(key): scalar_text(value) for key, value in row.items()}) for row in mappings)
def drops(row: Mapping[str, str], combination: Mapping[str, str]) -> bool:
"""GitHub removes a combination that carries every value one `exclude` row names."""
return all(combination.get(key) == value for key, value in row.items())
def extends(row: Mapping[str, str], combination: Mapping[str, str]) -> bool:
"""GitHub folds an `include` row into a combination only where it overwrites no listed value."""
return all(combination[key] == value for key, value in row.items() if key in combination)
def extended(combination: Mapping[str, str], rows: Sequence[Mapping[str, str]]) -> Mapping[str, str]:
additions: Final = {key: value for row in rows if extends(row, combination) for key, value in row.items()}
return MappingProxyType({**combination, **additions})
def crossed_values(listed: Sequence[tuple[str, tuple[str, ...]]]) -> tuple[Mapping[str, str], ...]:
if not listed:
return ()
return tuple(
MappingProxyType(dict(zip((key for key, _ in listed), values)))
for values in itertools.product(*(values for _, values in listed))
)
def matrix_combinations(job: Job) -> tuple[Mapping[str, str], ...] | Opaque:
"""One mapping per job the matrix produces, `exclude` applied before `include` as GitHub does."""
if not isinstance(job.strategy, Mapping):
return Opaque("its whole `strategy` comes from an expression")
matrix: Final = job.strategy.get("matrix")
if matrix is None:
return ()
if not isinstance(matrix, Mapping):
return Opaque("the matrix itself comes from an expression")
listed: Final = listed_values(matrix)
if isinstance(listed, Opaque):
return listed
rows: Final = directive_rows(matrix, "include")
if isinstance(rows, Opaque):
return rows
dropped: Final = directive_rows(matrix, "exclude")
if isinstance(dropped, Opaque):
return dropped
kept: Final = tuple(
combination for combination in crossed_values(listed) if not any(drops(row, combination) for row in dropped)
)
standalone: Final = tuple(row for row in rows if not any(extends(row, combination) for combination in kept))
return (*(extended(combination, rows) for combination in kept), *standalone)
def scanned(state: tuple[int, bool], char: str) -> tuple[int, bool]:
depth, quoted = state
if char == "'":
return depth, not quoted
if quoted:
return depth, quoted
return depth + int(char == "(") - int(char == ")"), quoted
def split_outside(text: str, token: str) -> tuple[str, ...]:
"""`text` cut on every `token` that sits outside quotes and parentheses."""
states: Final = tuple(itertools.accumulate(text, scanned, initial=(0, False)))
cuts: Final = tuple(
index
for index in range(len(text) - len(token) + 1)
if text.startswith(token, index) and states[index] == (0, False)
)
starts: Final = (0, *(cut + len(token) for cut in cuts))
return tuple(text[start:end] for start, end in zip(starts, (*cuts, len(text))))
def formatted(template: str, arguments: Sequence[str]) -> str | None:
"""GitHub's `format()` fills `{0}`-style holes and escapes braces, so anything richer resolves to nothing."""
residue: Final = GITHUB_PLACEHOLDER.sub("", template)
if "{" in residue or "}" in residue:
return None
try:
return template.format(*arguments)
except (IndexError, KeyError, ValueError):
return None
def value_of(text: str, values: Mapping[str, str]) -> str | None:
expression: Final = text.strip()
literal: Final = LITERAL.match(expression)
if literal is not None:
return literal.group("text")
reference: Final = MATRIX_REF.match(expression)
if reference is not None:
return values.get(reference.group("key"))
call: Final = FORMAT_CALL.match(expression)
if call is None:
return None
arguments: Final = tuple(value_of(part, values) for part in split_outside(call.group("args"), ","))
resolved: Final = tuple(argument for argument in arguments if argument is not None)
if not resolved or len(resolved) != len(arguments):
return None
return formatted(resolved[0], resolved[1:])
def holds(condition: str, values: Mapping[str, str]) -> bool | None:
comparison: Final = COMPARISON.match(condition.strip())
if comparison is None:
return None
left: Final = value_of(comparison.group("left"), values)
right: Final = value_of(comparison.group("right"), values)
if left is None or right is None:
return None
return (left == right) == (comparison.group("operator") == "==")
def evaluate(body: str, values: Mapping[str, str]) -> str | None:
"""The single string this expression yields, or None when its shape is not understood."""
branches: Final = tuple(split_outside(alternative, "&&") for alternative in split_outside(body, "||"))
outcomes: Final = tuple(tuple(holds(part, values) for part in branch[:-1]) for branch in branches)
if any(outcome is None for branch in outcomes for outcome in branch):
return None
taken: Final = next((branch[-1] for branch, outcome in zip(branches, outcomes) if all(outcome)), None)
return None if taken is None else value_of(taken, values)
def resolved_span(span: re.Match[str], values: Mapping[str, str]) -> str:
substitution: Final = evaluate(span.group("body"), values)
return span.group(0) if substitution is None else substitution
def rendered(template: str, values: Mapping[str, str]) -> str:
return EXPRESSION.sub(lambda span: resolved_span(span, values), template)
def comparable(name: str) -> bool:
"""A name still holding an expression is one GitHub resolves per job, so it is nothing to compare."""
return EXPRESSION.search(name) is None
def run_wide(name: str) -> bool:
"""A name whose leftover expressions one workflow run fills in the same way for every job in it."""
return all(RUN_WIDE.match(span.group("body").strip()) is not None for span in EXPRESSION.finditer(name))
def settled(names: Sequence[str]) -> Names:
unresolved: Final = tuple(name for name in names if not comparable(name))
return Names(
tuple(name for name in names if comparable(name)),
tuple(f"its name stays `{name}`" for name in unresolved if not run_wide(name)),
tuple(name for name in unresolved if run_wide(name)),
)
def expand(template: str, job: Job) -> Names:
combinations: Final = matrix_combinations(job)
if isinstance(combinations, Opaque):
return Names((), (combinations.reason,))
over: Final = combinations or (NO_MATRIX,)
return settled(tuple(rendered(template, values) for values in over))
def suffixed(job_id: str, combination: Mapping[str, str]) -> str:
"""The name GitHub gives a job with no `name:`, its id plus the combination it runs."""
return f"{job_id} ({', '.join(combination.values())})" if combination else job_id
def published_names(job_id: str, job: Job) -> Names:
if job.name is not None:
return expand(scalar_text(job.name), job)
combinations: Final = matrix_combinations(job)
if isinstance(combinations, Opaque):
return Names((), (combinations.reason,))
suffixes: Final = tuple(dict.fromkeys(suffixed(job_id, values) for values in combinations))
return Names(suffixes or (job_id,))
def callee_path(job: Job) -> str | None:
if job.uses is None or not job.uses.startswith(LOCAL_CALL_PREFIX):
return None
return job.uses[len(LOCAL_CALL_PREFIX) :].split("@")[0]
def joined(groups: Sequence[Names]) -> Names:
return Names(
tuple(name for group in groups for name in group.known),
tuple(reason for group in groups for reason in group.unknown),
tuple(name for group in groups for name in group.local),
)
def tagged(names: Names) -> tuple[tuple[str, bool], ...]:
"""Each name a job publishes beside whether only its own workflow's run settles it."""
return (*((name, False) for name in names.known), *((name, True) for name in names.local))
def call_blocker(job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str]) -> str | None:
path: Final = callee_path(job)
if path is None:
return "it calls a reusable workflow outside this repository"
if path in callers:
return f"its call to {path} loops back on itself"
return None if path in workflows else f"it calls {path}, which this checkout does not hold"
def job_names(job_id: str, job: Job, workflows: Mapping[str, Workflow], callers: frozenset[str] = NO_CALLERS) -> Names:
prefixes: Final = published_names(job_id, job)
if job.uses is None:
return prefixes
blocker: Final = call_blocker(job, workflows, callers)
if blocker is not None:
return Names((), (*prefixes.unknown, blocker))
path: Final = callee_path(job) or ""
suffixes: Final = joined(
tuple(
job_names(callee_id, callee_job, workflows, callers | {path})
for callee_id, callee_job in workflows[path].jobs.items()
)
)
composed: Final = tuple(
(f"{prefix} / {suffix}", prefix_local or suffix_local)
for prefix, prefix_local in tagged(prefixes)
for suffix, suffix_local in tagged(suffixes)
)
return Names(
tuple(name for name, is_local in composed if not is_local),
(*prefixes.unknown, *suffixes.unknown),
tuple(name for name, is_local in composed if is_local),
)
def readable(sources: Mapping[str, str]) -> Mapping[str, tuple[Workflow, object]]:
parsed: Final = {rel: parse(source) for rel, source in sources.items()}
return MappingProxyType({rel: entry for rel, entry in parsed.items() if not isinstance(entry, Unreadable)})
def unreadable(sources: Mapping[str, str]) -> tuple[str, ...]:
parsed: Final = {rel: parse(source) for rel, source in sources.items()}
return tuple(
f"{rel} sits in the workflows directory but {entry.reason}, so none of its jobs were checked."
for rel, entry in sorted(parsed.items())
if isinstance(entry, Unreadable)
)
def scanned_jobs(sources: Mapping[str, str]) -> Iterator[tuple[str, str, Names]]:
parsed: Final = readable(sources)
workflows: Final = {rel: workflow for rel, (workflow, _) in parsed.items()}
for rel, (workflow, raw_on) in parsed.items():
if not publishes_check_runs(raw_on):
continue
for job_id, job in workflow.jobs.items():
yield rel, job_id, job_names(job_id, job, workflows)
def published(sources: Mapping[str, str]) -> Iterator[tuple[str, str]]:
for rel, job_id, names in scanned_jobs(sources):
for name in names.known:
yield name, f"{rel} job `{job_id}`"
def blind_spots(sources: Mapping[str, str]) -> tuple[str, ...]:
"""Jobs whose published names GitHub decides at run time, which no offline sweep can compare."""
return tuple(
f"{rel} job `{job_id}` publishes a name this check cannot work out because {reason}."
for rel, job_id, names in scanned_jobs(sources)
for reason in sorted(names.unknown)
)
def owners_by_name(sources: Mapping[str, str]) -> Iterator[tuple[str, tuple[str, ...]]]:
for name, pairs in itertools.groupby(sorted(published(sources)), key=operator.itemgetter(0)):
yield name, tuple(owner for _, owner in pairs)
def clash(name: str, owners: Sequence[str]) -> str | None:
"""Why one name is ambiguous, whether two jobs carry it or one job repeats it over its matrix."""
jobs: Final = tuple(dict.fromkeys(owners))
if len(jobs) > 1:
return (
f"`{name}` is published by {len(jobs)} jobs: {', '.join(jobs)}. A required status check matching "
f"that name cannot say which job proves it; give one of them a distinct `name:` or job id."
)
if len(owners) > 1:
return (
f"`{name}` is published {len(owners)} times by {jobs[0]}, once per matrix combination. A required "
f"status check matching that name cannot say which run proves it; put a matrix value in its `name:`."
)
return None
def local_published(sources: Mapping[str, str]) -> Iterator[tuple[tuple[str, str], str]]:
"""Names their own workflow's run settles, keyed by the file whose run settles them."""
for rel, job_id, names in scanned_jobs(sources):
for name in names.local:
yield (rel, name), f"job `{job_id}`"
def local_clash(rel: str, name: str, owners: Sequence[str]) -> str | None:
"""Why one workflow's own run lands several of its jobs on one check run."""
if len(owners) < 2:
return None
jobs: Final = tuple(dict.fromkeys(owners))
return (
f"`{name}` is published {len(owners)} times inside {rel}, by {', '.join(jobs)}. One run fills that "
f"expression in the same way throughout, so they all land on one check run; make the names differ."
)
def local_clashes(sources: Mapping[str, str]) -> tuple[str, ...]:
grouped: Final = itertools.groupby(sorted(local_published(sources)), key=operator.itemgetter(0))
found: Final = tuple(local_clash(rel, name, tuple(owner for _, owner in pairs)) for (rel, name), pairs in grouped)
return tuple(message for message in found if message is not None)
def collisions(sources: Mapping[str, str]) -> tuple[str, ...]:
found: Final = tuple(clash(name, owners) for name, owners in owners_by_name(sources))
return (*(message for message in found if message is not None), *local_clashes(sources))
def workflow_sources() -> Mapping[str, str]:
"""Repo-relative posix paths to text, the keys `uses: ./...` resolves against."""
return {path.relative_to(REPO_ROOT).as_posix(): path.read_text() for path in sorted(WORKFLOWS_DIR.glob("*.y*ml"))}
def report(header: str, problems: Sequence[str]) -> None:
if problems:
print(f"ERROR: {header}:\n - " + "\n - ".join(problems), file=sys.stderr)
def exit_code(sources: Mapping[str, str]) -> int:
unread: Final = unreadable(sources)
found: Final = collisions(sources)
blind: Final = blind_spots(sources)
if blind:
print("NOTE: names left out of the comparison:\n - " + "\n - ".join(blind))
report("Some workflows could not be read", unread)
report("Check-run names are not unique", found)
if unread or found:
return 1
print(f"Check-run names are unique across {len(sources)} workflows")
return 0
def main() -> int:
return exit_code(workflow_sources())
if __name__ == "__main__":
sys.exit(main())

880
tests/code_coverage_tests/test_workflow_job_name_collisions.py Normal file
View file

@ -0,0 +1,880 @@
from typing import Final
from check_workflow_job_name_collisions import (
Unreadable,
blind_spots,
callee_path,
collisions,
exit_code,
parse,
published,
unreadable,
workflow_sources,
)
REUSABLE_BASE: Final = """on:
workflow_call:
jobs:
run:
name: >-
${{ matrix.python-version == '3.12' && 'Run tests'
|| format('Run tests (Python {0})', matrix.python-version) }}
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.12", "3.13"]
"""
SHARD_CALLER: Final = """on: pull_request
jobs:
unit:
name: ${{ matrix.shard }}
uses: ./.github/workflows/base.yml
strategy:
matrix:
include:
- shard: core-utils
"""
CORRELATED_ROWS: Final = """on: pull_request
jobs:
unit:
name: ${{ matrix.shard }} on ${{ matrix.test-path }}
runs-on: ubuntu-latest
strategy:
matrix:
include:
- shard: core-utils
test-path: tests/core
- shard: proxy
test-path: tests/proxy
"""
LISTED_PLUS_ROW: Final = """on: pull_request
jobs:
unit:
name: ${{ matrix.python-version }} ${{ matrix.label }}
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.12", "3.13"]
include:
- label: fast
"""
NAMELESS_MATRIX: Final = """on: pull_request
jobs:
build:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.12", "3.13"]
"""
EXCLUDED_PAIR: Final = """on: pull_request
jobs:
unit:
name: ${{ matrix.os }}-${{ matrix.python-version }}
runs-on: ubuntu-latest
strategy:
matrix:
os: [ubuntu, macos]
python-version: ["3.12", "3.13"]
exclude:
- os: macos
python-version: "3.13"
"""
EXCLUDED_KEY: Final = """on: pull_request
jobs:
unit:
name: ${{ matrix.os }}-${{ matrix.python-version }}
runs-on: ubuntu-latest
strategy:
matrix:
os: [ubuntu, macos]
python-version: ["3.12", "3.13"]
exclude:
- os: macos
"""
BOOLEAN_MATRIX: Final = """on: pull_request
jobs:
unit:
name: cache ${{ matrix.cached }}
runs-on: ubuntu-latest
strategy:
matrix:
cached: [true, false]
"""
UNFILLABLE_FORMAT: Final = """on: pull_request
jobs:
unit:
name: ${{ format('{0} {1}', matrix.shard) }}
runs-on: ubuntu-latest
strategy:
matrix:
shard: [core]
"""
def test_every_workflow_in_the_repo_publishes_a_unique_check_run_name() -> None:
assert collisions(workflow_sources()) == ()
def test_every_workflow_in_the_repo_parses_into_jobs() -> None:
unparsed: Final = tuple(
rel
for rel, source in workflow_sources().items()
if isinstance(entry := parse(source), Unreadable) or not entry[0].jobs
)
assert unparsed == ()
def test_every_local_reusable_call_in_the_repo_resolves_to_a_workflow() -> None:
sources: Final = workflow_sources()
parsed: Final = tuple(entry for text in sources.values() if not isinstance(entry := parse(text), Unreadable))
unresolved: Final = tuple(
job.uses
for workflow, _ in parsed
for job in workflow.jobs.values()
if (callee := callee_path(job)) is not None and callee not in sources
)
assert unresolved == ()
def test_two_jobs_falling_back_to_the_same_job_id_collide() -> None:
sources: Final = {
"a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n",
"b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n",
}
found: Final = collisions(sources)
assert len(found) == 1
assert "`test` is published by 2 jobs" in found[0]
assert "a.yml job `test`" in found[0] and "b.yml job `test`" in found[0]
def test_an_explicit_name_overrides_the_job_id_and_clears_the_collision() -> None:
sources: Final = {
"a.yml": "on: pull_request\njobs:\n test:\n name: Sweep tests\n runs-on: ubuntu-latest\n",
"b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n",
}
assert collisions(sources) == ()
def test_an_explicit_name_matching_another_job_id_collides() -> None:
sources: Final = {
"a.yml": "on: pull_request\njobs:\n sweep:\n name: test\n runs-on: ubuntu-latest\n",
"b.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n",
}
found: Final = collisions(sources)
assert len(found) == 1
assert "`test` is published by 2 jobs" in found[0]
def test_two_callers_of_one_reusable_workflow_collide_on_a_shared_matrix_value() -> None:
base: Final = "on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n"
caller: Final = (
"on: pull_request\n"
"jobs:\n"
" {job}:\n"
" name: ${{{{ matrix.shard }}}}\n"
" uses: ./.github/workflows/base.yml\n"
" strategy:\n"
" matrix:\n"
" include:\n"
" - shard: {shard}\n"
)
sources: Final = {
".github/workflows/base.yml": base,
"unit.yml": caller.format(job="unit", shard="proxy-auth"),
"proxy-db.yml": caller.format(job="proxy-db", shard="proxy-auth"),
}
found: Final = collisions(sources)
assert len(found) == 1
assert "`proxy-auth / Run tests` is published by 2 jobs" in found[0]
def test_distinct_matrix_values_through_one_reusable_workflow_do_not_collide() -> None:
base: Final = "on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n"
caller: Final = (
"on: pull_request\n"
"jobs:\n"
" {job}:\n"
" name: ${{{{ matrix.shard }}}}\n"
" uses: ./.github/workflows/base.yml\n"
" strategy:\n"
" matrix:\n"
" include:\n"
" - shard: {shard}\n"
)
sources: Final = {
".github/workflows/base.yml": base,
"unit.yml": caller.format(job="unit", shard="proxy-auth"),
"proxy-db.yml": caller.format(job="proxy-db", shard="budgets"),
}
assert collisions(sources) == ()
def test_a_reusable_caller_does_not_collide_with_a_plain_job_of_the_same_name() -> None:
sources: Final = {
".github/workflows/base.yml": (
"on:\n workflow_call:\njobs:\n run:\n name: Run tests\n runs-on: ubuntu-latest\n"
),
"unit.yml": (
"on: pull_request\n"
"jobs:\n"
" unit:\n"
" name: ${{ matrix.shard }}\n"
" uses: ./.github/workflows/base.yml\n"
" strategy:\n"
" matrix:\n"
" include:\n"
" - shard: proxy-behavior\n"
),
"postgres.yml": (
"on: pull_request\n"
"jobs:\n"
" postgres:\n"
" name: ${{ matrix.shard }}\n"
" runs-on: ubuntu-latest\n"
" strategy:\n"
" matrix:\n"
" include:\n"
" - shard: proxy-behavior\n"
),
}
assert collisions(sources) == ()
def test_a_workflow_call_only_workflow_publishes_nothing_of_its_own() -> None:
sources: Final = {
"base.yml": "on:\n workflow_call:\njobs:\n run:\n runs-on: ubuntu-latest\n",
"other.yml": "on:\n workflow_call:\njobs:\n run:\n runs-on: ubuntu-latest\n",
}
assert collisions(sources) == ()
def test_a_workflow_call_workflow_that_also_runs_on_pull_request_still_publishes() -> None:
sources: Final = {
"base.yml": "on:\n workflow_call:\n pull_request:\njobs:\n run:\n runs-on: ubuntu-latest\n",
"other.yml": "on: pull_request\njobs:\n run:\n runs-on: ubuntu-latest\n",
}
found: Final = collisions(sources)
assert len(found) == 1
assert "`run` is published by 2 jobs" in found[0]
def test_a_matrix_list_supplies_values_the_same_way_include_rows_do() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\n"
"jobs:\n"
" build:\n"
" name: Analyze (${{ matrix.language }})\n"
" runs-on: ubuntu-latest\n"
" strategy:\n"
" matrix:\n"
" language: [python, go]\n"
),
"b.yml": "on: pull_request\njobs:\n go:\n name: Analyze (go)\n runs-on: ubuntu-latest\n",
}
found: Final = collisions(sources)
assert len(found) == 1
assert "`Analyze (go)` is published by 2 jobs" in found[0]
def test_two_workflows_sharing_a_run_wide_template_are_not_called_a_collision() -> None:
template: Final = (
"on: pull_request\njobs:\n {job}:\n name: ${{{{ github.event_name }}}}-build\n runs-on: ubuntu-latest\n"
)
sources: Final = {
"a.yml": template.format(job="one"),
"b.yml": template.format(job="two"),
}
assert collisions(sources) == ()
assert blind_spots(sources) == ()
def test_two_jobs_of_one_workflow_sharing_a_run_wide_template_are_a_collision() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n"
" one:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n"
" two:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n"
),
}
found: Final = collisions(sources)
assert len(found) == 1
assert "is published 2 times inside a.yml, by job `one`, job `two`" in found[0]
assert exit_code(sources) == 1
def test_a_run_wide_template_carrying_a_matrix_value_does_not_collide_inside_one_workflow() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n"
" one:\n name: ${{ github.event_name }}-${{ matrix.shard }}\n runs-on: ubuntu-latest\n"
" strategy:\n matrix:\n shard: [core, extras]\n"
),
}
assert collisions(sources) == ()
assert blind_spots(sources) == ()
def test_a_run_wide_name_repeated_over_a_matrix_by_one_job_is_a_collision() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n"
" one:\n name: ${{ github.event_name }}-build\n runs-on: ubuntu-latest\n"
" strategy:\n matrix:\n shard: [core, extras]\n"
),
}
found: Final = collisions(sources)
assert len(found) == 1
assert "is published 2 times inside a.yml, by job `one`" in found[0]
def test_a_name_reading_the_job_it_sits_in_stays_out_of_the_comparison() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n"
" one:\n name: ${{ github.job }}-build\n runs-on: ubuntu-latest\n"
" two:\n name: ${{ github.job }}-build\n runs-on: ubuntu-latest\n"
),
}
assert collisions(sources) == ()
assert len(blind_spots(sources)) == 2
def test_a_run_wide_caller_name_collides_through_the_workflow_it_calls() -> None:
sources: Final = {
".github/workflows/a.yml": (
"on: pull_request\njobs:\n"
" one:\n name: ${{ github.event_name }}\n uses: ./.github/workflows/c.yml\n"
" two:\n name: ${{ github.event_name }}\n uses: ./.github/workflows/c.yml\n"
),
".github/workflows/c.yml": "on:\n workflow_call:\njobs:\n build:\n runs-on: ubuntu-latest\n",
}
found: Final = collisions(sources)
assert len(found) == 1
assert "github.event_name }} / build` is published 2 times inside .github/workflows/a.yml" in found[0]
def test_a_run_wide_name_inside_a_called_workflow_collides_under_the_caller() -> None:
sources: Final = {
".github/workflows/a.yml": ("on: pull_request\njobs:\n one:\n uses: ./.github/workflows/c.yml\n"),
".github/workflows/c.yml": (
"on:\n workflow_call:\njobs:\n"
" build:\n name: ${{ github.event_name }}\n runs-on: ubuntu-latest\n"
" lint:\n name: ${{ github.event_name }}\n runs-on: ubuntu-latest\n"
),
}
found: Final = collisions(sources)
assert len(found) == 1
assert "`one / ${{ github.event_name }}` is published 2 times inside .github/workflows/a.yml" in found[0]
def test_a_name_reading_the_workflow_it_sits_in_is_not_called_a_collision() -> None:
template: Final = (
"on: pull_request\njobs:\n {job}:\n name: ${{{{ github.workflow }}}} / build\n runs-on: ubuntu-latest\n"
)
sources: Final = {
"a.yml": template.format(job="one"),
"b.yml": template.format(job="two"),
}
assert collisions(sources) == ()
assert exit_code(sources) == 0
def test_a_format_call_python_accepts_but_github_does_not_publishes_nothing_to_compare() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n one:\n name: ${{ format('{0.real}', matrix.shard) }}\n"
" runs-on: ubuntu-latest\n strategy:\n matrix:\n shard: [core]\n"
),
}
assert collisions(sources) == ()
assert blind_spots(sources) != ()
def test_a_format_call_padding_its_argument_publishes_nothing_to_compare() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n one:\n name: ${{ format('{0:>8}', matrix.shard) }}\n"
" runs-on: ubuntu-latest\n strategy:\n matrix:\n shard: [core]\n"
),
"b.yml": "on: pull_request\njobs:\n two:\n name: ' core'\n runs-on: ubuntu-latest\n",
}
assert collisions(sources) == ()
assert blind_spots(sources) != ()
def test_an_exclude_row_that_is_not_a_mapping_is_reported_rather_than_skipped() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n"
" strategy:\n matrix:\n v: [1, 2]\n exclude:\n - oops\n"
),
"b.yml": "on: pull_request\njobs:\n other:\n name: build (1)\n runs-on: ubuntu-latest\n",
}
assert collisions(sources) == ()
assert blind_spots(sources) != ()
def test_an_exclude_row_holding_a_non_scalar_never_drops_every_combination() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n"
" strategy:\n matrix:\n v: [1, 2]\n exclude:\n - cfg: {k: 1}\n"
),
"b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n",
}
assert collisions(sources) == ()
assert blind_spots(sources) != ()
def test_two_jobs_sharing_a_template_that_reads_per_job_are_not_called_a_collision() -> None:
template: Final = (
"on: pull_request\njobs:\n {job}:\n name: ${{{{ matrix.shard }}}}\n runs-on: ubuntu-latest\n"
)
sources: Final = {
"a.yml": template.format(job="one"),
"b.yml": template.format(job="two"),
}
assert collisions(sources) == ()
assert len(blind_spots(sources)) == 2
def test_a_file_that_is_not_a_workflow_is_reported_rather_than_skipped() -> None:
sources: Final = {
"notes.yml": "just a string\n",
"a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n",
}
found: Final = unreadable(sources)
assert len(found) == 1
assert "notes.yml" in found[0]
assert collisions(sources) == ()
def test_a_workflow_holding_a_job_shape_github_would_reject_is_reported() -> None:
sources: Final = {"a.yml": "on: pull_request\njobs:\n test:\n uses: [not, a, string]\n"}
found: Final = unreadable(sources)
assert len(found) == 1
assert "a.yml" in found[0]
def test_a_conditional_name_expands_to_the_branch_each_matrix_value_takes() -> None:
names: Final = frozenset(
name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER})
)
assert names == frozenset({"core-utils / Run tests", "core-utils / Run tests (Python 3.13)"})
def test_a_conditional_name_never_publishes_the_branch_its_condition_rules_out() -> None:
names: Final = frozenset(
name for name, _ in published({".github/workflows/base.yml": REUSABLE_BASE, "unit.yml": SHARD_CALLER})
)
assert "core-utils / Run tests (Python 3.12)" not in names
def test_a_conditional_reusable_name_collides_with_a_plain_job_publishing_the_same_name() -> None:
sources: Final = {
".github/workflows/base.yml": REUSABLE_BASE,
"unit.yml": SHARD_CALLER,
"postgres.yml": ("on: pull_request\njobs:\n legacy:\n name: core-utils / Run tests\n"),
}
found: Final = collisions(sources)
assert len(found) == 1
assert "`core-utils / Run tests` is published by 2 jobs" in found[0]
def test_a_name_reading_two_matrix_keys_publishes_only_the_pairs_each_include_row_holds() -> None:
names: Final = frozenset(name for name, _ in published({"unit.yml": CORRELATED_ROWS}))
assert names == frozenset({"core-utils on tests/core", "proxy on tests/proxy"})
def test_a_name_reading_two_matrix_keys_never_publishes_a_pair_no_include_row_holds() -> None:
names: Final = frozenset(name for name, _ in published({"unit.yml": CORRELATED_ROWS}))
assert "core-utils on tests/proxy" not in names
assert "proxy on tests/core" not in names
def test_an_include_row_carrying_no_listed_key_extends_every_listed_combination() -> None:
names: Final = frozenset(name for name, _ in published({"unit.yml": LISTED_PLUS_ROW}))
assert names == frozenset({"3.12 fast", "3.13 fast"})
def test_every_workflow_in_the_repo_resolves_every_expression_in_its_job_names() -> None:
unresolved: Final = tuple(f"{owner}: {name}" for name, owner in published(workflow_sources()) if "${{" in name)
assert unresolved == ()
def test_a_matrix_job_with_no_name_publishes_the_id_and_values_github_appends() -> None:
names: Final = frozenset(name for name, _ in published({"a.yml": NAMELESS_MATRIX}))
assert names == frozenset({"build (3.12)", "build (3.13)"})
def test_a_matrix_job_with_no_name_does_not_collide_with_a_plain_job_carrying_its_id() -> None:
sources: Final = {
"a.yml": NAMELESS_MATRIX,
"b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n",
}
assert collisions(sources) == ()
def test_a_matrix_job_with_no_name_collides_with_the_suffixed_name_github_writes() -> None:
sources: Final = {
"a.yml": NAMELESS_MATRIX,
"b.yml": "on: pull_request\njobs:\n legacy:\n name: build (3.13)\n runs-on: ubuntu-latest\n",
}
found: Final = collisions(sources)
assert len(found) == 1
assert "`build (3.13)` is published by 2 jobs" in found[0]
def test_an_excluded_combination_publishes_no_check_run() -> None:
names: Final = frozenset(name for name, _ in published({"unit.yml": EXCLUDED_PAIR}))
assert names == frozenset({"ubuntu-3.12", "ubuntu-3.13", "macos-3.12"})
def test_an_exclude_row_naming_one_key_drops_every_combination_carrying_it() -> None:
names: Final = frozenset(name for name, _ in published({"unit.yml": EXCLUDED_KEY}))
assert names == frozenset({"ubuntu-3.12", "ubuntu-3.13"})
def test_a_boolean_matrix_value_renders_the_way_github_writes_it() -> None:
names: Final = frozenset(name for name, _ in published({"unit.yml": BOOLEAN_MATRIX}))
assert names == frozenset({"cache true", "cache false"})
def test_a_format_call_its_arguments_cannot_fill_publishes_nothing_to_compare() -> None:
sources: Final = {"unit.yml": UNFILLABLE_FORMAT}
assert frozenset(name for name, _ in published(sources)) == frozenset()
assert "its name stays" in blind_spots(sources)[0]
def test_a_call_to_a_workflow_outside_the_repo_is_reported_rather_than_guessed() -> None:
sources: Final = {
"a.yml": "on: pull_request\njobs:\n unit:\n uses: BerriAI/other/.github/workflows/base.yml@main\n",
"b.yml": "on: pull_request\njobs:\n unit:\n runs-on: ubuntu-latest\n",
}
assert frozenset(name for name, _ in published(sources)) == frozenset({"unit"})
assert collisions(sources) == ()
assert "outside this repository" in blind_spots(sources)[0]
def test_a_chain_of_local_reusable_calls_publishes_every_level_of_the_chain() -> None:
sources: Final = {
".github/workflows/leaf.yml": (
"on:\n workflow_call:\njobs:\n run:\n name: Leaf\n runs-on: ubuntu-latest\n"
),
".github/workflows/mid.yml": (
"on:\n workflow_call:\njobs:\n call:\n name: Mid\n uses: ./.github/workflows/leaf.yml\n"
),
"top.yml": "on: pull_request\njobs:\n top:\n name: Top\n uses: ./.github/workflows/mid.yml\n",
}
names: Final = frozenset(name for name, _ in published(sources))
assert names == frozenset({"Top / Mid / Leaf"})
def test_a_job_name_that_is_not_a_string_still_publishes_the_value_github_renders() -> None:
sources: Final = {
"a.yml": "on: pull_request\njobs:\n sweep:\n name: 2024\n runs-on: ubuntu-latest\n",
"b.yml": 'on: pull_request\njobs:\n other:\n name: "2024"\n runs-on: ubuntu-latest\n',
}
found: Final = collisions(sources)
assert len(found) == 1
assert "`2024` is published by 2 jobs" in found[0]
def test_the_check_fails_when_a_file_in_the_workflows_directory_cannot_be_read() -> None:
assert exit_code({"notes.yml": "just a string\n"}) == 1
def test_the_check_fails_when_two_jobs_publish_one_check_run_name() -> None:
plain: Final = "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n"
assert exit_code({"a.yml": plain, "b.yml": plain}) == 1
def test_the_check_passes_when_every_file_reads_and_every_name_is_unique() -> None:
sources: Final = {
"a.yml": "on: pull_request\njobs:\n test:\n runs-on: ubuntu-latest\n",
"b.yml": "on: pull_request\njobs:\n sweep:\n runs-on: ubuntu-latest\n",
}
assert exit_code(sources) == 0
def test_two_callers_of_one_reusable_workflow_named_from_its_inputs_do_not_collide() -> None:
sources: Final = {
".github/workflows/callee.yml": (
"on:\n workflow_call:\njobs:\n run:\n name: ${{ inputs.suite }}\n runs-on: ubuntu-latest\n"
),
"caller.yml": (
"on: pull_request\njobs:\n"
" alpha:\n name: A\n uses: ./.github/workflows/callee.yml\n with:\n suite: alpha\n"
" beta:\n name: A\n uses: ./.github/workflows/callee.yml\n with:\n suite: beta\n"
),
}
assert collisions(sources) == ()
assert len(blind_spots(sources)) == 2
def test_a_matrix_that_is_itself_an_expression_never_collapses_onto_the_bare_job_id() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n build:\n strategy:\n"
" matrix: ${{ fromJson(needs.plan.outputs.matrix) }}\n runs-on: ubuntu-latest\n"
),
"b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n",
}
assert frozenset(name for name, _ in published(sources)) == frozenset({"build"})
assert collisions(sources) == ()
assert "the matrix itself comes from an expression" in blind_spots(sources)[0]
def test_a_matrix_listing_objects_never_collapses_onto_the_bare_job_id() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n build:\n strategy:\n matrix:\n target:\n"
" - os: ubuntu\n - os: windows\n runs-on: ubuntu-latest\n"
),
"b.yml": "on: pull_request\njobs:\n build:\n runs-on: ubuntu-latest\n",
}
assert frozenset(name for name, _ in published(sources)) == frozenset({"build"})
assert collisions(sources) == ()
assert "not plain scalars" in blind_spots(sources)[0]
def test_a_call_to_a_workflow_file_the_checkout_does_not_hold_is_reported() -> None:
sources: Final = {"a.yml": "on: pull_request\njobs:\n unit:\n uses: ./.github/workflows/gone.yml\n"}
assert collisions(sources) == ()
assert "which this checkout does not hold" in blind_spots(sources)[0]
def test_reusable_workflows_calling_each_other_in_a_loop_are_reported_not_followed() -> None:
sources: Final = {
".github/workflows/a.yml": (
"on:\n workflow_call:\njobs:\n call:\n name: A\n uses: ./.github/workflows/b.yml\n"
),
".github/workflows/b.yml": (
"on:\n workflow_call:\njobs:\n call:\n name: B\n uses: ./.github/workflows/a.yml\n"
),
"top.yml": "on: pull_request\njobs:\n top:\n name: Top\n uses: ./.github/workflows/a.yml\n",
}
assert collisions(sources) == ()
assert any("loops back on itself" in spot for spot in blind_spots(sources))
def test_a_caller_still_publishes_the_callee_jobs_it_can_read() -> None:
sources: Final = {
".github/workflows/callee.yml": (
"on:\n workflow_call:\njobs:\n"
" lint:\n name: Lint\n runs-on: ubuntu-latest\n"
" suite:\n name: ${{ inputs.suite }}\n runs-on: ubuntu-latest\n"
),
"caller.yml": "on: pull_request\njobs:\n call:\n name: A\n uses: ./.github/workflows/callee.yml\n",
}
assert frozenset(name for name, _ in published(sources)) == frozenset({"A / Lint"})
assert len(blind_spots(sources)) == 1
def test_a_name_the_check_cannot_work_out_is_reported_without_failing_the_check() -> None:
sources: Final = {
"a.yml": "on: pull_request\njobs:\n unit:\n uses: BerriAI/other/.github/workflows/base.yml@main\n",
}
assert blind_spots(sources) != ()
assert exit_code(sources) == 0
def test_a_caller_whose_own_name_is_unreadable_publishes_none_of_its_callee_names() -> None:
sources: Final = {
".github/workflows/callee.yml": (
"on:\n workflow_call:\njobs:\n lint:\n name: Lint\n runs-on: ubuntu-latest\n"
),
"caller.yml": (
"on: pull_request\njobs:\n call:\n name: ${{ matrix.suite }}\n"
" uses: ./.github/workflows/callee.yml\n"
),
"other.yml": "on: pull_request\njobs:\n plain:\n name: Lint\n runs-on: ubuntu-latest\n",
}
assert frozenset(name for name, _ in published(sources)) == frozenset({"Lint"})
assert collisions(sources) == ()
assert "its name stays" in blind_spots(sources)[0]
def test_an_include_row_naming_a_listed_key_extends_only_the_combinations_it_matches() -> None:
sources: Final = {
"unit.yml": (
"on: pull_request\njobs:\n unit:\n"
" name: ${{ matrix.python-version }} ${{ matrix.label }}\n"
" runs-on: ubuntu-latest\n strategy:\n matrix:\n"
' python-version: ["3.12", "3.13"]\n'
" include:\n"
' - python-version: "3.12"\n'
" label: fast\n"
)
}
assert frozenset(name for name, _ in published(sources)) == frozenset({"3.12 fast"})
assert len(blind_spots(sources)) == 1
def test_a_job_whose_whole_strategy_is_an_expression_is_reported_rather_than_rejecting_the_file() -> None:
sources: Final = {
"plan.yml": (
"on: pull_request\njobs:\n plan:\n name: Plan\n runs-on: ubuntu-latest\n"
" fan:\n strategy: ${{ fromJSON(needs.plan.outputs.strategy) }}\n runs-on: ubuntu-latest\n"
)
}
assert unreadable(sources) == ()
assert frozenset(name for name, _ in published(sources)) == frozenset({"Plan"})
assert "`strategy` comes from an expression" in blind_spots(sources)[0]
assert exit_code(sources) == 0
def test_one_job_publishing_one_name_for_every_matrix_combination_is_a_collision() -> None:
sources: Final = {
"unit.yml": (
"on: pull_request\njobs:\n build:\n name: Run tests\n runs-on: ubuntu-latest\n"
' strategy:\n matrix:\n python-version: ["3.12", "3.13"]\n'
)
}
found: Final = collisions(sources)
assert len(found) == 1
assert "`Run tests` is published 2 times by unit.yml job `build`" in found[0]
assert exit_code(sources) == 1
def test_a_name_carrying_a_matrix_value_publishes_one_name_per_combination_without_colliding() -> None:
sources: Final = {
"unit.yml": (
"on: pull_request\njobs:\n build:\n name: Run tests ${{ matrix.python-version }}\n"
' runs-on: ubuntu-latest\n strategy:\n matrix:\n python-version: ["3.12", "3.13"]\n'
)
}
assert frozenset(name for name, _ in published(sources)) == frozenset({"Run tests 3.12", "Run tests 3.13"})
assert collisions(sources) == ()
assert exit_code(sources) == 0
def test_a_file_that_is_not_valid_yaml_is_reported_rather_than_raising() -> None:
sources: Final = {"broken.yml": "jobs:\n build: [\n"}
assert unreadable(sources) == (
"broken.yml sits in the workflows directory but it does not read as one YAML "
"document, so none of its jobs were checked.",
)
assert exit_code(sources) == 1
def test_a_file_holding_two_yaml_documents_is_reported_rather_than_raising() -> None:
sources: Final = {"two.yml": "on: pull_request\n---\non: push\n"}
assert len(unreadable(sources)) == 1
assert exit_code(sources) == 1
def test_an_exclude_that_is_itself_an_expression_is_reported_rather_than_ignored() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n one:\n name: build\n runs-on: ubuntu-latest\n"
" strategy:\n matrix:\n python: ['3.11', '3.12']\n"
" exclude: ${{ fromJson(vars.SKIP) }}\n"
),
}
found: Final = blind_spots(sources)
assert collisions(sources) == ()
assert len(found) == 1
assert "a matrix `exclude` is itself an expression" in found[0]
def test_an_include_that_is_itself_an_expression_is_reported_rather_than_ignored() -> None:
sources: Final = {
"a.yml": (
"on: pull_request\njobs:\n one:\n name: build-${{ matrix.python }}\n runs-on: ubuntu-latest\n"
" strategy:\n matrix:\n python: ['3.11']\n"
" include: ${{ fromJson(vars.EXTRA) }}\n"
),
}
found: Final = blind_spots(sources)
assert collisions(sources) == ()
assert len(found) == 1
assert "a matrix `include` is itself an expression" in found[0]

2
tests/e2e/coverage_registry/mgmt.yaml
View file

@ -5,6 +5,8 @@
- {id: mgmt.key.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:2462", rationale: "Budget/model changes persist"}
- {id: mgmt.key.update.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:2462", rationale: "Non-admin cannot escalate perms"}
- {id: mgmt.key.update.happy_path, module: mgmt, tier: P1, surface: ui, assertions: [happy_path], source: "key_management_endpoints.py:2462", rationale: "Key edit through the dashboard"}
- {id: mgmt.key.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "key_management_endpoints.py:2829", rationale: "A partial /key/update changes only the field it names; alias, models, limits, budget window, team and metadata read back unchanged on every gateway replica"}
- {id: mgmt.key.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "key_management_endpoints.py:2829", rationale: "An explicit null on /key/update clears max_budget and budget_duration, and the derived budget_reset_at with it, on every gateway replica"}
- {id: mgmt.key.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3122", rationale: "Deletion revokes future calls"}
- {id: mgmt.key.delete.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:3122", rationale: "Non-owner cannot delete"}
- {id: mgmt.key.info.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3380", rationale: "Info reflects all writes"}

2
tests/e2e/coverage_registry/reliability.yaml
View file

@ -7,7 +7,7 @@
- {id: reliability.retry.timeout.succeeds_within_retries, module: reliability, tier: P0, behavior: retry, variant: timeout, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:44", rationale: "Timeout retried per policy"}
- {id: reliability.retry.429.succeeds_within_retries, module: reliability, tier: P0, behavior: retry, variant: "429", assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:46", rationale: "429 retried per RateLimitErrorRetries policy"}
- {id: reliability.retry.auth.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: auth, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:42", rationale: "Transient auth glitch retry"}
- {id: reliability.retry.context_window.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: context_window, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:51", rationale: "Multi-attempt on context error"}
- {id: reliability.retry.context_window.succeeds_within_retries, module: reliability, tier: P1, behavior: retry, variant: context_window, assertions: [succeeds_within_retries], exercised_on: [chat_completions, messages], source: "get_retry_from_policy.py:51", fail_before_fix: proven, rationale: "A context-window 400 under BadRequestErrorRetries retries onto a sibling deployment in the same model group, instead of coming straight back as the 400 the deployment that just refused it returned"}
- {id: reliability.cooldown.5xx.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "5xx", assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:40", rationale: "Deployment cools after repeated 5xx, recovers after cooldown_time"}
- {id: reliability.cooldown.429.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "429", assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:69", rationale: "Cools on 429, avoids hammering exhausted provider"}
- {id: reliability.cooldown.auth.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: auth, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages], source: "cooldown_handlers.py:74", rationale: "Cools on 401 auth error"}

299
tests/e2e/management/test_key_lifecycle_e2e.py Normal file
View file

@ -0,0 +1,299 @@
"""Live e2e: one virtual key walked through its whole lifecycle, read back on every
gateway replica.
Create, read, partial update, clear, enforce, delete: one method per step, and every
step creates its own team and key (both deleted on teardown) so a step reruns or skips
on its own. Writes go through the control plane; read-backs poll every URL in
PROXY_REPLICA_URLS until each replica converges, because a write that is visible on the
gateway that took it and stale on its neighbour is exactly the failure this file exists
to catch. Revocation is the slowest of those: a deleted key stays usable on the other
replicas until their auth cache entry expires, so the delete step polls each of them
rather than asserting once.
"""
from __future__ import annotations
import time
from collections.abc import Callable, Iterator, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_http import Result, StreamingResponse, Success, UnknownApiError, unwrap
from lifecycle import ResourceManager
from management_client import MODEL_ACCESS_DENIED_MARKER, ManagementClient
from models import (
CLEAR,
ChatBody,
ChatMessage,
KeyGenerateBody,
KeyGenerateResponse,
KeyInfo,
KeyInfoParams,
KeyInfoResponse,
KeyMetadata,
KeyUpdateBody,
LiteLLMParamsBody,
TeamNewBody,
)
from proxy_client import Converged, NotConverged, Poller, await_converged, await_converged_everywhere
from transport import Transport
pytestmark = pytest.mark.e2e
BACKING_MODEL: Final = "gpt-4o-mini"
DENIED_MODEL: Final = "gpt-5.5"
MAX_BUDGET: Final = 25.0
TPM_LIMIT: Final = 313131
RPM_LIMIT: Final = 323232
UPDATED_RPM_LIMIT: Final = 424242
BUDGET_DURATION: Final = "30d"
@dataclass(frozen=True, slots=True)
class CreatedKey:
written: KeyGenerateBody
response: KeyGenerateResponse
@property
def key(self) -> str:
return self.response.key
@pytest.fixture(scope="module")
def mock_deployment(client: ManagementClient) -> Iterator[str]:
"""A deployment that answers from a canned response, so the enforcement step needs no
provider key. The alias carries a unique marker, like every other model this suite
registers, so concurrent runs never share one model group."""
model_name: Final = f"e2e-key-lifecycle-{unique_marker()}"
model_id: Final = client.proxy.create_model(model_name, LiteLLMParamsBody(model=BACKING_MODEL, mock_response="ok"))
try:
yield model_name
finally:
client.proxy.delete_model(model_id)
def _await[T](client: ManagementClient, poller: Poller[T], converged: Callable[[T], bool], failure: str) -> T:
outcome: Final = await_converged(
poller,
converged=converged,
timeout=client.proxy.poll_timeout,
interval=client.proxy.poll_interval,
now=time.monotonic,
sleep=time.sleep,
)
match outcome:
case Converged(result=result):
return result
case NotConverged(last_result=last):
pytest.fail(f"{failure}; last outcome: {last}")
def _chat_poller(transport: Transport, key: str, model: str) -> Poller[StreamingResponse]:
return lambda: transport.send(
"/chat/completions",
headers=transport.bearer(key),
json=ChatBody(
model=model,
messages=[ChatMessage(role="user", content=f"say hi {unique_marker()}")],
max_tokens=16,
),
)
def _create_key(client: ManagementClient, resources: ResourceManager, model: str) -> CreatedKey:
marker: Final = unique_marker()
team_id: Final = client.create_team(TeamNewBody(team_alias=f"e2e-key-lifecycle-team-{marker}"))
resources.defer(lambda: client.delete_team(team_id))
written: Final = KeyGenerateBody(
key_alias=f"e2e-key-lifecycle-{marker}",
models=[model],
max_budget=MAX_BUDGET,
tpm_limit=TPM_LIMIT,
rpm_limit=RPM_LIMIT,
budget_duration=BUDGET_DURATION,
metadata=KeyMetadata(tag=marker),
team_id=team_id,
)
response: Final = unwrap(client.generate_key(written))
resources.defer(lambda: client.proxy.delete_key(response.key))
return CreatedKey(written=written, response=response)
def _key_info_everywhere(
client: ManagementClient, key: str, settled: Callable[[KeyInfo], bool]
) -> Mapping[str, KeyInfo]:
def converged(result: Result[KeyInfoResponse]) -> bool:
return isinstance(result, Success) and settled(result.data.info)
reads: Final = client.proxy.read_back_everywhere(
"/key/info", params=KeyInfoParams(key=key), response_type=KeyInfoResponse, converged=converged
)
return MappingProxyType({replica: unwrap(read).info for replica, read in reads.items()})
def _is_key_not_found(result: Result[KeyInfoResponse]) -> bool:
return isinstance(result, UnknownApiError) and result.status_code == 404
def _assert_reads_back(info: KeyInfo, expected: KeyGenerateBody, replica: str) -> None:
for field, observed, wanted in (
("key_alias", info.key_alias, expected.key_alias),
("models", info.models, expected.models),
("max_budget", info.max_budget, expected.max_budget),
("tpm_limit", info.tpm_limit, expected.tpm_limit),
("rpm_limit", info.rpm_limit, expected.rpm_limit),
("budget_duration", info.budget_duration, expected.budget_duration),
("team_id", info.team_id, expected.team_id),
("metadata", info.metadata, expected.metadata),
):
assert observed == wanted, f"{replica}: /key/info reports {field}={observed!r}, expected {wanted!r}"
def _poll_chat_ok(client: ManagementClient, key: str, model: str) -> None:
_ = _await(
client,
_chat_poller(client.proxy.transport, key, model),
lambda outcome: outcome.ok,
f"chat on {model} never succeeded for the key before the deadline",
)
def _warm_every_replica(client: ManagementClient, key: str, model: str) -> None:
"""Serve one call from every replica, so each has the key in its auth cache. Without
this the revocation check below would only prove a replica rejects a key it never
knew, which is true of any random string."""
for replica, transport in client.proxy.replicas.items():
_ = _await(
client,
_chat_poller(transport, key, model),
lambda outcome: outcome.ok,
f"{replica}: chat on {model} never succeeded for the key before the deadline",
)
def _assert_chat_rejected_everywhere(client: ManagementClient, key: str, model: str) -> None:
outcomes: Final = await_converged_everywhere(
{replica: _chat_poller(transport, key, model) for replica, transport in client.proxy.replicas.items()},
converged=lambda outcome: outcome.status_code == 401,
timeout=client.proxy.poll_timeout,
interval=client.proxy.poll_interval,
now=time.monotonic,
sleep=time.sleep,
)
for replica, outcome in outcomes.items():
assert isinstance(outcome, Converged), (
f"{replica}: the deleted key was still accepted on chat after "
f"{client.proxy.poll_timeout}s, last status {outcome.last_result.status_code}"
)
class TestKeyLifecycle:
def test_create_echoes_every_field_written(
self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
) -> None:
created: Final = _create_key(client, resources, mock_deployment)
response: Final = created.response
for field, observed, wanted in (
("key_alias", response.key_alias, created.written.key_alias),
("models", response.models, created.written.models),
("max_budget", response.max_budget, created.written.max_budget),
("tpm_limit", response.tpm_limit, created.written.tpm_limit),
("rpm_limit", response.rpm_limit, created.written.rpm_limit),
("budget_duration", response.budget_duration, created.written.budget_duration),
("team_id", response.team_id, created.written.team_id),
("metadata", response.metadata, created.written.metadata),
):
assert observed == wanted, f"/key/generate echoed {field}={observed!r}, sent {wanted!r}"
def test_read_reflects_the_create_on_every_replica(
self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
) -> None:
created: Final = _create_key(client, resources, mock_deployment)
infos: Final = _key_info_everywhere(
client, created.key, lambda info: info.key_alias == created.written.key_alias
)
for replica, info in infos.items():
_assert_reads_back(info, created.written, replica)
assert info.budget_reset_at is not None, (
f"{replica}: /key/info reports no budget_reset_at for budget_duration={BUDGET_DURATION!r}"
)
@pytest.mark.covers("mgmt.key.update.preserves_unrelated_fields")
def test_partial_update_changes_only_the_named_field(
self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
) -> None:
created: Final = _create_key(client, resources, mock_deployment)
before: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == RPM_LIMIT)
_ = unwrap(client.update_key(KeyUpdateBody(key=created.key, rpm_limit=UPDATED_RPM_LIMIT)))
after: Final = _key_info_everywhere(client, created.key, lambda info: info.rpm_limit == UPDATED_RPM_LIMIT)
for replica, info in after.items():
_assert_reads_back(info, created.written.model_copy(update={"rpm_limit": UPDATED_RPM_LIMIT}), replica)
assert info.budget_reset_at == before[replica].budget_reset_at, (
f"{replica}: budget_reset_at moved from {before[replica].budget_reset_at!r} to "
f"{info.budget_reset_at!r} on a /key/update that did not name budget_duration"
)
@pytest.mark.covers("mgmt.key.update.clear_persists")
def test_explicit_null_clears_the_budget_and_its_reset_time(
self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
) -> None:
created: Final = _create_key(client, resources, mock_deployment)
_ = _key_info_everywhere(client, created.key, lambda info: info.max_budget == MAX_BUDGET)
_ = unwrap(client.update_key(KeyUpdateBody(key=created.key, max_budget=CLEAR, budget_duration=CLEAR)))
cleared: Final = _key_info_everywhere(client, created.key, lambda info: info.max_budget is None)
for replica, info in cleared.items():
assert info.budget_duration is None, (
f"{replica}: budget_duration={info.budget_duration!r} survived an explicit null"
)
assert info.budget_reset_at is None, (
f"{replica}: clearing budget_duration left budget_reset_at={info.budget_reset_at!r}"
)
_assert_reads_back(
info, created.written.model_copy(update={"max_budget": None, "budget_duration": None}), replica
)
def test_key_serves_its_model_and_is_denied_others(
self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
) -> None:
created: Final = _create_key(client, resources, mock_deployment)
_poll_chat_ok(client, created.key, mock_deployment)
denied: Final = client.chat_status(created.key, DENIED_MODEL, f"say hi {unique_marker()}")
assert denied.status_code == 403, (
f"chat on {DENIED_MODEL!r} outside the key's model list must be denied 403, got "
f"{denied.status_code}: {denied.body[:300]}"
)
assert MODEL_ACCESS_DENIED_MARKER in denied.body, (
f"403 body must be a model-access denial, got: {denied.body[:300]}"
)
def test_delete_revokes_info_and_chat_on_every_replica(
self, client: ManagementClient, resources: ResourceManager, mock_deployment: str
) -> None:
"""The teardown's deferred delete fires again on the already-deleted key by
design: the deferred cleanup must survive this test failing before the
in-body delete, and a repeat /key/delete is a cheap no-op the warn-only
teardown absorbs."""
created: Final = _create_key(client, resources, mock_deployment)
_warm_every_replica(client, created.key, mock_deployment)
client.delete_key_strict(created.key)
_ = client.proxy.read_back_everywhere(
"/key/info",
params=KeyInfoParams(key=created.key),
response_type=KeyInfoResponse,
converged=_is_key_not_found,
)
_assert_chat_rejected_everywhere(client, created.key, mock_deployment)

37
tests/e2e/models.py
View file

@ -8,9 +8,9 @@ from __future__ import annotations
from collections.abc import Sequence
from datetime import datetime
from typing import Literal
from typing import Final, Literal
from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_validator
from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator
# ---------- keys ----------
@ -49,6 +49,7 @@ class KeyMetadata(BaseModel):
logging: list[KeyLoggingCallback] | None = None
priority: str | None = None
batch_enqueued_token_limit: int | None = None
tag: str | None = None
class ObjectPermission(BaseModel):
@ -81,6 +82,14 @@ class KeyGenerateBody(BaseModel):
class KeyGenerateResponse(BaseModel):
key: str
key_alias: str | None = None
models: list[str] = []
max_budget: float | None = None
tpm_limit: int | None = None
rpm_limit: int | None = None
budget_duration: str | None = None
team_id: str | None = None
metadata: KeyMetadata | None = None
class KeyRegenerateBody(BaseModel):
@ -122,6 +131,7 @@ class KeyInfo(BaseModel):
blocked: bool | None = None
spend: float | None = None
max_budget: float | None = None
budget_duration: str | None = None
budget_reset_at: str | None = None
budget_id: str | None = None
litellm_budget_table: LiteLLMBudgetTable | None = None
@ -291,6 +301,7 @@ class RouterSettingsOverride(BaseModel):
context_window_fallbacks: list[dict[str, list[str]]] | None = None
content_policy_fallbacks: list[dict[str, list[str]]] | None = None
num_retries: int | None = None
model_group_retry_policy: dict[str, dict[str, int]] | None = None
enable_tag_filtering: bool | None = None
@ -918,12 +929,34 @@ class CredentialCreateResponse(BaseModel):
# ---------- key / team / user / organization management ----------
class Cleared(BaseModel):
"""An explicit JSON null in a merge-patch body. The transport drops `None` fields
before sending (`exclude_none`), so `None` means "leave the stored value alone"; a
field set to `CLEAR` reaches the wire as `null`, which tells the proxy to clear it."""
model_config = ConfigDict(frozen=True)
@model_serializer
def _as_null(self) -> None:
return None
CLEAR: Final = Cleared()
class KeyUpdateBody(BaseModel):
"""POST /key/update is a merge patch: a field left `None` is dropped from the body and
keeps its stored value, `CLEAR` sends an explicit null that clears it (`budget_duration`
clears `budget_reset_at` with it), and `metadata` replaces the stored metadata wholesale."""
key: str
models: list[str] | None = None
key_alias: str | None = None
tpm_limit: int | None = None
rpm_limit: int | None = None
max_budget: float | Cleared | None = None
budget_duration: str | Cleared | None = None
metadata: KeyMetadata | None = None
class KeyBlockBody(BaseModel):

152
tests/e2e/proxy_client.py
View file

@ -16,6 +16,8 @@ from datetime import datetime
from types import MappingProxyType
from typing import Final
from pydantic import BaseModel
from e2e_http import (
AnthropicHeaders,
AuthHeaders,
@ -149,9 +151,7 @@ def await_servable(
last_result: Result[ModelsListResponse] | None = None
while True:
t = now()
phase_deadline = (
started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds
)
phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds
remaining = phase_deadline - t
if remaining <= 0:
if (
@ -164,9 +164,7 @@ def await_servable(
poll_timeout = min(request_timeout, remaining)
last_result = list_models(poll_timeout)
listed = isinstance(last_result, Success) and any(
entry.id == model_name for entry in last_result.data.data
)
listed = isinstance(last_result, Success) and any(entry.id == model_name for entry in last_result.data.data)
t = now()
if not listed:
first_seen_at = None
@ -179,9 +177,7 @@ def await_servable(
elif t - first_seen_at >= db_sync_seconds:
return Servable()
phase_deadline = (
started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds
)
phase_deadline = started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds
wait = min(interval, phase_deadline - now())
if wait > 0:
sleep(wait)
@ -239,6 +235,88 @@ def servable_timeout_message(
)
type Poller[T] = Callable[[], T]
@dataclass(frozen=True, slots=True)
class Converged[T]:
result: T
@dataclass(frozen=True, slots=True)
class NotConverged[T]:
"""The deadline passed without a read satisfying the predicate; `last_result` is
the final read, so the caller can tell a stale body from a failed request."""
last_result: T
type ConvergeOutcome[T] = Converged[T] | NotConverged[T]
def await_converged[T](
poll: Poller[T],
*,
converged: Callable[[T], bool],
timeout: float,
interval: float,
now: Callable[[], float],
sleep: Callable[[float], None],
) -> ConvergeOutcome[T]:
"""Poll until a read satisfies `converged` or `timeout` elapses.
Polls before testing the deadline, so a zero or already-spent budget still gets one
attempt, and sleeps only min(interval, time left), so the attempt that lands exactly
on the deadline is taken rather than skipped. Clock and sleep are injected."""
deadline: Final = now() + timeout
while True:
result = poll()
if converged(result):
return Converged(result=result)
remaining = deadline - now()
if remaining <= 0:
return NotConverged(last_result=result)
sleep(min(interval, remaining))
def await_converged_everywhere[T](
pollers: Mapping[str, Poller[T]],
*,
converged: Callable[[T], bool],
timeout: float,
interval: float,
now: Callable[[], float],
sleep: Callable[[float], None],
) -> Mapping[str, ConvergeOutcome[T]]:
"""`await_converged` against every replica in turn, each with the full budget, so a
replica that lags behind the one a write landed on is polled until it catches up
rather than failing on its first stale read."""
return MappingProxyType(
{
replica: await_converged(
poll, converged=converged, timeout=timeout, interval=interval, now=now, sleep=sleep
)
for replica, poll in pollers.items()
}
)
def first_lagging_replica[T](
outcomes: Mapping[str, ConvergeOutcome[T]],
) -> tuple[str, NotConverged[T]] | None:
return next(
((replica, outcome) for replica, outcome in outcomes.items() if isinstance(outcome, NotConverged)),
None,
)
def converge_timeout_message(*, what: str, replica: str, timeout: float, last_result: object) -> str:
return (
f"{what} on {replica} never converged within {timeout}s of the write "
f"(control/data-plane propagation issue); last read: {last_result}"
)
@dataclass(frozen=True, slots=True)
class ProxyClient:
transport: Transport
@ -290,6 +368,52 @@ class ProxyClient:
)
).info
def read_back_everywhere[R: BaseModel](
self,
path: str,
*,
params: BaseModel,
response_type: type[R],
converged: Callable[[Result[R]], bool],
) -> Mapping[str, Result[R]]:
"""GET `path` under the master key on every replica in PROXY_REPLICA_URLS (the
data-plane URL alone when the stack exports no per-gateway addresses), polling
each to poll_timeout until its read satisfies `converged`. Returns that read per
replica, or fails naming the first replica that never converged and its last
read. Behind a load balancer the single address proves one replica converged,
not all of them; only per-gateway addresses make this a fleet-wide proof."""
outcomes: Final = await_converged_everywhere(
{
url: self._body_poller(transport, path, params, response_type)
for url, transport in self.replicas.items()
},
converged=converged,
timeout=self.poll_timeout,
interval=self.poll_interval,
now=time.monotonic,
sleep=time.sleep,
)
lagging: Final = first_lagging_replica(outcomes)
if lagging is not None:
replica, outcome = lagging
raise AssertionError(
converge_timeout_message(
what=f"GET {path}",
replica=replica,
timeout=self.poll_timeout,
last_result=outcome.last_result,
)
)
return MappingProxyType(
{replica: outcome.result for replica, outcome in outcomes.items() if isinstance(outcome, Converged)}
)
@staticmethod
def _body_poller[R: BaseModel](
transport: Transport, path: str, params: BaseModel, response_type: type[R]
) -> Poller[Result[R]]:
return lambda: transport.get(path, headers=transport.master, params=params, response_type=response_type)
def model_info(self) -> list[ModelInfoEntry]:
"""Every configured deployment with the price the proxy resolved for it
(config override merged over cost-map defaults)."""
@ -320,9 +444,7 @@ class ProxyClient:
response_type=FileListResponse,
)
def list_fine_tuning_jobs(
self, key: str, params: FineTuningJobsParams
) -> Result[FineTuningJobsResponse]:
def list_fine_tuning_jobs(self, key: str, params: FineTuningJobsParams) -> Result[FineTuningJobsResponse]:
return self.transport.get(
"/v1/fine_tuning/jobs",
headers=self.transport.bearer(key),
@ -375,7 +497,11 @@ class ProxyClient:
)
).model_id
written_at = time.monotonic()
self._await_model_servable(body.model_name, listed_for)
try:
self._await_model_servable(body.model_name, listed_for)
except BaseException:
self.delete_model(model_id)
raise
settle_propagation(written_at)
return model_id

19
tests/e2e/router/reliability_support.py
View file

@ -73,10 +73,25 @@ def create_always_timing_out_deployment(proxy: ProxyClient, name: str) -> str:
)
def create_always_picked_small_context_deployment(proxy: ProxyClient, name: str) -> str:
"""The always-picked half of a retry pair on the smallest-context model OpenAI
still serves: it holds all of the model group's shuffle weight, so an oversized
prompt opens on it and earns a real context-window refusal, which never benches
a deployment, so only the retry itself can steer the request off it."""
return proxy.register_model(
ModelNewBody(
model_name=name,
litellm_params=LiteLLMParamsBody(model=SMALL_CONTEXT_MODEL, api_key=REAL_KEY, weight=1),
model_info=ModelInfoBody(),
)
)
def create_zero_weight_backup_deployment(proxy: ProxyClient, name: str) -> str:
"""The other half of a retry pair: healthy, but weight 0, so the weighted shuffle
never opens on it. It is reachable only once its sibling is benched and the
weighted pick falls through to a uniform one over what is left."""
never opens on it. It is reachable only once its sibling is out of the running,
benched by a cooldown or skipped by the retry, and the weighted pick falls through
to a uniform one over what is left."""
return proxy.register_model(
ModelNewBody(
model_name=name,

88
tests/e2e/router/test_reliability_retries_e2e.py
View file

@ -1,13 +1,17 @@
"""Live e2e: a request that fails on its first deployment is retried inside its own
model group and still comes back a completion.
The model group is a pair: an always-timing-out deployment that holds all of the
group's shuffle weight, and a healthy backup at weight 0. The weighted pick always
opens on the timing-out one, its first Timeout benches it (an
`allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`), and the retry falls
through to the only deployment left. So the customer sees a completion and the
proxy reports that it took a retry to get there, with no random first pick in the
middle of it.
Each model group is a pair: a deployment that always refuses and holds all of the
group's shuffle weight, plus a healthy backup at weight 0. The weighted pick always
opens on the refusing one, so the customer sees a completion only if the retry
lands on the backup, and the proxy reports that it took a retry to get there, with
no random first pick in the middle of it.
The timeout pair relies on cooldown: the first Timeout benches the timing-out
deployment (an `allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`) and the
retry falls through to the only deployment left. The context-window pair cannot:
a 400 never benches a deployment, so the retry policy's `BadRequestErrorRetries`
has to steer the retry off the deployment that just refused the prompt.
"""
from __future__ import annotations
@ -16,20 +20,48 @@ import pytest
from complexity_router_client import ComplexityRouterClient
from e2e_config import unique_marker
from e2e_http import StreamingResponse
from lifecycle import ResourceManager
from models import RouterSettingsOverride
from reliability_support import (
chat_override,
completion_tokens_of,
content_of,
create_always_picked_small_context_deployment,
create_always_timing_out_deployment,
create_zero_weight_backup_deployment,
finish_reason_of,
oversized_prompt,
)
pytestmark = pytest.mark.e2e
def assert_retry_landed_on_backup(resp: StreamingResponse) -> None:
assert resp.status_code == 200, (
f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}"
)
attempted = resp.headers.get("x-litellm-attempted-retries")
assert attempted is not None, "response is missing the x-litellm-attempted-retries header"
assert int(attempted) >= 1, (
f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never "
"opened on the refusing deployment, so this proves nothing about retries"
)
content = content_of(resp)
finish_reason = finish_reason_of(resp)
completion_tokens = completion_tokens_of(resp) or 0
assert isinstance(content, str), (
f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})"
)
assert content or (finish_reason == "length" and completion_tokens > 0), (
f"the retry returned empty content with finish_reason={finish_reason!r}, "
f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget "
f"was spent on non-visible reasoning (body={resp.body[:300]})"
)
class TestReliabilityRetries:
@pytest.mark.covers("reliability.retry.timeout.succeeds_within_retries")
def test_timeout_on_first_deployment_succeeds_on_retry(
@ -49,25 +81,27 @@ class TestReliabilityRetries:
override=RouterSettingsOverride(num_retries=2),
)
assert resp.status_code == 200, (
f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}"
assert_retry_landed_on_backup(resp)
@pytest.mark.covers("reliability.retry.context_window.succeeds_within_retries")
def test_context_window_refusal_on_first_deployment_succeeds_on_retry(
self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str
) -> None:
group = f"reliability-retry-{unique_marker()}"
small_context = create_always_picked_small_context_deployment(client.proxy, group)
resources.defer(lambda: client.proxy.delete_model(small_context))
backup = create_zero_weight_backup_deployment(client.proxy, group)
resources.defer(lambda: client.proxy.delete_model(backup))
resp = chat_override(
client.proxy,
scoped_key,
group,
oversized_prompt(unique_marker()),
override=RouterSettingsOverride(
num_retries=2,
model_group_retry_policy={group: {"BadRequestErrorRetries": 2}},
),
)
attempted = resp.headers.get("x-litellm-attempted-retries")
assert attempted is not None, "response is missing the x-litellm-attempted-retries header"
assert int(attempted) >= 1, (
f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never "
"opened on the timing-out deployment, so this proves nothing about retries"
)
content = content_of(resp)
finish_reason = finish_reason_of(resp)
completion_tokens = completion_tokens_of(resp) or 0
assert isinstance(content, str), (
f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})"
)
assert content or (finish_reason == "length" and completion_tokens > 0), (
f"the retry returned empty content with finish_reason={finish_reason!r}, "
f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget "
f"was spent on non-visible reasoning (body={resp.body[:300]})"
)
assert_retry_landed_on_backup(resp)

114
tests/e2e/test_proxy_client.py
View file

@ -1,9 +1,11 @@
"""Harness coverage for the model barrier that gates on every replica.
"""Harness coverage for the barriers that gate on every replica.
No proxy needed and no ``e2e`` marker: this pins that a model registered through
the control plane only counts as servable once every configured replica lists it
on /v1/models, which is what keeps a two-gateway stack from handing a test a
model that one gateway has not reloaded yet. The fakes are plain pollers and an
on /v1/models, and that a management write only counts as read back once every
replica's read satisfies the caller's predicate, which is what keeps a two-gateway
stack from handing a test a model or a key that one gateway has not caught up on
yet. The fakes are plain pollers standing in for each replica's transport plus an
injected clock, so nothing here monkeypatches anything.
"""
@ -12,18 +14,33 @@ from __future__ import annotations
from collections.abc import Iterable, Mapping
from dataclasses import dataclass
from itertools import chain, repeat
from types import MappingProxyType
from typing import Final
import pytest
from e2e_config import parse_replica_urls
from e2e_http import Success
from models import ModelListEntry, ModelsListResponse
from proxy_client import ModelsPoller, NotServableOn, Servable, await_servable_everywhere
from e2e_http import Result, Success
from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse
from proxy_client import (
Poller,
ConvergeOutcome,
Converged,
ModelsPoller,
NotConverged,
NotServableOn,
Servable,
await_converged_everywhere,
await_servable_everywhere,
first_lagging_replica,
converge_timeout_message,
)
MODEL: Final = "gpt-under-test"
TIMEOUT: Final = 10.0
INTERVAL: Final = 2.0
RPM_BEFORE_UPDATE: Final = 100
RPM_AFTER_UPDATE: Final = 200
@dataclass
@ -78,6 +95,91 @@ class TestAwaitServableEverywhere:
assert _await(pollers) == Servable()
def _key_info(rpm_limit: int) -> Success[KeyInfoResponse]:
return Success(status_code=200, data=KeyInfoResponse(info=KeyInfo(rpm_limit=rpm_limit)))
def _reads(results: Iterable[Result[KeyInfoResponse]]) -> Poller[Result[KeyInfoResponse]]:
it: Final = iter(results)
return lambda: next(it)
def _updated(result: Result[KeyInfoResponse]) -> bool:
return isinstance(result, Success) and result.data.info.rpm_limit == RPM_AFTER_UPDATE
def _converge(
pollers: Mapping[str, Poller[Result[KeyInfoResponse]]], clock: FakeClock
) -> Mapping[str, ConvergeOutcome[Result[KeyInfoResponse]]]:
return await_converged_everywhere(
pollers,
converged=_updated,
timeout=TIMEOUT,
interval=INTERVAL,
now=clock.now,
sleep=clock.sleep,
)
class TestAwaitConvergedEverywhere:
def test_waits_for_the_replica_that_lags_behind_the_write(self) -> None:
clock: Final = FakeClock()
pollers: Final = MappingProxyType(
{
"gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))),
"gateway-2": _reads(
chain(repeat(_key_info(RPM_BEFORE_UPDATE), 2), repeat(_key_info(RPM_AFTER_UPDATE)))
),
}
)
outcomes: Final = _converge(pollers, clock)
assert outcomes == {
"gateway-1": Converged(result=_key_info(RPM_AFTER_UPDATE)),
"gateway-2": Converged(result=_key_info(RPM_AFTER_UPDATE)),
}
assert first_lagging_replica(outcomes) is None
assert clock.elapsed == 2 * INTERVAL
def test_names_the_replica_that_never_converges_with_its_last_read(self) -> None:
clock: Final = FakeClock()
pollers: Final = MappingProxyType(
{
"gateway-1": _reads(repeat(_key_info(RPM_AFTER_UPDATE))),
"gateway-2": _reads(repeat(_key_info(RPM_BEFORE_UPDATE))),
}
)
outcomes: Final = _converge(pollers, clock)
assert first_lagging_replica(outcomes) == (
"gateway-2",
NotConverged(last_result=_key_info(RPM_BEFORE_UPDATE)),
)
assert clock.elapsed == TIMEOUT
message: Final = converge_timeout_message(
what="GET /key/info",
replica="gateway-2",
timeout=TIMEOUT,
last_result=_key_info(RPM_BEFORE_UPDATE),
)
assert "gateway-2" in message and "/key/info" in message and str(RPM_BEFORE_UPDATE) in message
def test_each_replica_gets_its_own_full_budget(self) -> None:
"""A replica that converges late must not eat into the next replica's budget: both
need most of the timeout here, so one shared deadline would starve the second."""
clock: Final = FakeClock()
slow: Final = chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE)))
pollers: Final = MappingProxyType(
{
"gateway-1": _reads(slow),
"gateway-2": _reads(
chain(repeat(_key_info(RPM_BEFORE_UPDATE), 3), repeat(_key_info(RPM_AFTER_UPDATE)))
),
}
)
outcomes: Final = _converge(pollers, clock)
assert first_lagging_replica(outcomes) is None
assert clock.elapsed == 2 * 3 * INTERVAL
class TestParseReplicaUrls:
def test_splits_and_trims_the_gateway_addresses(self) -> None:
raw: Final = " http://127.0.0.1:4010/, http://127.0.0.1:4011 "

190
tests/e2e/ui/fixtures/migratedPages.ts
View file

@ -1,50 +1,142 @@
/**
* Source of truth for the App Router migration E2E suites.
*
* Add an entry (legacy sidebar page id -> route segment) once a page's migration
* has MERGED to the branch under test. Consumers pick it up automatically:
* - migration smoke (tests/migration/migratedPages.spec.ts), via MIGRATED_E2E_SEGMENTS:
* default mount: npm run e2e:migration
* server-root-path mount: SERVER_ROOT_PATH=/<root> npm run e2e:migration:root
* - navigation specs that assert per-page URLs (tests/navigation/sidebar.spec.ts)
*
* Keep this in lockstep with MIGRATED_PAGES in src/utils/migratedPages.ts.
*/
export const MIGRATED_E2E_PAGES: Record<string, string> = {
"api-keys": "api-keys",
models: "models-and-endpoints",
api_ref: "api-reference",
"llm-playground": "playground",
projects: "projects",
"access-groups": "access-groups",
budgets: "budgets",
workflows: "workflows",
"guardrails-monitor": "guardrails-monitor",
"mcp-servers": "mcp-servers",
"search-tools": "search-tools",
"tag-management": "tag-management",
"vector-stores": "vector-stores",
memory: "memory",
policies: "policies",
guardrails: "guardrails",
prompts: "prompts",
"tool-policies": "tool-policies",
skills: "skills",
caching: "caching",
"cost-tracking": "cost-tracking",
"transform-request": "transform-request",
"ui-theme": "ui-theme",
logs: "logs",
"admin-panel": "admin-panel",
"logging-and-alerts": "logging-and-alerts",
"model-hub-table": "model-hub-table",
new_usage: "usage",
usage: "old-usage",
agents: "agents",
"router-settings": "router-settings",
users: "users",
teams: "teams",
organizations: "organizations",
};
export type MigratedPage = Readonly<{
segment: string;
linkName: string | RegExp;
group?: string;
content: Readonly<{ role: "heading" | "tab" | "button"; name: string }> | Readonly<{ text: string }>;
unlicensedText?: string;
}>;
export const MIGRATED_E2E_SEGMENTS: string[] = [...new Set(Object.values(MIGRATED_E2E_PAGES))];
export const MIGRATED_E2E_PAGES: Readonly<Record<string, MigratedPage>> = {
"api-keys": { segment: "api-keys", linkName: "Virtual Keys", content: { role: "heading", name: "Virtual Keys" } },
models: {
segment: "models-and-endpoints",
linkName: "Models + Endpoints",
content: { role: "heading", name: "Model Management" },
},
api_ref: {
segment: "api-reference",
linkName: "API Reference",
content: { role: "heading", name: "OpenAI Compatible Proxy: API Reference" },
},
"llm-playground": { segment: "playground", linkName: "Playground", content: { role: "tab", name: "Chat" } },
projects: {
segment: "projects",
linkName: /^Projects(?: Beta)?$/,
content: { role: "heading", name: "Projects" },
},
"access-groups": {
segment: "access-groups",
linkName: "Access Groups",
content: { role: "heading", name: "Access Groups" },
},
budgets: { segment: "budgets", linkName: "Budgets", content: { role: "heading", name: "Budgets" } },
workflows: {
segment: "workflows",
linkName: "Workflow Runs",
group: "Agentic",
content: { text: "Workflow Runs" },
},
"guardrails-monitor": {
segment: "guardrails-monitor",
linkName: "Guardrails Monitor",
content: { role: "heading", name: "Guardrails Monitor" },
},
"mcp-servers": {
segment: "mcp-servers",
linkName: "MCP Servers",
content: { role: "heading", name: "MCP Servers" },
},
"search-tools": {
segment: "search-tools",
linkName: "Search Tools",
group: "Tools",
content: { role: "heading", name: "Search Tools" },
},
"tag-management": {
segment: "tag-management",
linkName: "Tag Management",
group: "Experimental",
content: { role: "heading", name: "Tag Management" },
},
"vector-stores": {
segment: "vector-stores",
linkName: "Vector Stores",
group: "Tools",
content: { role: "heading", name: "Vector Store Management" },
},
memory: { segment: "memory", linkName: "Memory", group: "Agentic", content: { role: "heading", name: "Memory" } },
policies: { segment: "policies", linkName: "Policies", content: { role: "tab", name: "Policy Simulator" } },
guardrails: { segment: "guardrails", linkName: "Guardrails", content: { role: "tab", name: "Guardrails" } },
prompts: {
segment: "prompts",
linkName: "Prompts",
group: "Experimental",
content: { role: "button", name: "Add New Prompt" },
},
"tool-policies": {
segment: "tool-policies",
linkName: "Tool Policies",
group: "Tools",
content: { role: "heading", name: "Tool Policies" },
},
skills: { segment: "skills", linkName: "Skills", content: { role: "heading", name: "Skills" } },
caching: { segment: "caching", linkName: "Response Cache", content: { role: "tab", name: "Cache Settings" } },
"cost-tracking": {
segment: "cost-tracking",
linkName: "Cost Tracking",
group: "Settings",
content: { text: "Cost Tracking Settings" },
},
"transform-request": {
segment: "transform-request",
linkName: "API Playground",
group: "Experimental",
content: { role: "heading", name: "Playground" },
},
"ui-theme": {
segment: "ui-theme",
linkName: "UI Theme",
group: "Settings",
content: { role: "heading", name: "UI Theme Customization" },
},
logs: { segment: "logs", linkName: "Logs", content: { role: "heading", name: "Request Logs" } },
"admin-panel": {
segment: "admin-panel",
linkName: "Admin Settings",
group: "Settings",
content: { role: "heading", name: "Admin Access" },
},
"logging-and-alerts": {
segment: "logging-and-alerts",
linkName: "Logging & Alerts",
group: "Settings",
content: { role: "tab", name: "Logging Callbacks" },
},
"model-hub-table": {
segment: "model-hub-table",
linkName: "AI Hub",
content: { role: "heading", name: "AI Hub" },
},
new_usage: { segment: "usage", linkName: "Usage", content: { role: "heading", name: "Usage View" } },
usage: {
segment: "old-usage",
linkName: "Old Usage",
group: "Experimental",
content: { role: "tab", name: "All Up" },
},
agents: { segment: "agents", linkName: "Agents", group: "Agentic", content: { role: "heading", name: "Agents" } },
"router-settings": {
segment: "router-settings",
linkName: "Router Settings",
group: "Settings",
content: { role: "heading", name: "Routing Settings" },
},
users: { segment: "users", linkName: "Internal Users", content: { role: "tab", name: "Users" } },
teams: { segment: "teams", linkName: "Teams", content: { role: "heading", name: "Teams" } },
organizations: {
segment: "organizations",
linkName: "Organizations",
content: { text: "Click on an organization ID to view its details." },
unlicensedText: "This is a LiteLLM Enterprise feature, and requires a valid key to use. Get a trial key here.",
},
};

24
tests/e2e/ui/helpers/navigation.ts
View file

@ -1,5 +1,29 @@
import { Page } from "../fixtures/pages";
import { Page as PlaywrightPage, expect } from "@playwright/test";
import { UI_BASE_URL } from "../constants";
export const sidebarLink = (page: PlaywrightPage, name: string | RegExp) =>
page.getByRole("complementary").getByRole("link", { name, exact: true });
export async function clickSidebarLink(page: PlaywrightPage, name: string | RegExp, groupName?: string): Promise<void> {
const link = sidebarLink(page, name);
if (groupName && !(await link.isVisible())) {
const group = page.getByRole("complementary").getByRole("button", { name: groupName, exact: true });
await expect(group).toBeVisible();
if ((await group.getAttribute("aria-expanded")) === "false") {
await group.click();
}
}
await link.click();
}
export async function expectUiRoute(page: PlaywrightPage, segment: string): Promise<void> {
const root = (process.env.SERVER_ROOT_PATH ?? "").replace(/\/+$/, "");
const expected = new URL(`${root}/ui/${segment}`, UI_BASE_URL);
await expect(page, `navigate to ${expected.pathname}`).toHaveURL(
(url) => url.origin === expected.origin && url.pathname.replace(/\/+$/, "") === expected.pathname,
);
}
/**
* Navigates to a specific page using the page query parameter.

31
tests/e2e/ui/tests/migration/README.md
View file

@ -1,17 +1,25 @@
# App Router migration smoke
A growing E2E smoke for pages migrated from the legacy `?page=` switch to App
Router path routes. For each migrated page it clicks the page's sidebar link, checks
the URL is the path route and the page renders, reloads it, then clicks off to a
legacy page and back to confirm navigation still works. It runs in two situations:
the default mount and a non-root `SERVER_ROOT_PATH` mount.
Router path routes. For each page it clicks the sidebar link by its accessible
name, verifies the destination's content, reloads it, then visits Virtual Keys
and returns. It runs at the default mount and a non-root `SERVER_ROOT_PATH` mount
Link selection does not depend on `href` formatting. URL assertions compare the
origin and pathname, allowing a trailing slash, query string, and fragment while
rejecting another route or mount. Reloads must return a successful document,
and each journey must finish without uncaught browser errors
## Adding a page
When a page's migration merges, add its route segment to
`tests/e2e/ui/fixtures/migratedPages.ts` (keep it in lockstep with `MIGRATED_PAGES`
in `ui/litellm-dashboard/src/utils/migratedPages.ts`). Both suites pick it up
automatically.
Add an entry to `tests/e2e/ui/fixtures/migratedPages.ts`, keyed by the legacy page
ID. Specify its route segment, accessible link name, sidebar group if collapsed,
and distinctive visible content such as a heading or tab. Keep expectations
independent of the application's route table so an incorrect destination fails
the test. Both navigation suites use this fixture
For a licensed-only page, `unlicensedText` describes the expected upgrade notice.
The authenticated session's license claim determines which content must render
## Running
@ -31,4 +39,9 @@ SERVER_ROOT_PATH=/litellm npm run e2e:migration:root
```
`globalSetup` logs in once per role; the admin storage state is reused for these
tests. Under a non-root mount it logs in at `${SERVER_ROOT_PATH}/ui/login`.
tests. Under a non-root mount it logs in at `${SERVER_ROOT_PATH}/ui/login`
`tests/navigation/sidebar.spec.ts` also checks the navigation helpers against
equivalent link formats on the live dashboard and a deep link containing a query
string and fragment. The link-format cases change only the rendered `href`
attribute to exercise the locator contract; destination pages and APIs remain live

126
tests/e2e/ui/tests/migration/migratedPages.spec.ts
View file

@ -1,105 +1,73 @@
import { test, expect, type Page } from "@playwright/test";
import { MIGRATED_E2E_SEGMENTS } from "../../fixtures/migratedPages";
import { MIGRATED_E2E_PAGES, type MigratedPage } from "../../fixtures/migratedPages";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { dismissFeedbackPopup } from "../../helpers/navigation";
import { clickSidebarLink, dismissFeedbackPopup, expectUiRoute, sidebarLink } from "../../helpers/navigation";
import { proxyIsPremium } from "../../helpers/premium";
/**
* App Router migration smoke as a user journey: start where the proxy lands you,
* click a migrated page in the sidebar, confirm it routed and rendered, reload it
* (the check a wrong server_root_path breaks), bounce to a legacy page and back,
* and, once two pages are migrated, navigate directly between two migrated pages.
*
* Driven by MIGRATED_E2E_SEGMENTS, so it grows as pages are migrated. Set
* SERVER_ROOT_PATH (e.g. "/litellm") to exercise the non-root mount; leave it
* unset for the default mount. Boot the proxy with the matching value first.
*/
const ROOT = process.env.SERVER_ROOT_PATH ?? "";
const ROOT = (process.env.SERVER_ROOT_PATH ?? "").replace(/\/+$/, "");
const apiKeys = MIGRATED_E2E_PAGES["api-keys"];
const esc = (s: string) => s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
const pathRe = (segment: string) => new RegExp(`${esc(ROOT)}/ui/${esc(segment)}/?($|\\?)`);
// Scope nav lookups to the sidebar (a `complementary` landmark). The top bar
// now renders a breadcrumb whose current-page item is also a "Virtual Keys"
// link, so an unscoped locator would match two elements.
const sidebar = (page: Page) => page.getByRole("complementary");
const virtualKeysLink = (page: Page) => sidebar(page).getByRole("link", { name: "Virtual Keys", exact: true });
/** The dashboard shell is present (sidebar rendered); page didn't 404 / crash. */
async function expectRendered(page: Page) {
await expect(virtualKeysLink(page)).toBeVisible({ timeout: 20_000 });
async function expectContent(page: Page, destination: MigratedPage): Promise<void> {
await expect(sidebarLink(page, apiKeys.linkName)).toBeVisible({ timeout: 20_000 });
const main = page.getByRole("main");
if (destination.unlicensedText && !proxyIsPremium()) {
await expect(main.getByText(destination.unlicensedText, { exact: true })).toBeVisible();
return;
}
const content = destination.content;
const landmark =
"role" in content
? main.getByRole(content.role, { name: content.name, exact: true })
: main.getByText(content.text, { exact: true });
await expect(landmark).toBeVisible();
}
/**
* Click a migrated page's sidebar link. Migrated items render as <a href=".../ui/<segment>">;
* nested ones live under collapsible groups whose children only render while the
* group is open, so expand collapsed groups until the link is clickable.
*/
async function clickSidebar(page: Page, segment: string) {
const link = sidebar(page).locator(`a[href$="/ui/${segment}"]`).first();
const collapsedGroups = sidebar(page).getByRole("button", { expanded: false });
for (let i = 0; i < 8 && !(await link.isVisible().catch(() => false)); i++) {
const stillCollapsed = await collapsedGroups.count();
if (stillCollapsed === 0) break;
await collapsedGroups.first().click();
await expect(collapsedGroups).toHaveCount(stillCollapsed - 1);
}
await link.click();
async function navigateToDestination(page: Page, destination: MigratedPage): Promise<void> {
await clickSidebarLink(page, destination.linkName, destination.group);
await expectUiRoute(page, destination.segment);
await dismissFeedbackPopup(page);
await expectContent(page, destination);
}
test.use({ storageState: ADMIN_STORAGE_PATH });
test.describe("App Router migrated pages", () => {
for (const segment of MIGRATED_E2E_SEGMENTS) {
test(`${segment}: sidebar nav, reload, and round-trip via the api-keys landing`, async ({ page }) => {
for (const destination of Object.values(MIGRATED_E2E_PAGES)) {
test(`${destination.segment}: sidebar nav, reload, and round-trip via the api-keys landing`, async ({ page }) => {
const pageErrors: string[] = [];
page.on("pageerror", (e) => pageErrors.push(String(e)));
page.on("pageerror", (error) => pageErrors.push(String(error)));
// 1. Start where the proxy lands us.
await page.goto(`${ROOT}/ui/`);
const landing = await page.goto(`${ROOT}/ui/`);
expect(landing?.ok(), "dashboard document loads successfully").toBe(true);
await dismissFeedbackPopup(page);
await expectRendered(page);
await expectContent(page, apiKeys);
// 2. Click the migrated page in the sidebar -> path route + rendered.
await clickSidebar(page, segment);
await expect(page).toHaveURL(pathRe(segment));
await expectRendered(page);
// 3. Reload the path route directly; a wrong server_root_path 404s here.
await page.reload();
await navigateToDestination(page, destination);
const reloaded = await page.reload();
expect(reloaded?.ok(), `${destination.segment} document loads on reload`).toBe(true);
await expectUiRoute(page, destination.segment);
await dismissFeedbackPopup(page);
await expect(page).toHaveURL(pathRe(segment));
await expectRendered(page);
// 4. Click the Virtual Keys sidebar link to the api-keys landing (now a path route), then back.
await virtualKeysLink(page).click();
await expect(page).toHaveURL(pathRe("api-keys"));
await dismissFeedbackPopup(page);
await expectRendered(page);
// 5. Click back to the migrated page.
await clickSidebar(page, segment);
await expect(page).toHaveURL(pathRe(segment));
await expectRendered(page);
expect(pageErrors, `page errors during ${segment} journey`).toEqual([]);
await expectContent(page, destination);
await navigateToDestination(page, apiKeys);
await navigateToDestination(page, destination);
expect(pageErrors, `page errors during ${destination.segment} journey`).toEqual([]);
});
}
test("navigates directly between two migrated pages", async ({ page }) => {
test.skip(MIGRATED_E2E_SEGMENTS.length < 2, "needs >= 2 migrated pages");
const [first, second] = MIGRATED_E2E_SEGMENTS;
const pageErrors: string[] = [];
page.on("pageerror", (e) => pageErrors.push(String(e)));
page.on("pageerror", (error) => pageErrors.push(String(error)));
await page.goto(`${ROOT}/ui/`);
const landing = await page.goto(`${ROOT}/ui/`);
expect(landing?.ok(), "dashboard document loads successfully").toBe(true);
await dismissFeedbackPopup(page);
await expectContent(page, apiKeys);
await clickSidebar(page, first);
await expect(page).toHaveURL(pathRe(first));
await expectRendered(page);
await clickSidebar(page, second);
await expect(page).toHaveURL(pathRe(second));
await expectRendered(page);
// Back to the first migrated page.
await clickSidebar(page, first);
await expect(page).toHaveURL(pathRe(first));
await expectRendered(page);
expect(pageErrors, "page errors during migrated -> migrated nav").toEqual([]);
for (const destination of [apiKeys, MIGRATED_E2E_PAGES.models, apiKeys]) {
await navigateToDestination(page, destination);
}
expect(pageErrors, "page errors during migrated page navigation").toEqual([]);
});
});

59
tests/e2e/ui/tests/navigation/sidebar.spec.ts
View file

@ -3,7 +3,13 @@ import { Role } from "../../fixtures/roles";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { Page } from "../../fixtures/pages";
import { menuLabelToPage } from "../../fixtures/menuMappings";
import { navigateToPage } from "../../helpers/navigation";
import {
clickSidebarLink,
dismissFeedbackPopup,
expectUiRoute,
navigateToPage,
sidebarLink,
} from "../../helpers/navigation";
import { MIGRATED_E2E_PAGES } from "../../fixtures/migratedPages";
import type { Page as PlaywrightPage } from "@playwright/test";
@ -11,7 +17,7 @@ const sidebarButtons = {
[Role.ProxyAdmin]: [
"Virtual Keys",
"Playground",
"Models",
"Models + Endpoints",
"Usage",
"Teams",
"Internal Users",
@ -22,9 +28,9 @@ const sidebarButtons = {
/** Migrated pages live at a path route; legacy pages keep the ?page= query param. */
async function expectPageUrl(page: PlaywrightPage, pageKey: string): Promise<void> {
const migratedSegment = MIGRATED_E2E_PAGES[pageKey];
if (migratedSegment) {
await expect(page).toHaveURL(new RegExp(`/ui/${migratedSegment}/?($|\\?)`));
const migratedPage = MIGRATED_E2E_PAGES[pageKey];
if (migratedPage) {
await expectUiRoute(page, migratedPage.segment);
} else {
await expect(page).toHaveURL(new RegExp(`[?&]page=${pageKey}(&|$)`));
}
@ -51,12 +57,7 @@ for (const { role, storage } of roles) {
throw new Error(`No page mapping found for menu label: ${buttonLabel}`);
}
// Sidebar items are links inside the `complementary` landmark; scoping
// there avoids the top-bar breadcrumb, which also links the page name.
const tab = page.getByRole("complementary").getByRole("link", { name: buttonLabel });
await expect(tab).toBeVisible();
await tab.click();
await clickSidebarLink(page, buttonLabel);
await expectPageUrl(page, expectedPage);
}
@ -81,5 +82,41 @@ for (const { role, storage } of roles) {
await navigateToPage(page, Page.LlmPlayground);
await expectPageUrl(page, Page.LlmPlayground);
});
for (const format of ["without trailing slash", "absolute with query and fragment", "relative"] as const) {
test(`sidebar locator tolerates hrefs ${format}`, async ({ page }) => {
await page.goto("/ui/");
await dismissFeedbackPopup(page);
const link = sidebarLink(page, "Models + Endpoints");
await expect(link).toBeVisible();
const destination = new URL("/ui/models-and-endpoints/", page.url());
const href =
format === "without trailing slash"
? destination.pathname.replace(/\/$/, "")
: format === "relative"
? "./models-and-endpoints/"
: `${destination.href}?source=navigation-smoke#overview`;
await link.evaluate((element, value) => element.setAttribute("href", value), href);
await expect(link).toHaveAttribute("href", href);
await clickSidebarLink(page, "Models + Endpoints");
await expectUiRoute(page, "models-and-endpoints");
await expect(
page.getByRole("main").getByRole("heading", { name: "Model Management", exact: true }),
).toBeVisible();
});
}
test("route assertion tolerates a query string and fragment on a deep link", async ({ page }) => {
const response = await page.goto("/ui/models-and-endpoints/?source=navigation-smoke#overview");
expect(response?.ok()).toBe(true);
await expectUiRoute(page, "models-and-endpoints");
await expect(
page.getByRole("main").getByRole("heading", { name: "Model Management", exact: true }),
).toBeVisible();
expect(new URL(page.url()).search).toBe("?source=navigation-smoke");
expect(new URL(page.url()).hash).toBe("#overview");
});
});
}

145
tests/e2e/ui/tests/proxy-admin/teamGuardrailRemoval.spec.ts Normal file
View file

@ -0,0 +1,145 @@
import { test, expect, type APIRequestContext, type Page as PlaywrightPage } from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { Page } from "../../fixtures/pages";
import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation";
import { proxyIsPremium } from "../../helpers/premium";
import { readBack } from "../../helpers/roundTrip";
import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic";
const auth = () => ({ Authorization: `Bearer ${masterKey()}` });
async function guardrailId(request: APIRequestContext, name: string): Promise<string | undefined> {
const res = await request.get("/v2/guardrails/list", { headers: auth() });
expect(res.ok(), `GET /v2/guardrails/list (${res.status()})`).toBe(true);
const rows = (await res.json()).guardrails as { guardrail_id: string; guardrail_name: string | null }[];
return rows.find((row) => row.guardrail_name === name)?.guardrail_id;
}
async function teamGuardrails(page: PlaywrightPage, teamId: string): Promise<string[]> {
const body = await readBack<{ team_info: { metadata: { guardrails?: string[] } | null } }>(
page,
`/team/info?team_id=${encodeURIComponent(teamId)}`,
);
return body.team_info.metadata?.guardrails ?? [];
}
async function keywordPromptStatus(request: APIRequestContext, apiKey: string, keyword: string): Promise<number> {
const res = await request.post("/v1/chat/completions", {
headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" },
data: { model: CHAT_MODEL_A, messages: [{ role: "user", content: `please tell me about ${keyword}` }] },
});
return res.status();
}
test.describe("Proxy Admin - Team guardrail removal", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
test("Clearing a team's only guardrail on the Settings tab lets blocked traffic through again", async ({
page,
request,
}) => {
test.skip(!proxyIsPremium(), "proxy under test is unlicensed, so team guardrails are premium-gated");
const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`;
const guardrailName = `e2e-team-guardrail-${stamp}`;
const bannedKeyword = `e2eteamban${stamp}`;
const teamAlias = `e2e-guardrail-team-${stamp}`;
let teamId = "";
let teamKey = "";
try {
const guardrailRes = await request.post("/guardrails", {
headers: auth(),
data: {
guardrail: {
guardrail_name: guardrailName,
litellm_params: {
guardrail: "litellm_content_filter",
mode: "pre_call",
default_on: false,
blocked_words: [{ keyword: bannedKeyword, action: "BLOCK" }],
},
},
},
});
expect(
guardrailRes.ok(),
`POST /guardrails failed (${guardrailRes.status()}): ${await guardrailRes.text()}`,
).toBe(true);
const teamRes = await request.post("/team/new", {
headers: auth(),
data: { team_alias: teamAlias, models: [CHAT_MODEL_A], metadata: { guardrails: [guardrailName] } },
});
expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true);
teamId = (await teamRes.json()).team_id as string;
const keyRes = await request.post("/key/generate", { headers: auth(), data: { team_id: teamId } });
expect(keyRes.ok(), `POST /key/generate failed (${keyRes.status()}): ${await keyRes.text()}`).toBe(true);
teamKey = (await keyRes.json()).key as string;
await expect
.poll(async () => keywordPromptStatus(request, teamKey, bannedKeyword), {
message: "the team's guardrail never started refusing the banned keyword",
timeout: 60_000,
})
.toBe(400);
await navigateToPage(page, Page.Teams);
await dismissFeedbackPopup(page);
await clickTeamId(page, teamId);
await page.getByRole("tab", { name: "Settings" }).click();
await page.getByRole("button", { name: "Edit Settings" }).click();
const chip = page.locator('[data-slot="combobox-chip"]').filter({ hasText: guardrailName });
await expect(chip).toBeVisible({ timeout: 10_000 });
await chip.locator('[data-slot="combobox-chip-remove"]').click();
await expect(chip).toHaveCount(0, { timeout: 10_000 });
await page.getByRole("button", { name: "Save Changes" }).click();
await expect
.poll(async () => teamGuardrails(page, teamId), {
message: "the team still carries a guardrail in /team/info after the save",
timeout: 20_000,
})
.toEqual([]);
await page.reload();
await page.getByRole("tab", { name: "Settings" }).click();
await page.getByRole("button", { name: "Edit Settings" }).click();
await expect(page.getByRole("combobox", { name: "Select guardrails" })).toBeVisible({ timeout: 15_000 });
await expect(
page.locator('[data-slot="combobox-chip"]').filter({ hasText: guardrailName }),
"the removed guardrail is gone from the Settings tab after a reload",
).toHaveCount(0);
await expect
.poll(async () => keywordPromptStatus(request, teamKey, bannedKeyword), {
message: "the team key is still refused for a keyword whose guardrail was removed",
timeout: 60_000,
})
.toBe(200);
const served = await request.post("/v1/chat/completions", {
headers: { Authorization: `Bearer ${teamKey}`, "Content-Type": "application/json" },
data: {
model: CHAT_MODEL_A,
messages: [{ role: "user", content: `please tell me about ${bannedKeyword}` }],
},
});
expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT);
} finally {
if (teamKey) {
await request.post("/key/delete", { headers: auth(), data: { keys: [teamKey] } });
}
if (teamId) {
await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } });
}
const id = await guardrailId(request, guardrailName);
if (id) {
await request.delete(`/guardrails/${id}`, { headers: auth() });
}
}
});
});

129
tests/e2e/ui/tests/proxy-admin/teamMemberEdit.spec.ts Normal file
View file

@ -0,0 +1,129 @@
import { test, expect, type Page as PlaywrightPage } from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { Page } from "../../fixtures/pages";
import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation";
import { readBack } from "../../helpers/roundTrip";
import { CHAT_MODEL_A, masterKey } from "../../helpers/traffic";
interface TeamInfoResponse {
team_info: {
models: string[];
members_with_roles: { user_id?: string; role?: string }[];
};
team_memberships: {
user_id: string;
litellm_budget_table: { max_budget: number | null } | null;
}[];
}
const auth = () => ({ Authorization: `Bearer ${masterKey()}` });
async function teamInfo(page: PlaywrightPage, teamId: string): Promise<TeamInfoResponse> {
return readBack<TeamInfoResponse>(page, `/team/info?team_id=${encodeURIComponent(teamId)}`);
}
function roleOf(info: TeamInfoResponse, userId: string): string | undefined {
return info.team_info.members_with_roles.find((member) => member.user_id === userId)?.role;
}
function budgetOf(info: TeamInfoResponse, userId: string): number | null | undefined {
return info.team_memberships.find((membership) => membership.user_id === userId)?.litellm_budget_table?.max_budget;
}
function otherMembers(info: TeamInfoResponse, userId: string): string[] {
return info.team_info.members_with_roles
.filter((member) => member.user_id !== userId)
.map((member) => `${member.user_id}:${member.role}`)
.sort();
}
test.describe("Proxy Admin - Team member edit", () => {
test.use({ storageState: ADMIN_STORAGE_PATH });
const createdTeams: string[] = [];
const createdUsers: string[] = [];
test.afterEach(async ({ request }) => {
for (const teamId of createdTeams.splice(0)) {
await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } });
}
for (const userId of createdUsers.splice(0)) {
await request.post("/user/delete", { headers: auth(), data: { user_ids: [userId] } });
}
});
test("Editing a member's role and per-member budget persists and survives a reload", async ({ page, request }) => {
const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`;
const memberId = `e2e-member-edit-${stamp}`;
const teamAlias = `e2e-member-edit-team-${stamp}`;
createdUsers.push(memberId);
const created = await request.post("/user/new", {
headers: auth(),
data: { user_id: memberId, user_role: "internal_user", auto_create_key: false },
});
expect(created.ok(), `POST /user/new failed (${created.status()}): ${await created.text()}`).toBe(true);
const teamRes = await request.post("/team/new", {
headers: auth(),
data: {
team_alias: teamAlias,
models: [CHAT_MODEL_A],
members_with_roles: [{ user_id: memberId, role: "admin" }],
},
});
expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true);
const teamId = (await teamRes.json()).team_id as string;
createdTeams.push(teamId);
const before = await teamInfo(page, teamId);
expect(roleOf(before, memberId), "the member starts out as a team admin").toBe("admin");
expect(budgetOf(before, memberId) ?? null, "the member starts out with no per-member budget").toBeNull();
expect(
otherMembers(before, memberId).length,
"the team has another member for the edit to leave alone",
).toBeGreaterThan(0);
await navigateToPage(page, Page.Teams);
await dismissFeedbackPopup(page);
await clickTeamId(page, teamId);
await page.getByRole("tab", { name: "Members" }).click();
const memberRow = page.locator("tr", { hasText: memberId }).first();
await expect(memberRow).toBeVisible({ timeout: 10_000 });
await memberRow.getByTestId("edit-member").click();
const modal = page.getByRole("dialog", { name: "Edit Member" });
await expect(modal).toBeVisible({ timeout: 10_000 });
await modal.getByLabel(/^Role/).click();
await page.getByRole("option", { name: "User", exact: true }).click();
await modal.getByLabel(/Team Member Budget \(USD\)/).fill("5");
await modal.getByRole("button", { name: "Save Changes" }).click();
await expect(page.getByText("Team member updated successfully").first()).toBeVisible({ timeout: 10_000 });
await expect
.poll(
async () => {
const info = await teamInfo(page, teamId);
return [roleOf(info, memberId), budgetOf(info, memberId)];
},
{ message: "the member's role and budget never landed in /team/info", timeout: 20_000 },
)
.toEqual(["user", 5]);
await page.reload();
await page.getByRole("tab", { name: "Members" }).click();
const reloadedRow = page.locator("tr", { hasText: memberId }).first();
await expect(reloadedRow).toBeVisible({ timeout: 15_000 });
await expect(reloadedRow.getByText("user", { exact: true }), "role shown after a reload").toBeVisible();
await expect(reloadedRow.getByText("$5.00"), "per-member budget shown after a reload").toBeVisible();
const after = await teamInfo(page, teamId);
expect(after.team_info.models, "model access untouched by a member edit").toEqual(before.team_info.models);
expect(otherMembers(after, memberId), "the rest of the roster untouched by a member edit").toEqual(
otherMembers(before, memberId),
);
});
});

177
tests/e2e/ui/tests/team-admin/memberPermissions.spec.ts Normal file
View file

@ -0,0 +1,177 @@
import { test, expect, type Browser, type BrowserContext, type Page as PlaywrightPage } from "@playwright/test";
import { Page } from "../../fixtures/pages";
import { navigateToPage, dismissFeedbackPopup, clickTeamId } from "../../helpers/navigation";
import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, masterKey } from "../../helpers/traffic";
const PASSWORD = "E2e-Member-Perms-Pass-1!";
const auth = () => ({ Authorization: `Bearer ${masterKey()}` });
async function sessionKey(page: PlaywrightPage): Promise<string> {
const cookie = (await page.context().cookies()).find((candidate) => candidate.name === "token");
expect(cookie?.value, "logged-in session carries a token cookie").toBeTruthy();
const payload = JSON.parse(Buffer.from(cookie!.value.split(".")[1], "base64url").toString("utf-8")) as {
key?: string;
};
expect(payload.key, "session JWT carries the virtual key the dashboard calls with").toMatch(/^sk-/);
return payload.key!;
}
async function signIn(browser: Browser, email: string): Promise<BrowserContext> {
const context = await browser.newContext({ storageState: { cookies: [], origins: [] } });
const page = await context.newPage();
await page.goto("/ui/login");
await page.getByPlaceholder("Enter your username").fill(email);
await page.getByPlaceholder("Enter your password").fill(PASSWORD);
await page.getByRole("button", { name: "Login", exact: true }).click();
await expect(page.locator("a", { hasText: "Virtual Keys" })).toBeVisible({ timeout: 30_000 });
await dismissFeedbackPopup(page);
return context;
}
test.describe("Team Admin - Member permissions", () => {
test.use({ storageState: { cookies: [], origins: [] } });
test("Granting /key/generate lets a plain member create a team key that serves traffic", async ({
browser,
request,
}) => {
const stamp = `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`;
const adminId = `e2e-perm-admin-${stamp}`;
const memberId = `e2e-perm-member-${stamp}`;
const adminEmail = `${adminId}@test.local`;
const memberEmail = `${memberId}@test.local`;
const teamAlias = `e2e-perm-team-${stamp}`;
const createUser = async (userId: string, email: string): Promise<void> => {
const created = await request.post("/user/new", {
headers: auth(),
data: { user_id: userId, user_email: email, user_role: "internal_user", auto_create_key: false },
});
expect(created.ok(), `POST /user/new for ${userId} (${created.status()}): ${await created.text()}`).toBe(true);
const password = await request.post("/user/update", {
headers: auth(),
data: { user_id: userId, password: PASSWORD },
});
expect(password.ok(), `POST /user/update for ${userId} (${password.status()})`).toBe(true);
};
let teamId = "";
const createdKeys: string[] = [];
const contexts: BrowserContext[] = [];
try {
await createUser(adminId, adminEmail);
await createUser(memberId, memberEmail);
const teamRes = await request.post("/team/new", {
headers: auth(),
data: {
team_alias: teamAlias,
models: [CHAT_MODEL_A],
members_with_roles: [
{ user_id: adminId, role: "admin" },
{ user_id: memberId, role: "user" },
],
},
});
expect(teamRes.ok(), `POST /team/new failed (${teamRes.status()}): ${await teamRes.text()}`).toBe(true);
teamId = (await teamRes.json()).team_id as string;
const memberContext = await signIn(browser, memberEmail);
contexts.push(memberContext);
const memberPage = memberContext.pages()[0];
const memberSessionKey = await sessionKey(memberPage);
const refused = await memberPage.request.post("/key/generate", {
headers: { Authorization: `Bearer ${memberSessionKey}`, "Content-Type": "application/json" },
data: { team_id: teamId, key_alias: `e2e-perm-denied-${stamp}` },
});
expect(refused.status(), "a plain member cannot mint a team key before the grant").toBe(401);
expect(await refused.text()).toContain("/key/generate");
const adminContext = await signIn(browser, adminEmail);
contexts.push(adminContext);
const adminPage = adminContext.pages()[0];
await navigateToPage(adminPage, Page.Teams);
await clickTeamId(adminPage, teamId);
await adminPage.getByRole("tab", { name: "Member Permissions" }).click();
for (const route of ["/key/generate", "/key/update"]) {
await adminPage.getByRole("row").filter({ hasText: route }).getByRole("checkbox").check();
}
await adminPage.getByRole("button", { name: "Save Changes" }).click();
await expect(adminPage.getByText("Permissions updated successfully").first()).toBeVisible({ timeout: 10_000 });
await expect
.poll(
async () => {
const res = await request.get(`/team/permissions_list?team_id=${encodeURIComponent(teamId)}`, {
headers: auth(),
});
if (!res.ok()) return [];
return ((await res.json()).team_member_permissions ?? []) as string[];
},
{ message: "the granted permissions never landed in /team/permissions_list", timeout: 20_000 },
)
.toEqual(expect.arrayContaining(["/key/generate", "/key/update"]));
const keyAlias = `e2e-perm-key-${stamp}`;
await navigateToPage(memberPage, Page.ApiKeys);
await memberPage.getByRole("button", { name: /Create New Key/i }).click();
await expect(memberPage.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 });
await memberPage.getByLabel(/Key Name/).fill(keyAlias);
const teamSelect = memberPage.getByTestId("team-dropdown").getByRole("combobox");
await teamSelect.click();
await memberPage.keyboard.type(teamAlias);
await memberPage.getByRole("option", { name: teamAlias }).first().click();
await memberPage.getByRole("combobox", { name: "Select models" }).click();
await memberPage.getByRole("option", { name: "All Team Models", exact: true }).click();
await memberPage.keyboard.press("Escape");
await memberPage.getByRole("button", { name: "Create Key", exact: true }).click();
const saveDialog = memberPage.getByRole("dialog", { name: "Save your Key" });
await expect(saveDialog).toBeVisible({ timeout: 15_000 });
const apiKey = (await saveDialog.locator("pre").innerText()).trim();
expect(apiKey).toMatch(/^sk-/);
createdKeys.push(apiKey);
await memberPage.keyboard.press("Escape");
await expect
.poll(
async () => {
const res = await request.get(
`/key/list?team_id=${encodeURIComponent(teamId)}&return_full_object=true&size=100`,
{ headers: auth() },
);
if (!res.ok()) return null;
const row = ((await res.json()).keys as Record<string, unknown>[]).find(
(candidate) => candidate.key_alias === keyAlias,
);
return row ? [row.user_id, row.team_id] : null;
},
{ message: `key ${keyAlias} never appeared on the team with the member as its owner`, timeout: 20_000 },
)
.toEqual([memberId, teamId]);
const served = await request.post("/v1/chat/completions", {
headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" },
data: { model: CHAT_MODEL_A, messages: [{ role: "user", content: `member key ping ${stamp}` }] },
});
expect(served.status(), "the delegated key is a real key the gateway serves").toBe(200);
expect((await served.json()).choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT);
} finally {
for (const context of contexts) {
await context.close();
}
for (const key of createdKeys) {
await request.post("/key/delete", { headers: auth(), data: { keys: [key] } });
}
if (teamId) {
await request.post("/team/delete", { headers: auth(), data: { team_ids: [teamId] } });
}
await request.post("/user/delete", { headers: auth(), data: { user_ids: [adminId, memberId] } });
}
});
});

1
tests/rust-python-harness/shared/parity/compare.py
View file

@ -78,3 +78,4 @@ def assert_parity(baseline: Execution, candidate: Execution, baseline_user_agent
validate_harness(baseline, candidate, baseline_user_agent)
assert_request_parity(baseline.requests, candidate.requests)
assert_value_parity(baseline.report, candidate.report)
assert_value_parity(baseline.callbacks, candidate.callbacks, path="$.callbacks")

21
tests/rust-python-harness/shared/parity/models.py
View file

@ -37,6 +37,25 @@ class SDKError(BaseModel):
llm_provider: str | None
class CallbackObservation(BaseModel):
model_config = ConfigDict(frozen=True)
hook: Literal[
"log_success_event",
"async_log_success_event",
"log_failure_event",
"async_log_failure_event",
]
phase: Literal["success", "failure"]
model: str | None
call_type: str | None
litellm_call_id: str | None
metadata: JsonValue
kwargs: JsonValue
payload: JsonValue
error: SDKError | None
class SDKJsonChunk(BaseModel):
model_config = ConfigDict(frozen=True)
@ -119,6 +138,7 @@ class Execution(BaseModel):
requests: tuple[CapturedRequest, ...]
report: SDKReport
callbacks: tuple[CallbackObservation, ...] | None = None
class SDKCommand(BaseModel):
@ -133,6 +153,7 @@ class WorkerSuccess(BaseModel):
status: Literal["ok"] = "ok"
report: SDKReport
callbacks: tuple[CallbackObservation, ...] | None = None
class WorkerFailure(BaseModel):

10
tests/rust-python-harness/shared/parity/runner.py
View file

@ -39,9 +39,7 @@ class SubprocessRunner:
return (
sys.executable,
"-m",
".".join(
self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts
),
".".join(self.entrypoint.resolve().relative_to(PROJECT_ROOT).with_suffix("").parts),
"--parity-worker",
provider_url,
)
@ -113,7 +111,11 @@ class SubprocessWorker:
)
assert isinstance(result, WorkerSuccess)
try:
return Execution(requests=self.provider.take_requests(len(responses)), report=result.report)
return Execution(
requests=self.provider.take_requests(len(responses)),
report=result.report,
callbacks=result.callbacks,
)
except AssertionError:
self.provider.reset()
raise

52
tests/rust-python-harness/shared/parity/test_parity.py
View file

@ -7,7 +7,7 @@ import pytest
from pydantic import BaseModel, ConfigDict, JsonValue, PrivateAttr
from .compare import assert_model_parity, assert_parity
from .models import CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report
from .models import CallbackObservation, CapturedRequest, Execution, SDKError, SDKSuccess, sdk_error_report
SENTINEL: Final = "python-parity-fallback"
@ -69,6 +69,56 @@ def test_parity_rejects_response_difference() -> None:
assert_parity(python, rust, SENTINEL)
def test_parity_distinguishes_unobserved_callbacks_from_zero_events() -> None:
python: Final = _execution(user_agent=SENTINEL)
rust: Final = _execution(user_agent="litellm-rust").model_copy(update={"callbacks": ()})
with pytest.raises(AssertionError, match=r"\$\.callbacks"):
assert_parity(python, rust, SENTINEL)
def test_parity_rejects_callback_payload_difference() -> None:
observation: Final = CallbackObservation(
hook="log_success_event",
phase="success",
model="test-model",
call_type="ocr",
litellm_call_id="test-call",
metadata={"profile": "success"},
kwargs={"model": "test-model"},
payload={"model": "test-model", "pages": []},
error=None,
)
python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)})
rust: Final = _execution(user_agent="litellm-rust").model_copy(
update={"callbacks": (observation.model_copy(update={"payload": {"model": "changed", "pages": []}}),)}
)
with pytest.raises(AssertionError, match=r"\$\.callbacks"):
assert_parity(python, rust, SENTINEL)
def test_parity_rejects_callback_kwargs_difference() -> None:
observation: Final = CallbackObservation(
hook="log_success_event",
phase="success",
model="test-model",
call_type="ocr",
litellm_call_id="test-call",
metadata={"profile": "success"},
kwargs={"model": "test-model"},
payload={"model": "test-model", "pages": []},
error=None,
)
python: Final = _execution(user_agent=SENTINEL).model_copy(update={"callbacks": (observation,)})
rust: Final = _execution(user_agent="litellm-rust").model_copy(
update={"callbacks": (observation.model_copy(update={"kwargs": {"model": "changed"}}),)}
)
with pytest.raises(AssertionError, match=r"\$\.callbacks"):
assert_parity(python, rust, SENTINEL)
def test_parity_rejects_error_difference() -> None:
python: Final = Execution(
requests=(),

3
tests/rust-python-harness/shared/tracing/native.py
View file

@ -19,7 +19,8 @@ class _TraceEventPayload(BaseModel):
class TraceResponsePayload(BaseModel):
model_config = ConfigDict(strict=True, extra="forbid")
response: object
response: object = None
error: str | None = None
trace: tuple[_TraceEventPayload, ...] | list[_TraceEventPayload]

4
tests/rust-python-harness/strategies/e2e_parity/__init__.py
View file

@ -18,8 +18,8 @@ CASES: Final[tuple[CaseDefinition, ...]] = (
coverage=Coverage.PARTIAL,
module="tests.rust-python-harness.strategies.e2e_parity.sdk.ocr.test_sdk_parity",
note=(
"Recorded sync/async SDK parity; invalid-model provider errors differ, "
"and Reducto lacks a Rust contract."
"Recorded sync/async SDK parity with focused success/error callback profiles; "
"Reducto lacks a Rust contract, and known provider parity gaps remain."
),
),
surface="sdk",

295
tests/rust-python-harness/strategies/e2e_parity/sdk/ocr/test_sdk_parity.py
View file

@ -1,23 +1,31 @@
from __future__ import annotations
import asyncio
import datetime
import queue
import sys
import tempfile
import time
import traceback
from collections.abc import Callable, Coroutine, Generator
from collections.abc import Callable, Coroutine, Generator, Mapping
from contextlib import contextmanager
from enum import Enum
from functools import partial
from pathlib import Path
from typing import Annotated, Final, Literal, cast
from urllib.parse import urlsplit, urlunsplit
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from .....shared.parity.compare import assert_parity
from .....shared.parity.fixtures.store import fixture_id, recorded_fixtures
from .....shared.parity.models import (
JSON_VALUE_ADAPTER,
CallbackObservation,
Execution,
SDKCommand,
SDKError,
SDKReport,
@ -39,6 +47,9 @@ from .fixtures.config import configured_fixture_directory
from .fixtures.models import OcrParityCase, OcrSdkInput
API_KEY: Final = "test-key"
CALLBACK_DELAY_SECONDS: Final = 0.05
CALLBACK_DRAIN_TIMEOUT_SECONDS: Final = 10.0
CALLBACK_TERMINALS: Final[tuple[Literal["success", "failure"], ...]] = ("success", "failure")
PYTHON_HTTP_SENTINEL: Final = "python-ocr-parity-fallback"
PYTHON_VARIANT: Final = ExecutionVariant(name="Python", environment=(("LITELLM_RUST", "0"),))
RUST_VARIANT: Final = ExecutionVariant(name="Rust", environment=(("LITELLM_RUST", "1"),))
@ -75,10 +86,148 @@ class InvalidOcrWorkerCase(BaseModel):
case: InvalidOcrCase
OcrWorkerCase = Annotated[RecordedOcrWorkerCase | InvalidOcrWorkerCase, Field(discriminator="kind")]
class CallbackOcrWorkerCase(BaseModel):
model_config = ConfigDict(frozen=True)
kind: Literal["callback"] = "callback"
case: OcrParityCase
terminal: Literal["success", "failure"]
OcrWorkerCase = Annotated[
RecordedOcrWorkerCase | InvalidOcrWorkerCase | CallbackOcrWorkerCase,
Field(discriminator="kind"),
]
OCR_WORKER_CASE_ADAPTER: Final[TypeAdapter[OcrWorkerCase]] = TypeAdapter(OcrWorkerCase)
class RecordingCallback(CustomLogger):
def __init__(self) -> None:
self.message_logging: Final = True
self.turn_off_message_logging: Final = False
self._observations: Final[queue.SimpleQueue[CallbackObservation]] = queue.SimpleQueue()
def _normalized_kwargs(self, value: object, key: str | None = None) -> JsonValue:
if value is None or isinstance(value, (bool, int, float)):
return value
if isinstance(value, str):
if key != "api_base":
return value
parsed: Final = urlsplit(value)
return urlunsplit(("", "", parsed.path, parsed.query, parsed.fragment))
if isinstance(value, datetime.datetime):
return "datetime"
if isinstance(value, Exception):
return sdk_error_report(value).model_dump(mode="json")
if isinstance(value, BaseModel):
return self._normalized_kwargs(value.model_dump(mode="json"), key)
if isinstance(value, Mapping):
if any(not isinstance(map_key, str) for map_key in value):
raise TypeError("callback kwarg mappings must use string keys")
return {
map_key: self._normalized_kwargs(map_value, map_key)
for map_key, map_value in value.items()
}
if isinstance(value, (list, tuple)):
return [self._normalized_kwargs(item) for item in value]
raise TypeError(f"unsupported callback kwarg type: {type(value)}")
def _record(
self,
hook: Literal[
"log_success_event",
"async_log_success_event",
"log_failure_event",
"async_log_failure_event",
],
phase: Literal["success", "failure"],
kwargs: dict[str, object],
response_obj: object,
) -> None:
raw_litellm_params: Final = kwargs.get("litellm_params")
litellm_params: Final[Mapping[str, object]] = (
cast(Mapping[str, object], raw_litellm_params) if isinstance(raw_litellm_params, Mapping) else {}
)
raw_metadata: Final = litellm_params.get("metadata")
metadata_mapping: Final[Mapping[str, object]] = (
cast(Mapping[str, object], raw_metadata) if isinstance(raw_metadata, Mapping) else {}
)
metadata: Final = JSON_VALUE_ADAPTER.validate_python(
{key: metadata_mapping[key] for key in ("callback_profile", "sdk_route") if key in metadata_mapping}
)
raw_error: Final = kwargs.get("exception")
error: Final = sdk_error_report(raw_error) if isinstance(raw_error, Exception) else None
normalized_kwargs: Final = self._normalized_kwargs(kwargs)
payload_source: Final = (
response_obj.model_dump(mode="json") if isinstance(response_obj, BaseModel) else response_obj
)
payload: Final = JSON_VALUE_ADAPTER.validate_python(payload_source)
raw_model: Final = kwargs.get("model")
raw_call_type: Final = kwargs.get("call_type")
raw_call_id: Final = kwargs.get("litellm_call_id")
self._observations.put(
CallbackObservation(
hook=hook,
phase=phase,
model=raw_model if isinstance(raw_model, str) else None,
call_type=str(raw_call_type) if raw_call_type is not None else None,
litellm_call_id=raw_call_id if isinstance(raw_call_id, str) else None,
metadata=metadata,
kwargs=normalized_kwargs,
payload=payload,
error=error,
)
)
def log_success_event(
self,
kwargs: dict[str, object],
response_obj: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
) -> None:
del start_time, end_time
time.sleep(CALLBACK_DELAY_SECONDS)
self._record("log_success_event", "success", kwargs, response_obj)
async def async_log_success_event(
self,
kwargs: dict[str, object],
response_obj: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
) -> None:
del start_time, end_time
await asyncio.sleep(CALLBACK_DELAY_SECONDS)
self._record("async_log_success_event", "success", kwargs, response_obj)
def log_failure_event(
self,
kwargs: dict[str, object],
response_obj: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
) -> None:
del start_time, end_time
time.sleep(CALLBACK_DELAY_SECONDS)
self._record("log_failure_event", "failure", kwargs, response_obj)
async def async_log_failure_event(
self,
kwargs: dict[str, object],
response_obj: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
) -> None:
del start_time, end_time
await asyncio.sleep(CALLBACK_DELAY_SECONDS)
self._record("async_log_failure_event", "failure", kwargs, response_obj)
def observations(self) -> tuple[CallbackObservation, ...]:
observations: Final = tuple(self._observations.get_nowait() for _ in range(self._observations.qsize()))
return tuple(sorted(observations, key=lambda observation: observation.hook))
INVALID_OCR_CASES: Final = (
InvalidOcrCase(
name="unsupported_provider",
@ -229,6 +378,106 @@ def _execute_invalid_sdk_case(
return _execute_sdk_call(call_kwargs, route, event_loop)
def _callback_call_id(route: SDKRoute, terminal: Literal["success", "failure"]) -> str:
return f"ocr-callback-{route.value}-{terminal}"
def _callback_metadata(route: SDKRoute, terminal: Literal["success", "failure"]) -> dict[str, str]:
return {"callback_profile": terminal, "sdk_route": route.value}
def _drain_callback_delivery(route: SDKRoute, event_loop: asyncio.AbstractEventLoop) -> None:
if route is SDKRoute.AOCR:
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
async def drain_async_callbacks() -> None:
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=CALLBACK_DRAIN_TIMEOUT_SECONDS)
await GLOBAL_LOGGING_WORKER.stop()
event_loop.run_until_complete(drain_async_callbacks())
from litellm.litellm_core_utils.thread_pool_executor import executor
executor.shutdown(wait=True, cancel_futures=False)
def _execute_callback_sdk_case(
case: OcrParityCase,
route: SDKRoute,
terminal: Literal["success", "failure"],
mock_url: str,
event_loop: asyncio.AbstractEventLoop,
) -> WorkerSuccess:
callback: Final = RecordingCallback()
call_kwargs: Final = {
**_call_kwargs(case.litellm_input, mock_url, route),
"callbacks": [callback],
"litellm_call_id": _callback_call_id(route, terminal),
"litellm_trace_id": _callback_call_id(route, terminal),
"metadata": _callback_metadata(route, terminal),
}
report: Final = _execute_sdk_call(call_kwargs, route, event_loop)
_drain_callback_delivery(route, event_loop)
return WorkerSuccess(report=report, callbacks=callback.observations())
def _assert_callback_lifecycle(
execution: Execution,
route: SDKRoute,
terminal: Literal["success", "failure"],
) -> None:
observations: Final = execution.callbacks
assert observations is not None, f"{route.value} {terminal} callbacks were not observed"
expected_hooks: Final = (
(f"log_{terminal}_event",)
if route is SDKRoute.OCR
else ("async_log_success_event",)
if terminal == "success"
else ("async_log_failure_event", "log_failure_event")
)
actual_hooks: Final = tuple(observation.hook for observation in observations)
assert actual_hooks == expected_hooks, (
f"{route.value} {terminal} expected callback hooks {expected_hooks}, received {actual_hooks}"
)
expected_call_id: Final = _callback_call_id(route, terminal)
expected_metadata: Final = _callback_metadata(route, terminal)
for observation in observations:
assert observation.phase == terminal
assert observation.call_type == route.value
assert observation.litellm_call_id == expected_call_id
assert observation.metadata == expected_metadata
assert observation.model
if terminal == "success":
assert isinstance(execution.report, SDKSuccess)
assert observation.payload == execution.report.response
assert observation.error is None
else:
assert isinstance(execution.report, SDKError)
assert observation.payload is None
assert observation.error is not None
assert observation.error.exception_type
assert observation.error.message
assert observation.error.status_code is not None
assert observation.error.status_code >= 400
def _check_callback_ocr_sdk_parity(
case: OcrParityCase,
route: SDKRoute,
terminal: Literal["success", "failure"],
case_file: Path,
runner: SubprocessRunner,
) -> None:
with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers:
python_worker, rust_worker = workers
python: Final = python_worker.execute(case_file, route.value, case.provider_responses)
rust: Final = rust_worker.execute(case_file, route.value, case.provider_responses)
_assert_callback_lifecycle(python, route, terminal)
_assert_callback_lifecycle(rust, route, terminal)
assert_parity(python, rust, PYTHON_HTTP_SENTINEL)
def _check_recorded_ocr_sdk_parity(
ocr_fixture: OcrParityCase,
route: SDKRoute,
@ -276,6 +525,25 @@ def _write_worker_case(directory: Path, index: int, case: OcrWorkerCase) -> Path
return case_file
def _callback_fixture(
fixtures: tuple[OcrParityCase, ...],
terminal: Literal["success", "failure"],
) -> OcrParityCase:
matching: Final = tuple(
fixture
for fixture in fixtures
if fixture.litellm_input.contract == "mistral"
and (
all(response.status_code < 400 for response in fixture.provider_responses)
if terminal == "success"
else any(response.status_code >= 400 for response in fixture.provider_responses)
)
)
if not matching:
raise AssertionError(f"no recorded Mistral OCR {terminal} fixture is available for callback parity")
return min(matching, key=lambda fixture: fixture_id(fixture.litellm_input, fixture.litellm_input.model))
@contextmanager
def parity_checks() -> Generator[tuple[E2ECheck, ...]]:
fixtures: Final = tuple(
@ -298,6 +566,17 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]:
_write_worker_case(directory, len(recorded_files) + index, InvalidOcrWorkerCase(case=case))
for index, case in enumerate(INVALID_OCR_CASES)
)
callback_cases: Final[tuple[tuple[Literal["success", "failure"], OcrParityCase], ...]] = tuple(
(terminal, _callback_fixture(fixtures, terminal)) for terminal in CALLBACK_TERMINALS
)
callback_files: Final = tuple(
_write_worker_case(
directory,
len(recorded_files) + len(invalid_files) + index,
CallbackOcrWorkerCase(case=case, terminal=terminal),
)
for index, (terminal, case) in enumerate(callback_cases)
)
with execution_worker_pair(runner, PYTHON_VARIANT, RUST_VARIANT) as workers:
recorded: Final = tuple(
E2ECheck(
@ -315,7 +594,15 @@ def parity_checks() -> Generator[tuple[E2ECheck, ...]]:
for case, case_file in zip(INVALID_OCR_CASES, invalid_files, strict=True)
for route in SDKRoute
)
yield (*recorded, *invalid)
callbacks: Final = tuple(
E2ECheck(
f"callback:{route.value}:{terminal}",
partial(_check_callback_ocr_sdk_parity, case, route, terminal, case_file, runner),
)
for (terminal, case), case_file in zip(callback_cases, callback_files, strict=True)
for route in SDKRoute
)
yield (*recorded, *invalid, *callbacks)
def _execute_worker_command(
@ -333,6 +620,8 @@ def _execute_worker_command(
return WorkerSuccess(report=_execute_sdk_case(recorded.litellm_input, route, mock_url, event_loop))
case InvalidOcrWorkerCase(case=invalid):
return WorkerSuccess(report=_execute_invalid_sdk_case(invalid, route, mock_url, event_loop))
case CallbackOcrWorkerCase(case=callback_case, terminal=terminal):
return _execute_callback_sdk_case(callback_case, route, terminal, mock_url, event_loop)
except Exception:
return WorkerFailure(error=traceback.format_exc())

1
tests/rust-python-harness/strategies/trace_parity/models.py
View file

@ -16,6 +16,7 @@ TraceFailureSource = Literal["python", "rust", "harness"]
class RouteFixture:
kwargs: dict[str, object]
provider_responses: tuple[RecordedHttpResponse, ...]
expected_failure: bool = False
@dataclass(frozen=True, slots=True)

74
tests/rust-python-harness/strategies/trace_parity/sdk/execution.py
View file

@ -2,15 +2,16 @@ from __future__ import annotations
import asyncio
from collections.abc import Awaitable
from dataclasses import dataclass
from pathlib import Path
from typing import Final, Protocol, cast
from ....shared.parity.replay import replay_server
from ....shared.reporting.models import Surface
from ....shared.tracing.native import native_trace_events
from ....shared.tracing.native import TraceResponsePayload, native_trace_events
from ....shared.tracing.profiler import FunctionTraceEvent, profile_python
from ....shared.tracing.steps import Engine, pipeline_projection
from ..models import RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario
from ..models import RouteFixture, RouteSpec, TraceExecutionFailure, TraceMode, TraceScenario
from ..reporting import TraceComparisonArtifact
@ -18,9 +19,22 @@ class SdkCall(Protocol):
def __call__(self, **kwargs: object) -> object: ...
@dataclass(frozen=True, slots=True)
class _CollectedTrace:
events: tuple[FunctionTraceEvent, ...]
error: str | None = None
def _invoke(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> object:
async def invoke_async() -> object:
return await cast(Awaitable[object], function(**kwargs))
try:
return await cast(Awaitable[object], function(**kwargs))
finally:
await asyncio.sleep(0)
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
await GLOBAL_LOGGING_WORKER.stop()
if asynchronous:
return asyncio.run(invoke_async())
@ -48,16 +62,30 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa
return cast(SdkCall, getattr(owner, spec.python_entrypoints[int(asynchronous)]))
def _python_invocation_error(function: SdkCall, kwargs: dict[str, object], *, asynchronous: bool) -> str | None:
try:
_invoke(function, kwargs, asynchronous=asynchronous)
except Exception as error:
return f"{type(error).__name__}: {error}"
return None
def _collect(
function: SdkCall, kwargs: dict[str, object], engine: Engine, *, asynchronous: bool
) -> tuple[FunctionTraceEvent, ...]:
function: SdkCall,
fixture: RouteFixture,
engine: Engine,
*,
asynchronous: bool,
) -> _CollectedTrace:
kwargs: Final = fixture.kwargs
if engine == "rust":
return native_trace_events(_invoke(function, kwargs, asynchronous=asynchronous))
payload: Final = TraceResponsePayload.model_validate(_invoke(function, kwargs, asynchronous=asynchronous))
return _CollectedTrace(native_trace_events(payload), payload.error)
import litellm
with profile_python(Path(litellm.__file__).parent, threads=True) as profiler:
_invoke(function, kwargs, asynchronous=asynchronous)
return tuple(profiler.events)
error: Final = _python_invocation_error(function, kwargs, asynchronous=asynchronous)
return _CollectedTrace(tuple(profiler.events), error)
def collect_trace(
@ -68,22 +96,30 @@ def collect_trace(
return function
try:
with replay_server() as provider:
fixture: Final = spec.fixture(engine, provider.url)
for response in fixture.provider_responses:
base_fixture: Final = spec.fixture(engine, provider.url)
for response in base_fixture.provider_responses:
provider.enqueue_response(response)
kwargs: Final = {
**fixture.kwargs,
"api_key": "test-key",
"api_base": provider.url,
**({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}),
}
events: Final = _collect(function, kwargs, engine, asynchronous=asynchronous)
fixture: Final = RouteFixture(
kwargs={
**base_fixture.kwargs,
"api_key": "test-key",
"api_base": provider.url,
**({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}),
},
provider_responses=base_fixture.provider_responses,
expected_failure=base_fixture.expected_failure,
)
collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous)
provider.take_requests(len(fixture.provider_responses))
except Exception as error:
return TraceExecutionFailure(engine, f"{type(error).__name__}: {error}")
if not events:
if fixture.expected_failure and collected.error is None:
return TraceExecutionFailure(engine, "call succeeded but the scenario expects failure")
if not fixture.expected_failure and collected.error is not None:
return TraceExecutionFailure(engine, collected.error)
if not collected.events:
return TraceExecutionFailure(engine, "trace is empty")
return events
return collected.events
def _failure_message(result: tuple[FunctionTraceEvent, ...] | TraceExecutionFailure) -> str | None:

91
tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py
View file

@ -19,6 +19,20 @@ COMMON_MAPPINGS: Final = (
mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"),
)
SUCCESS_CALLBACK_SYNC_MAPPING: Final = mapping(
rust_span="success_callback",
python_frame=r"BoundedLoggingThreadPoolExecutor\.submit$",
)
SUCCESS_CALLBACK_ASYNC_MAPPING: Final = mapping(
rust_span="success_callback",
python_frame=r"Logging\.async_success_handler$",
)
FAILURE_CALLBACK_MAPPING: Final = mapping(
rust_span="failure_callback",
python_frame=r"Logging\.(?:async_)?failure_handler$",
)
IGNORED_SUCCESS_CALLBACK_MAPPING: Final = mapping(rust_span="success_callback")
SYNC_MAPPINGS: Final = (
*COMMON_MAPPINGS,
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
@ -39,6 +53,21 @@ ASYNC_MAPPINGS: Final = (
),
)
CALLBACK_SUCCESS_SYNC_MAPPINGS: Final = (*SYNC_MAPPINGS, SUCCESS_CALLBACK_SYNC_MAPPING)
CALLBACK_SUCCESS_ASYNC_MAPPINGS: Final = (*ASYNC_MAPPINGS, SUCCESS_CALLBACK_ASYNC_MAPPING)
CALLBACK_FAILURE_SYNC_MAPPINGS: Final = (
*COMMON_MAPPINGS,
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
FAILURE_CALLBACK_MAPPING,
)
CALLBACK_FAILURE_ASYNC_MAPPINGS: Final = (
*COMMON_MAPPINGS,
mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"),
FAILURE_CALLBACK_MAPPING,
)
AZURE_COMMON_MAPPINGS: Final = (
*COMMON_MAPPINGS[:7],
mapping(
@ -99,6 +128,34 @@ def _mistral_fixture(engine: Engine, _base_url: str) -> RouteFixture:
return _fixture(engine, "mistral/mistral-ocr-latest")
def _callback_fixture(engine: Engine, *, failure: bool) -> RouteFixture:
fixture: Final = _fixture(engine, "mistral/mistral-ocr-latest")
provider_responses: Final = (
(
RecordedHttpResponse.from_bytes(
400,
(HttpHeader(name="content-type", value="application/json"),),
b'{"message":"trace callback provider failure"}',
),
)
if failure
else fixture.provider_responses
)
return RouteFixture(
kwargs=fixture.kwargs,
provider_responses=provider_responses,
expected_failure=failure,
)
def _mistral_callback_success_fixture(engine: Engine, _base_url: str) -> RouteFixture:
return _callback_fixture(engine, failure=False)
def _mistral_callback_failure_fixture(engine: Engine, _base_url: str) -> RouteFixture:
return _callback_fixture(engine, failure=True)
def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture:
return _fixture(
engine,
@ -304,36 +361,50 @@ TRACE_SUITE: Final = TraceSuite(
name="mistral",
fixture=_mistral_fixture,
mappings=COMMON_MAPPINGS,
sync_mappings=SYNC_MAPPINGS,
async_mappings=ASYNC_MAPPINGS,
sync_mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
async_mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
),
TraceScenario(
name="mistral-callback-success",
fixture=_mistral_callback_success_fixture,
mappings=COMMON_MAPPINGS,
sync_mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS,
async_mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS,
),
TraceScenario(
name="mistral-callback-failure",
fixture=_mistral_callback_failure_fixture,
mappings=(*COMMON_MAPPINGS, FAILURE_CALLBACK_MAPPING),
sync_mappings=CALLBACK_FAILURE_SYNC_MAPPINGS,
async_mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS,
),
TraceScenario(
name="azure-ai",
fixture=_azure_fixture,
mappings=AZURE_COMMON_MAPPINGS,
sync_mappings=AZURE_SYNC_MAPPINGS,
async_mappings=AZURE_ASYNC_MAPPINGS,
sync_mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
async_mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
),
TraceScenario(
name="azure-document-intelligence",
fixture=_azure_document_intelligence_fixture,
mappings=DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS,
sync_mappings=DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS,
async_mappings=DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS,
sync_mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
async_mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
),
TraceScenario(
name="vertex-ai",
fixture=_vertex_fixture,
mappings=VERTEX_COMMON_MAPPINGS,
sync_mappings=VERTEX_SYNC_MAPPINGS,
async_mappings=VERTEX_ASYNC_MAPPINGS,
sync_mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
async_mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
),
TraceScenario(
name="vertex-deepseek",
fixture=_vertex_deepseek_fixture,
mappings=DEEPSEEK_COMMON_MAPPINGS,
sync_mappings=DEEPSEEK_SYNC_MAPPINGS,
async_mappings=DEEPSEEK_ASYNC_MAPPINGS,
sync_mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
async_mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
),
),
)

4
tests/rust-python-harness/strategies/unit_tests_mapping/cases/ocr.py
View file

@ -230,6 +230,10 @@ _HOST_ONLY_BRIDGE_EXCLUSIONS: Final = tuple(
"test_ocr_exception_type_uses_resolved_provider_context",
"Python wraps bridge exceptions into public errors.",
),
(
"test_rust_upstream_error_uses_ocr_provider_error_mapping",
"Python maps native upstream errors through the selected OCR provider config.",
),
("test_aocr_routes_to_async_rust_when_enabled", "Python selects and invokes the async native bridge."),
("test_aocr_exception_type_uses_resolved_provider_context", "Python wraps async bridge exceptions."),
("test_ocr_forwards_timeout_to_rust", "Python converts and forwards explicit timeouts."),

26
tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py
View file

@ -126,16 +126,6 @@ def test_load_rust_messages_returns_injected_impl():
assert rust_messages.load_rust_messages() is bridge
def test_bare_rust_still_toggles_ocr():
from litellm.rust_bridge.ocr import rust_ocr_enabled
litellm.rust(True)
assert rust_ocr_enabled() is True
litellm.rust(False)
assert rust_ocr_enabled() is False
def test_load_rust_amessages_returns_injected_impl():
bridge = RecordingAsyncMessages()
litellm.rust(True)
@ -214,7 +204,7 @@ async def test_amessages_wrapper_forwards_args():
def _gate(**overrides):
kwargs = {
"custom_llm_provider": "azure_ai",
"litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True),
"litellm_params": GenericLiteLLMParams(api_key="sk-azure"),
"has_agentic_hook": False,
"model": "claude-sonnet-4-5",
"api_key": "sk-azure",
@ -282,18 +272,6 @@ async def test_gate_uses_process_enable_without_request_override():
assert bridge.calls[0]["custom_llm_provider"] == "azure_ai"
@pytest.mark.asyncio
async def test_gate_skips_rust_when_flag_false():
bridge = ExplodingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False))
assert response is None
assert bridge.calls == 0
@pytest.mark.asyncio
async def test_gate_invokes_rust_for_native_anthropic_provider():
bridge = RecordingAsyncMessages()
@ -302,7 +280,7 @@ async def test_gate_invokes_rust_for_native_anthropic_provider():
response = await _gate(
custom_llm_provider="anthropic",
litellm_params=GenericLiteLLMParams(api_key="sk-ant", rust=True),
litellm_params=GenericLiteLLMParams(api_key="sk-ant"),
api_key="sk-ant",
api_base="https://api.anthropic.com",
headers={"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"},

57
tests/test_litellm/batches/test_batch_utils.py
View file

@ -21,11 +21,12 @@ from types import MappingProxyType
import httpx
import pytest
import respx
from openai.types.batch import BatchRequestCounts
import litellm
import litellm.batches.batch_utils as bu
from litellm.types.utils import Usage
from litellm.types.utils import LiteLLMBatch, Usage
# --------------------------------------------------------------------------- #
# Builders for batch OUTPUT file rows.
@ -1718,3 +1719,57 @@ def test_unparsable_bedrock_batch_usage_warns(caplog):
assert usage.total_tokens == 0
assert "does not understand" in caplog.text
assert "inputTextTokenCount" in caplog.text
# --------------------------------------------------------------------------- #
# batch_cost_is_final
# --------------------------------------------------------------------------- #
def _retrieved_batch(
status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None
) -> LiteLLMBatch:
return LiteLLMBatch(
id="batch_abc",
completion_window="24h",
created_at=1,
endpoint="/v1/chat/completions",
input_file_id="file-in",
object="batch",
status="validating",
output_file_id=output_file_id,
request_counts=counts,
).model_copy(update={"status": status})
class TestBatchCostIsFinal:
"""Every retrieve of one batch writes the same spend row, so the first retrieve
that prices it decides the row for good. A poll before the output exists must
therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048)."""
@pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"])
def test_in_flight_batch_is_not_final(self, status):
assert bu.batch_cost_is_final(_retrieved_batch(status)) is False
@pytest.mark.parametrize("status", ["completed", "complete"])
def test_completed_with_output_is_final(self, status):
assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True
def test_completed_without_output_and_unknown_counts_is_not_final(self):
assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False
def test_completed_without_output_and_zero_counts_is_not_final(self):
counts = BatchRequestCounts(total=0, completed=0, failed=0)
assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False
def test_completed_without_output_but_successful_lines_is_not_final(self):
counts = BatchRequestCounts(total=2, completed=2, failed=0)
assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False
@pytest.mark.parametrize("status", ["completed", "complete"])
def test_completed_without_output_and_every_line_failed_is_final(self, status):
counts = BatchRequestCounts(total=2, completed=0, failed=2)
assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True
@pytest.mark.parametrize("status", ["failed", "expired", "cancelled"])
def test_other_terminal_statuses_are_final(self, status):
assert bu.batch_cost_is_final(_retrieved_batch(status)) is True

195
tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
View file

@ -3952,3 +3952,198 @@ def test_convert_chat_completion_messages_to_responses_api_tool_result_with_tool
function_call_output = next(item for item in response if item.get("type") == "function_call_output")
assert function_call_output["output"] == [{"type": "input_text", "text": "1 tool found"}]
def _litellm_encoded_response_id(upstream_id: str) -> str:
from litellm.responses.utils import ResponsesAPIRequestUtils
return ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="azure", model_id="deployment-1", response_id=upstream_id
)
def test_transform_response_keeps_upstream_id_and_provider_extras():
from unittest.mock import Mock
from litellm.completion_extras.litellm_responses_transformation.transformation import (
LiteLLMResponsesTransformationHandler,
)
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import ModelResponse, Usage
content_filters = [
{"blocked": False, "source_type": "prompt", "content_filter_results": {"hate": {"filtered": False}}}
]
raw_response = ResponsesAPIResponse.model_validate(
{
"id": _litellm_encoded_response_id("resp_azure_123"),
"created_at": 1734366691,
"object": "response",
"model": "gpt-5.6",
"status": "completed",
"output": [
{
"type": "function_call",
"id": "fc_1",
"call_id": "call_1",
"name": "lookup_weather",
"arguments": '{"city": "Seattle"}',
"status": "completed",
}
],
"parallel_tool_calls": True,
"tool_choice": "auto",
"tools": [],
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
"service_tier": "default",
"content_filters": content_filters,
"max_tool_calls": None,
"background": False,
"top_logprobs": 0,
"store": True,
}
)
model_response = ModelResponse(
id="chatcmpl-local",
created=1734366691,
model=None,
object="chat.completion",
choices=[],
usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0),
)
result = LiteLLMResponsesTransformationHandler().transform_response(
model="gpt-5.6",
raw_response=raw_response,
model_response=model_response,
logging_obj=Mock(),
request_data={"model": "gpt-5.6"},
messages=[{"role": "user", "content": "What is the weather in Seattle?"}],
optional_params={},
litellm_params={},
encoding=Mock(),
)
dumped = result.model_dump()
assert dumped["id"] == "resp_azure_123"
assert dumped["object"] == "chat.completion"
assert dumped["service_tier"] == "default"
assert dumped["content_filters"] == content_filters
assert "max_tool_calls" not in dumped, "a null provider field must not appear as a null top-level key"
assert "output" not in dumped and "status" not in dumped, (
"Responses schema fields must not leak into the chat response"
)
assert not {"background", "top_logprobs", "store"} & dumped.keys(), (
"Responses API bookkeeping must not ride along as chat metadata"
)
assert dumped["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "lookup_weather"
def test_bridged_response_is_priced_by_the_reported_service_tier():
from unittest.mock import Mock
from litellm.completion_extras.litellm_responses_transformation.transformation import (
LiteLLMResponsesTransformationHandler,
)
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import ModelResponse
raw_response = ResponsesAPIResponse.model_validate(
{
"id": "resp_flex",
"created_at": 1734366691,
"object": "response",
"model": "gpt-5.4",
"status": "completed",
"output": [
{
"type": "message",
"id": "msg_1",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "hi", "annotations": []}],
}
],
"usage": {"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100},
"service_tier": "flex",
}
)
result = LiteLLMResponsesTransformationHandler().transform_response(
model="gpt-5.4",
raw_response=raw_response,
model_response=ModelResponse(),
logging_obj=Mock(),
request_data={"model": "gpt-5.4"},
messages=[{"role": "user", "content": "hi"}],
optional_params={},
litellm_params={},
encoding=Mock(),
)
pricing = litellm.model_cost["gpt-5.4"]
flex_cost = 1000 * pricing["input_cost_per_token_flex"] + 100 * pricing["output_cost_per_token_flex"]
standard_cost = 1000 * pricing["input_cost_per_token"] + 100 * pricing["output_cost_per_token"]
cost = litellm.completion_cost(completion_response=result, custom_llm_provider="openai")
assert cost == pytest.approx(flex_cost)
assert cost < standard_cost
def test_streaming_chunks_carry_the_upstream_response_id():
from litellm.completion_extras.litellm_responses_transformation.transformation import (
OpenAiResponsesToChatCompletionStreamIterator,
)
iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True)
encoded_id = _litellm_encoded_response_id("resp_azure_stream")
events = [
{"type": "response.created", "response": {"id": encoded_id, "output": []}},
{"type": "response.output_text.delta", "delta": "Hel"},
{"type": "response.completed", "response": {"id": encoded_id, "output": [{"type": "message"}]}},
]
ids = [iterator.chunk_parser(event).id for event in events]
assert ids == ["resp_azure_stream"] * len(events), f"streamed chunks did not carry the upstream id: {ids}"
def test_streaming_final_chunk_carries_provider_metadata():
from unittest.mock import MagicMock
from litellm.completion_extras.litellm_responses_transformation.transformation import (
OpenAiResponsesToChatCompletionStreamIterator,
)
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True)
content_filters = [{"blocked": False, "source_type": "completion", "content_filter_results": {}}]
events = [
{"type": "response.created", "response": {"id": "resp_azure_stream", "output": []}},
{"type": "response.output_text.delta", "delta": "Hello"},
{
"type": "response.completed",
"response": {
"id": "resp_azure_stream",
"output": [{"type": "message"}],
"usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4},
"service_tier": "default",
"content_filters": content_filters,
"background": False,
},
},
]
stream = CustomStreamWrapper(
completion_stream=iter([iterator.chunk_parser(event) for event in events]),
model="gpt-5.6",
custom_llm_provider="azure",
logging_obj=MagicMock(),
)
chunks = [chunk.model_dump() for chunk in stream]
assert chunks[-1]["choices"][0]["finish_reason"] == "stop"
assert chunks[-1]["service_tier"] == "default"
assert chunks[-1]["content_filters"] == content_filters
assert "background" not in chunks[-1]
assert all("service_tier" not in chunk for chunk in chunks[:-1])

43
tests/test_litellm/conftest.py
View file

@ -7,9 +7,12 @@
# 4. Added proper cleanup in fixtures
# 5. Added worker-specific isolation for parallel execution
import base64
import importlib
import os
from pathlib import Path
from types import SimpleNamespace
import httpx
import pytest
import asyncio
@ -595,3 +598,43 @@ def pytest_sessionfinish(session, exitstatus):
_close_handler_if_needed(getattr(litellm, "aclient", None))
_close_handler_if_needed(getattr(litellm, "client", None))
_run_coroutine_if_needed(close_litellm_async_clients())
ONE_PIXEL_PNG = base64.b64decode(
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
)
@pytest.fixture
def async_only_image_fetch(monkeypatch):
from litellm.litellm_core_utils.prompt_templates import factory, image_handling
from litellm.llms.gemini.chat import transformation as gemini_chat_transformation
fetch = SimpleNamespace(
fetched=[],
base64_png=base64.b64encode(ONE_PIXEL_PNG).decode(),
data_url="data:image/png;base64," + base64.b64encode(ONE_PIXEL_PNG).decode(),
)
def forbid_sync_fetch(client, url, **kwargs):
raise litellm.ImageFetchError(f"sync image fetch ran on the event loop: {url}")
async def serve_png(client, url, **kwargs):
fetch.fetched.append(url)
return httpx.Response(
200,
content=ONE_PIXEL_PNG,
headers={"content-type": "image/png"},
request=httpx.Request("GET", url),
)
def forbid_sync_convert(url, *args, **kwargs):
if url.startswith(("http://", "https://")):
raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}")
return url
monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch)
monkeypatch.setattr(image_handling, "async_safe_get", serve_png)
for module in (image_handling, factory, gemini_chat_transformation):
monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert)
return fetch

124
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -2008,7 +2008,14 @@ def test_generic_cost_per_token_azure_gpt56(_local_model_cost_map,
assert round(completion_cost, 10) == round(output_cost * completion_tokens, 10)
@pytest.mark.parametrize("model,zone_multiplier", [("azure/gpt-6-astra", 1.0), ("azure/us/gpt-6-astra", 1.1)])
@pytest.mark.parametrize(
"model,custom_llm_provider,zone_multiplier",
[
("azure/gpt-6-astra", "azure", 1.0),
("azure/us/gpt-6-astra", "azure", 1.1),
("azure_ai/gpt-6-astra", "azure_ai", 1.0),
],
)
@pytest.mark.parametrize(
"prompt_tokens,input_side_multiplier,output_multiplier",
[(100000, 1.0, 1.0), (300000, 2.0, 1.5)],
@ -2016,6 +2023,7 @@ def test_generic_cost_per_token_azure_gpt56(_local_model_cost_map,
def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet(
_local_model_cost_map,
model,
custom_llm_provider,
zone_multiplier,
prompt_tokens,
input_side_multiplier,
@ -2023,7 +2031,8 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet(
):
"""Microsoft Foundry sells gpt-6-astra at the OpenAI rates: $10 input, $1 cache read, $12.50 cache write,
$50 output per 1M tokens on Standard Global, with the input side doubling and output 1.5x above 272K
prompt tokens. Standard US Data Zone carries the usual 10% uplift on every rate.
prompt tokens. Standard US Data Zone carries the usual 10% uplift on every rate. A Foundry
deployment reached through the azure_ai route bills the same Standard Global sheet.
"""
cached_tokens = 50000
cache_write_tokens = 40000
@ -2041,7 +2050,7 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet(
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=usage,
custom_llm_provider="azure",
custom_llm_provider=custom_llm_provider,
)
input_side = zone_multiplier * input_side_multiplier
@ -2051,6 +2060,18 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet(
assert completion_cost == pytest.approx(zone_multiplier * output_multiplier * completion_tokens * 5e-5)
def test_generic_cost_per_token_azure_ai_gpt_6_astra_flex_bills_the_standard_rate(_local_model_cost_map):
usage = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100)
standard = generic_cost_per_token(model="azure_ai/gpt-6-astra", usage=usage, custom_llm_provider="azure_ai")
flex = generic_cost_per_token(
model="azure_ai/gpt-6-astra", usage=usage, custom_llm_provider="azure_ai", service_tier="flex"
)
assert flex == standard
assert standard == pytest.approx((1000 * 1e-05, 100 * 5e-05))
@pytest.mark.parametrize(
"model,expected_none,expected_xhigh,expected_minimal",
[
@ -4766,3 +4787,100 @@ def test_route_image_generation_cost_falls_back_to_requested_size(monkeypatch, r
)
assert cost == expected_cost
def test_generic_cost_per_token_bills_reasoning_nested_in_text_tokens_once(_local_model_cost_map: None) -> None:
"""
Realtime usage (OpenAI and Azure) reports output_tokens == text_tokens + audio_tokens with
reasoning_tokens already counted inside text_tokens, so reasoning must not be billed on top.
"""
model = "gpt-realtime-2.1-mini"
usage = Usage(
prompt_tokens=346,
completion_tokens=29,
total_tokens=375,
prompt_tokens_details=PromptTokensDetailsWrapper(
text_tokens=152, image_tokens=194, audio_tokens=0, cached_tokens=128
),
completion_tokens_details=CompletionTokensDetailsWrapper(
text_tokens=29, audio_tokens=0, reasoning_tokens=19
),
)
prompt_cost, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai")
breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage)
info = litellm.get_model_info(model=model, custom_llm_provider="openai")
assert completion_cost == pytest.approx(29 * info["output_cost_per_token"])
assert completion_cost - breakdown.reasoning_cost == pytest.approx(10 * info["output_cost_per_token"])
assert prompt_cost == pytest.approx(
24 * info["input_cost_per_token"]
+ 128 * info["cache_read_input_token_cost"]
+ 194 * info["input_cost_per_image_token"]
)
def test_generic_cost_per_token_keeps_billing_reasoning_reported_beside_text_tokens(
_local_model_cost_map: None,
) -> None:
"""Providers whose text_tokens exclude reasoning (text + reasoning == completion) stay billed in full."""
model = "gpt-realtime-2.1-mini"
usage = Usage(
prompt_tokens=100,
completion_tokens=44,
total_tokens=144,
completion_tokens_details=CompletionTokensDetailsWrapper(
text_tokens=25, audio_tokens=0, reasoning_tokens=19
),
)
_, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai")
info = litellm.get_model_info(model=model, custom_llm_provider="openai")
assert completion_cost == pytest.approx(44 * info["output_cost_per_token"])
def test_generic_cost_per_token_strips_only_the_reasoning_share_when_text_over_reports(
_local_model_cost_map: None,
) -> None:
"""Text over-reported past the reasoning share keeps its extra tokens billed; only the nested reasoning is netted out."""
model = "gpt-realtime-2.1-mini"
usage = Usage(
prompt_tokens=120,
completion_tokens=100,
total_tokens=220,
completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=100, audio_tokens=70, reasoning_tokens=10),
)
_, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai")
breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage)
info = litellm.get_model_info(model=model, custom_llm_provider="openai")
assert breakdown.reasoning_cost == pytest.approx(10 * info["output_cost_per_token"])
assert completion_cost == pytest.approx(
100 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"]
)
def test_generic_cost_per_token_bills_nested_reasoning_once_beside_audio_output(_local_model_cost_map: None) -> None:
"""Audio-output realtime usage nests reasoning inside text_tokens next to audio_tokens; text is billed net of it."""
model = "gpt-realtime-2.1-mini"
usage = Usage(
prompt_tokens=120,
completion_tokens=100,
total_tokens=220,
completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=30, audio_tokens=70, reasoning_tokens=20),
)
_, completion_cost = generic_cost_per_token(model=model, usage=usage, custom_llm_provider="openai")
breakdown = get_token_type_cost_breakdown(model=model, custom_llm_provider="openai", usage=usage)
info = litellm.get_model_info(model=model, custom_llm_provider="openai")
assert breakdown.reasoning_cost == pytest.approx(20 * info["output_cost_per_token"])
assert completion_cost == pytest.approx(
30 * info["output_cost_per_token"] + 70 * info["output_cost_per_audio_token"]
)

29
tests/test_litellm/litellm_core_utils/test_get_litellm_params.py
View file

@ -215,32 +215,3 @@ class TestMetadataFallsBackToLitellmMetadata:
assert result["metadata"] is not litellm_metadata
result["metadata"].pop("trace_id")
assert litellm_metadata == {"trace_id": "trace-1"}
class TestRustOptIn:
"""`rust: true` is a litellm param, so it has to reach `litellm_params`.
`all_litellm_params` keeps it out of the provider body; without it also
being carried into `litellm_params` the chat completions handlers cannot
see the opt-in and the Rust path is silently never taken.
"""
def test_rust_is_an_optional_kwargs_key(self):
assert "rust" in _OPTIONAL_KWARGS_KEYS
def test_rust_is_forwarded_from_completion_kwargs(self):
from litellm.litellm_core_utils.get_litellm_params import FORWARDED_KWARGS_KEYS
assert "rust" in FORWARDED_KWARGS_KEYS
def test_rust_survives_into_litellm_params(self):
params = get_litellm_params(rust=True)
assert params["rust"] is True
def test_rust_is_absent_when_the_deployment_did_not_set_it(self):
assert "rust" not in get_litellm_params()
def test_rust_stays_out_of_the_provider_body(self):
from litellm.types.utils import all_litellm_params
assert "rust" in all_litellm_params

272
tests/test_litellm/litellm_core_utils/test_image_handling.py
View file

@ -1,3 +1,7 @@
import asyncio
import copy
import time
import uuid
from unittest.mock import patch
import pytest
@ -7,9 +11,13 @@ import litellm
from litellm import constants
from litellm.litellm_core_utils.prompt_templates import image_handling
from litellm.litellm_core_utils.prompt_templates.image_handling import (
MAX_CONCURRENT_REMOTE_MEDIA_FETCHES,
RemoteMedia,
async_convert_url_to_base64,
async_inline_remote_media,
convert_url_to_base64,
)
from litellm.litellm_core_utils.url_utils import SSRFError
@pytest.fixture(autouse=True)
@ -107,9 +115,7 @@ class StreamingLargeImageClient:
request=Request("GET", url),
)
# Mock the iter_bytes method to return our generator
response.iter_bytes = lambda chunk_size=8192: generate_chunks(
size_bytes, chunk_size
)
response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size)
return response
@ -207,9 +213,7 @@ def test_streaming_download_handles_petabyte_file(monkeypatch):
"""
# Simulate a 1 petabyte file (1,000,000 GB)
# Without streaming protection, this would cause OOM or hang indefinitely
client = StreamingLargeImageClient(
size_mb=1_000_000_000, include_content_length=False
)
client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False)
monkeypatch.setattr(litellm, "module_level_client", client)
with pytest.raises(litellm.ImageFetchError) as excinfo:
@ -268,3 +272,259 @@ def test_image_size_limit_disabled(monkeypatch):
assert "Image URL download is disabled" in str(excinfo.value)
assert "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0" in str(excinfo.value)
async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_only_image_fetch):
image_url = f"http://img.example/{uuid.uuid4()}.png"
pdf_url = f"http://docs.example/{uuid.uuid4()}.pdf"
messages = [
{"role": "system", "content": "be terse"},
{
"role": "user",
"content": [
{"type": "text", "text": "what is this?"},
{"type": "image_url", "image_url": {"url": image_url, "detail": "low"}},
{"type": "image_url", "image_url": image_url},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
{"type": "file", "file": {"file_id": pdf_url}},
{"type": "file", "file": {"file_id": image_url, "format": "image/png"}},
{"type": "document", "source": {"type": "url", "url": pdf_url}, "title": "the doc"},
{"type": "image", "source": {"type": "url", "url": image_url}},
{"type": "document", "source": {"type": "file", "file_id": "file_abc"}},
],
},
]
snapshot = copy.deepcopy(messages)
inlined = await async_inline_remote_media(messages)
data_url = async_only_image_fetch.data_url
base64_png = async_only_image_fetch.base64_png
assert inlined[0] == {"role": "system", "content": "be terse"}
assert inlined[1]["content"] == [
{"type": "text", "text": "what is this?"},
{"type": "image_url", "image_url": {"url": data_url, "detail": "low"}},
{"type": "image_url", "image_url": data_url},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
{"type": "file", "file": {"format": "application/pdf", "file_data": data_url}},
{"type": "file", "file": {"format": "image/png", "file_data": data_url}},
{
"type": "document",
"source": {"type": "base64", "media_type": "application/pdf", "data": base64_png},
"title": "the doc",
},
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": base64_png}},
{"type": "document", "source": {"type": "file", "file_id": "file_abc"}},
]
assert sorted(async_only_image_fetch.fetched) == sorted([image_url, pdf_url])
assert messages == snapshot
async def test_async_inline_remote_media_inlines_only_the_parts_the_predicate_accepts(async_only_image_fetch):
files_api_prefix = "https://generativelanguage.googleapis.com/v1beta/files/"
files_api_pdf = f"{files_api_prefix}{uuid.uuid4().hex}"
hinted_image = f"https://img.example/{uuid.uuid4()}.png"
plain_image = f"https://img.example/{uuid.uuid4()}.png"
hinted_document = f"https://docs.example/{uuid.uuid4()}.pdf"
seen = []
def inline_unhinted_outside_files_api(media: RemoteMedia) -> bool:
seen.append(media)
return not media.url.startswith(files_api_prefix) and "format" not in media.fields
messages = [
{
"role": "user",
"content": [
{"type": "file", "file": {"file_id": files_api_pdf}},
{"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}},
{"type": "image_url", "image_url": {"url": plain_image}},
{"type": "image_url", "image_url": plain_image},
{"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}},
],
}
]
snapshot = copy.deepcopy(messages)
inlined = await async_inline_remote_media(messages, should_inline=inline_unhinted_outside_files_api)
assert inlined[0]["content"] == [
{"type": "file", "file": {"file_id": files_api_pdf}},
{"type": "image_url", "image_url": {"url": hinted_image, "format": "image/png"}},
{"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}},
{"type": "image_url", "image_url": async_only_image_fetch.data_url},
{"type": "document", "source": {"type": "url", "url": hinted_document, "format": "application/pdf"}},
]
assert async_only_image_fetch.fetched == [plain_image]
assert [(media.url, dict(media.fields)) for media in seen[:5]] == [
(files_api_pdf, {"file_id": files_api_pdf}),
(hinted_image, {"url": hinted_image, "format": "image/png"}),
(plain_image, {"url": plain_image}),
(plain_image, {}),
(hinted_document, {"type": "url", "url": hinted_document, "format": "application/pdf"}),
]
assert messages == snapshot
async def test_async_inline_remote_media_inlines_a_shared_url_only_where_the_predicate_accepts_it(
async_only_image_fetch,
):
shared = f"https://img.example/{uuid.uuid4()}.png"
messages = [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": shared, "format": "image/png"}},
{"type": "image_url", "image_url": {"url": shared}},
],
}
]
inlined = await async_inline_remote_media(messages, should_inline=lambda media: "format" not in media.fields)
assert inlined[0]["content"] == [
{"type": "image_url", "image_url": {"url": shared, "format": "image/png"}},
{"type": "image_url", "image_url": {"url": async_only_image_fetch.data_url}},
]
assert async_only_image_fetch.fetched == [shared]
async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fails(monkeypatch):
missing = f"http://img.example/{uuid.uuid4()}-missing.png"
slow = f"http://img.example/{uuid.uuid4()}-slow.png"
slow_fetch_outcomes = []
async def serve(client, url, **kwargs):
if url == missing:
return Response(404, request=Request("GET", url))
try:
await asyncio.sleep(5)
except asyncio.CancelledError:
slow_fetch_outcomes.append("cancelled")
raise
slow_fetch_outcomes.append("finished")
return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url))
monkeypatch.setattr(image_handling, "async_safe_get", serve)
messages = [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": missing}},
{"type": "image_url", "image_url": {"url": slow}},
],
}
]
started = time.perf_counter()
with pytest.raises(litellm.ImageFetchError, match="Status code: 404"):
await async_inline_remote_media(messages)
assert slow_fetch_outcomes == ["cancelled"]
assert time.perf_counter() - started < 1
_SSRF_VERDICTS = (
SSRFError(
"URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, "
"add the host to `user_url_allowed_hosts` in general_settings."
),
SSRFError("DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known"),
SSRFError("No addresses found for 'internal.example'"),
)
def _assert_verdict_free_messages(messages, url):
assert len(messages) == len(_SSRF_VERDICTS)
assert len(set(messages)) == 1, "a caller must not be able to tell a blocked host from one that does not resolve"
message = messages[0]
assert "The proxy could not resolve this host or its URL policy rejected it" in message
assert "user_url_allowed_hosts" in message
assert url in message
assert "10.0.0.8" not in message
assert "DNS" not in message
assert "No addresses" not in message
async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch):
attempts = []
messages = []
url = f"http://internal.example/{uuid.uuid4()}.png"
for verdict in _SSRF_VERDICTS:
async def block(client, fetched_url, verdict=verdict, **kwargs):
attempts.append(fetched_url)
raise verdict
monkeypatch.setattr(image_handling, "async_safe_get", block)
with pytest.raises(litellm.ImageFetchError) as raised:
await async_convert_url_to_base64(url)
messages.append(raised.value.message)
assert attempts == [url] * len(_SSRF_VERDICTS)
_assert_verdict_free_messages(messages, url)
def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch):
attempts = []
messages = []
url = f"http://internal.example/{uuid.uuid4()}.png"
for verdict in _SSRF_VERDICTS:
def block(client, fetched_url, verdict=verdict, **kwargs):
attempts.append(fetched_url)
raise verdict
monkeypatch.setattr(image_handling, "safe_get", block)
with pytest.raises(litellm.ImageFetchError) as raised:
convert_url_to_base64(url)
messages.append(raised.value.message)
assert attempts == [url] * len(_SSRF_VERDICTS)
_assert_verdict_free_messages(messages, url)
async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch):
in_flight = {"now": 0, "peak": 0}
async def serve_png_slowly(client, url, **kwargs):
in_flight["now"] += 1
in_flight["peak"] = max(in_flight["peak"], in_flight["now"])
await asyncio.sleep(0.01)
in_flight["now"] -= 1
return Response(200, content=b"\x89PNG", headers={"content-type": "image/png"}, request=Request("GET", url))
monkeypatch.setattr(image_handling, "async_safe_get", serve_png_slowly)
urls = [f"https://img.example/{uuid.uuid4()}.png" for _ in range(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES + 5)]
messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}} for url in urls]}]
inlined = await async_inline_remote_media(messages)
assert in_flight["peak"] == MAX_CONCURRENT_REMOTE_MEDIA_FETCHES
assert all(part["image_url"]["url"].startswith("data:image/png;base64,") for part in inlined[0]["content"])
async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch):
messages = [
{"role": "user", "content": [{"type": "text", "text": "hi"}]},
{
"role": "user",
"content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}],
},
]
assert await async_inline_remote_media(messages) is messages
assert async_only_image_fetch.fetched == []
async def test_async_inline_remote_media_raises_image_fetch_error_when_the_fetch_fails(monkeypatch):
async def serve_404(client, url, **kwargs):
return Response(404, request=Request("GET", url))
monkeypatch.setattr(image_handling, "async_safe_get", serve_404)
url = f"http://img.example/{uuid.uuid4()}.png"
with pytest.raises(litellm.ImageFetchError, match="Status code: 404"):
await async_inline_remote_media([{"role": "user", "content": [{"type": "image_url", "image_url": url}]}])

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

@ -1,7 +1,9 @@
import asyncio
import contextlib
import datetime
import os
import sys
import asyncio
from typing import Literal
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -632,6 +634,86 @@ class TestRetrieveBatchCostPassesModelIdentity:
assert captured["model_info"]["input_cost_per_token"] == 0.0
class TestRetrieveBatchPricesOnlyFinalBatches:
"""Regression (LIT-7048): retrieving a provider-id batch priced it on every poll.
Every retrieve of one batch logs under the same spend row, so pricing a poll
that landed before the output existed wrote that row at $0 and pinned it there.
Only a final batch gets priced; an in-flight poll carries no cost at all.
"""
@staticmethod
def _logging_obj() -> LitellmLogging:
obj = LitellmLogging(
model="gpt-5.6-luna",
messages=[{"role": "user", "content": "Hey"}],
stream=False,
call_type="aretrieve_batch",
start_time=time.time(),
litellm_call_id="batch-call-2",
function_id="f",
)
obj.custom_llm_provider = "openai"
return obj
@staticmethod
def _batch(status: str, output_file_id: str | None):
from litellm.types.utils import LiteLLMBatch
return LiteLLMBatch(
id="batch_6a9c99e185588190877d391f8b9d7f8a",
completion_window="24h",
created_at=1,
endpoint="/v1/chat/completions",
input_file_id="file-in",
object="batch",
status="validating",
output_file_id=output_file_id,
).model_copy(update={"status": status})
@pytest.mark.asyncio
@pytest.mark.parametrize(
("status", "output_file_id"),
[("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None), ("complete", None)],
)
async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_file_id) -> None:
from litellm.litellm_core_utils import litellm_logging as logging_module
handle_completed_batch = AsyncMock()
monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch)
batch = self._batch(status, output_file_id)
await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None)
handle_completed_batch.assert_not_awaited()
assert "response_cost" not in batch._hidden_params
@pytest.mark.asyncio
async def test_completed_batch_with_output_is_priced(self, monkeypatch) -> None:
from litellm.batches.batch_utils import BatchCostUsageResult
from litellm.litellm_core_utils import litellm_logging as logging_module
from litellm.types.utils import Usage
handle_completed_batch = AsyncMock(
return_value=BatchCostUsageResult(
cost=8e-06,
usage=Usage(prompt_tokens=26, completion_tokens=9, total_tokens=35),
models=["gpt-5.6-luna"],
successful_requests=2,
failed_requests=0,
)
)
monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch)
batch = self._batch("completed", "file-out")
await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None)
handle_completed_batch.assert_awaited_once()
assert batch._hidden_params["response_cost"] == 8e-06
assert batch.usage is not None
assert batch.usage.total_tokens == 35
class TestAnthropicPassthroughCustomPricing:
"""Verify the Anthropic pass-through handler forwards custom pricing."""
@ -4391,6 +4473,39 @@ def test_handle_anthropic_messages_response_logging_translates_bare_responses_ap
assert result.usage.total_tokens == 18 # type: ignore[attr-defined]
def test_handle_anthropic_messages_response_logging_keeps_the_served_response_id():
from openai.types.responses import ResponseOutputMessage, ResponseOutputText
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
served_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id="deployment-1", response_id="resp_upstream"
)
logging_obj = _anthropic_messages_logging_obj()
result = logging_obj._handle_anthropic_messages_response_logging(
result=ResponsesAPIResponse(
id=served_id,
created_at=1700000000,
output=[
ResponseOutputMessage(
id="msg-1",
type="message",
role="assistant",
status="completed",
content=[ResponseOutputText(annotations=[], text="hi", type="output_text")],
)
],
usage=ResponseAPIUsage(input_tokens=2, output_tokens=1, total_tokens=3),
service_tier="flex",
)
)
assert isinstance(result, ModelResponse)
assert result.id == served_id, "the spend log row must keep the id the caller was served"
assert result.service_tier == "flex"
def test_handle_anthropic_messages_response_logging_passes_model_response_through():
"""Anthropic-native path already yields a ModelResponse; it must be returned unchanged."""
logging_obj = _anthropic_messages_logging_obj()
@ -6326,6 +6441,113 @@ def test_get_standard_logging_object_payload_survives_logging_obj_without_timing
assert payload["hidden_params"]["litellm_overhead_time_ms"] is None
@pytest.mark.parametrize(
("header_name", "header_source"),
(
("x-amzn-RequestId", "response"),
("x-request-id", "response"),
("request-id", "response"),
("x-ms-request-id", "response"),
("apim-request-id", "response"),
("x-goog-request-id", "response"),
("cf-ray", "response"),
("X-Request-Id", "litellm_response_headers"),
("X-MS-Request-ID", "headers"),
),
)
def test_failure_standard_logging_payload_captures_provider_request_id(
logging_obj: LitellmLogging,
header_name: str,
header_source: Literal["response", "litellm_response_headers", "headers"],
):
from litellm.litellm_core_utils.litellm_logging import (
get_standard_logging_object_payload,
)
request_id = "provider-request-123"
request = httpx.Request("POST", "https://provider.example/v1/chat/completions")
response = httpx.Response(429, headers={header_name: request_id}, request=request)
provider_error = httpx.HTTPStatusError("provider error", request=request, response=response)
if header_source == "litellm_response_headers":
response.headers.clear()
provider_error.litellm_response_headers = {header_name: request_id}
elif header_source == "headers":
response.headers.clear()
provider_error.headers = {header_name: request_id}
now = datetime.datetime.now()
payload = get_standard_logging_object_payload(
kwargs={"litellm_call_id": "call-1", "model": "test-model", "messages": []},
init_response_obj={},
start_time=now,
end_time=now,
logging_obj=logging_obj,
status="failure",
original_exception=provider_error,
)
assert payload is not None
assert payload["error_information"] is not None
assert payload["error_information"]["error_provider_request_id"] == request_id
def test_get_error_information_ignores_unsupported_headers() -> None:
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
request = httpx.Request("POST", "https://provider.example/v1/chat/completions")
response = httpx.Response(429, headers={"retry-after": "3"}, request=request)
provider_error = httpx.HTTPStatusError("provider error", request=request, response=response)
error_information = StandardLoggingPayloadSetup.get_error_information(provider_error)
assert error_information["error_provider_request_id"] is None
def test_get_error_information_uses_header_precedence_and_fallback() -> None:
from litellm.exceptions import RateLimitError
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
request = httpx.Request("POST", "https://provider.example/v1/chat/completions")
response = httpx.Response(
429,
headers={"x-request-id": "response-id", "x-amzn-requestid": "amazon-id"},
request=request,
)
provider_error = RateLimitError(
message="provider error",
llm_provider="test-provider",
model="test-model",
response=response,
headers={"retry-after": "3"},
)
error_information = StandardLoggingPayloadSetup.get_error_information(provider_error)
assert error_information["error_provider_request_id"] == "amazon-id"
def test_get_error_information_ignores_malformed_headers() -> None:
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
provider_error = Exception("provider error")
provider_error.headers = [("x-request-id", "provider-request-123")]
error_information = StandardLoggingPayloadSetup.get_error_information(provider_error)
assert error_information["error_provider_request_id"] is None
def test_get_provider_request_id_ignores_header_lookup_errors() -> None:
from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id
class HeaderLookupError(Exception):
@property
def response(self) -> object:
raise RuntimeError("headers unavailable")
assert _get_provider_request_id(HeaderLookupError("provider error")) is None
def test_get_standard_logging_object_payload_failure_status_keeps_overhead_none(logging_obj):
"""A post_call guardrail can fail the request after the upstream call succeeded; the failure
payload keeps litellm_overhead_time_ms None, matching responses that carry their own _hidden_params."""

35
tests/test_litellm/litellm_core_utils/test_url_utils.py
View file

@ -1,5 +1,9 @@
import asyncio
import socket
import threading
import time
import httpx
import pytest
import litellm
@ -535,3 +539,34 @@ def test_assert_same_origin_error_message_does_not_leak_hostnames():
detail = str(exc.value)
assert "attacker.example.com" not in detail
assert "api.internal-corp.example" not in detail
async def test_async_safe_get_resolves_dns_off_the_event_loop(monkeypatch):
loop_thread = threading.current_thread()
resolver_threads = []
def slow_getaddrinfo(host, port, *args, **kwargs):
resolver_threads.append(threading.current_thread())
time.sleep(0.4)
return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port or 443))]
monkeypatch.setattr(url_utils.socket, "getaddrinfo", slow_getaddrinfo)
class FakeClient:
async def get(self, url, **kwargs):
return httpx.Response(200, request=httpx.Request("GET", url))
ticks = [time.perf_counter()]
async def heartbeat():
while True:
await asyncio.sleep(0.01)
ticks.append(time.perf_counter())
beating = asyncio.create_task(heartbeat())
response = await url_utils.async_safe_get(FakeClient(), "https://img.example/a.png")
beating.cancel()
assert response.status_code == 200
assert resolver_threads and all(thread is not loop_thread for thread in resolver_threads)
assert max(b - a for a, b in zip(ticks, ticks[1:])) < 0.2

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