scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

refactor(rust-error): keep causes as sources in llms and the gateway crates

Browse source 
ErrorDetail loses its string variant, OCR provider variants become data,
host faults stop counting as the caller's, a tunnelled transport error
classifies the same on every stream path, and each gateway crate has one
Error with one status() and extractors that reject in its envelope.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-10-23 12:59:21 -07:00
parent 4432805bf7
commit 204f146dc9
101 changed files with 1732 additions and 1177 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm-rust/Cargo.lock generated
View file

@ -3764,6 +3764,7 @@ dependencies = [
"http-body-util",
"litellm-auth-types",
"litellm-config",
"litellm-error",
"litellm-gateway-auth",
"litellm-gateway-inference",
"litellm-gateway-mcp",
@ -3819,6 +3820,7 @@ dependencies = [
"litellm-auth",
"litellm-cache-memory",
"litellm-cache-response",
"litellm-error",
"litellm-gateway-auth",
"litellm-host",
"litellm-host-http",
@ -3849,6 +3851,7 @@ version = "0.1.0"
dependencies = [
"base64 0.22.1",
"litellm-auth-types",
"litellm-error",
"litellm-gateway-auth",
"rand 0.8.7",
"rstest",
@ -3866,6 +3869,7 @@ dependencies = [
"http 1.4.2",
"litellm-auth-types",
"litellm-config",
"litellm-error",
"litellm-secrets",
"moka",
"rmcp",
@ -3880,6 +3884,7 @@ dependencies = [
"tower",
"url",
"uuid",
"veil",
]
[[package]]
@ -3889,9 +3894,11 @@ dependencies = [
"axum",
"axum-login",
"base64 0.22.1",
"futures-util",
"governor",
"jsonwebtoken",
"litellm-auth-types",
"litellm-error",
"litellm-gateway-auth",
"rand 0.8.7",
"rstest",
@ -3916,6 +3923,7 @@ dependencies = [
"litellm-coroutine",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
]
@ -4004,6 +4012,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-error",
"litellm-framer",
"litellm-host",
"litellm-http",
@ -4058,6 +4067,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-error",
"litellm-framer",
"litellm-host",
"litellm-host-native",

4
litellm-rust/crates/gateway-auth/AGENTS.md
View file

@ -20,7 +20,9 @@ Local UI login continues using axum-login and tower-sessions. Only after session
Authentication evidence and principals contain no raw token, password, request body, or mutable accounting state. Session ownership is scoped by principal authority, subject, verifier, and credential ID, so separate credentials do not silently share an MCP session. Credential rotation may retain ownership when the verifier preserves a stable credential ID. Scope and expiry checks still run for each operation
Failures distinguish invalid or expired credentials, forbidden operations, unavailable authentication services, and missing server configuration/context. HTTP adapters map those outcomes to status codes. Inference keeps its API-specific error envelopes and MCP translates policy failures to its own protocol
Failures distinguish invalid or expired credentials, forbidden operations, unavailable authentication services, and missing server configuration/context. Bearer tokens, virtual keys and UI sessions all fail with the one `Error`, classified once by `status()`, and `is_denied()` tells a rejected caller from an auth failure. Inference keeps its API-specific error envelopes, the UI renders `Error` in its own envelope and MCP translates policy failures to its own protocol
The plain-text rejection body is the top-level message only, never the `Report` chain, because the caller is unauthenticated and the sources carry secret-manager and backend detail. That stays reachable through `source()` for logs
Virtual-key storage, JWT verification, OAuth2 introspection, trusted-proxy validation, SSO, and custom Python hook adapters are not implemented here yet. They should supply these contracts rather than bypassing the shared authorization boundary. Request-body-dependent custom hooks will need a bounded endpoint adapter after parsing

2
litellm-rust/crates/gateway-auth/Cargo.toml
View file

@ -10,6 +10,7 @@ axum = { workspace = true, features = ["matched-path"] }
axum-login.workspace = true
litellm-auth-types.workspace = true
litellm-config.workspace = true
litellm-error.workspace = true
litellm-secrets.workspace = true
sha2.workspace = true
subtle.workspace = true
@ -20,7 +21,6 @@ veil.workspace = true
[dev-dependencies]
futures-util.workspace = true
litellm-error.workspace = true
rstest.workspace = true
tokio.workspace = true
tower = { version = "0.5.3", features = ["util"] }

108
litellm-rust/crates/gateway-auth/src/error.rs
View file

@ -2,44 +2,59 @@ use axum::{
http::StatusCode,
response::{IntoResponse, Response},
};
#[derive(Debug, thiserror::Error)]
pub enum KeyError {
#[error("missing or invalid virtual key")]
Invalid,
#[error("virtual key has expired")]
Expired,
#[error("virtual key lookup unavailable")]
Lookup(#[source] Box<dyn std::error::Error + Send + Sync>),
}
use litellm_error::Source;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("operation is not permitted")]
Forbidden,
#[error("credential has expired")]
Expired,
#[error("authenticated request context is missing")]
MissingIdentity,
#[error("authentication service unavailable")]
Unavailable,
#[error("gateway auth not configured")]
Unconfigured,
#[error("missing or invalid bearer token")]
InvalidToken,
#[error("invalid UI credentials or session")]
InvalidSession,
#[error("credential has expired")]
Expired,
#[error("operation is not permitted")]
Forbidden,
#[error("authentication service unavailable")]
Unavailable(#[source] Source),
#[error("gateway authentication unavailable")]
Secret(#[from] litellm_secrets::Error),
#[error("UI session unavailable")]
Session(#[from] tower_sessions::session::Error),
#[error("gateway auth not configured")]
Unconfigured,
#[error("UI_USERNAME and UI_PASSWORD must be nonempty")]
UiUnconfigured,
#[error("authenticated request context is missing")]
MissingIdentity,
#[error("UI session layer is not installed")]
MissingSessionLayer,
}
impl Error {
pub fn status(&self) -> StatusCode {
match self {
Self::InvalidToken | Self::Expired => StatusCode::UNAUTHORIZED,
Self::InvalidToken | Self::InvalidSession | Self::Expired => StatusCode::UNAUTHORIZED,
Self::Forbidden => StatusCode::FORBIDDEN,
Self::Unavailable => StatusCode::SERVICE_UNAVAILABLE,
Self::MissingIdentity | Self::Unconfigured | Self::Secret(_) => {
StatusCode::INTERNAL_SERVER_ERROR
}
Self::Unavailable(_) => StatusCode::SERVICE_UNAVAILABLE,
Self::Secret(_)
| Self::Session(_)
| Self::Unconfigured
| Self::UiUnconfigured
| Self::MissingIdentity
| Self::MissingSessionLayer => StatusCode::INTERNAL_SERVER_ERROR,
}
}
pub fn is_denied(&self) -> bool {
match self {
Self::InvalidToken | Self::InvalidSession | Self::Expired | Self::Forbidden => true,
Self::Unavailable(_)
| Self::Secret(_)
| Self::Session(_)
| Self::Unconfigured
| Self::UiUnconfigured
| Self::MissingIdentity
| Self::MissingSessionLayer => false,
}
}
}
@ -50,24 +65,31 @@ impl IntoResponse for Error {
}
}
#[derive(Debug, thiserror::Error)]
pub enum UiAuthError {
#[error("UI_USERNAME and UI_PASSWORD must be nonempty")]
Unconfigured,
#[error("invalid UI credentials or session")]
Unauthorized,
#[error("UI authentication unavailable")]
Unavailable,
#[error("UI session unavailable")]
Session(#[from] tower_sessions::session::Error),
}
#[cfg(test)]
mod tests {
use rstest::rstest;
impl IntoResponse for UiAuthError {
fn into_response(self) -> Response {
let status = match self {
Self::Unauthorized => StatusCode::UNAUTHORIZED,
_ => StatusCode::INTERNAL_SERVER_ERROR,
};
(status, self.to_string()).into_response()
use super::*;
#[rstest]
#[tokio::test]
async fn rejection_body_hides_backend_causes() {
let error =
Error::Unavailable(Source::new(std::io::Error::other("private-backend-detail")));
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("private-backend-detail")
);
let response = error.into_response();
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
let body = axum::body::to_bytes(response.into_body(), 4096)
.await
.unwrap();
assert!(
!std::str::from_utf8(&body)
.unwrap()
.contains("private-backend-detail")
);
}
}

10
litellm-rust/crates/gateway-auth/src/keys.rs
View file

@ -2,7 +2,7 @@ use std::{future::Future, pin::Pin, time::SystemTime};
use litellm_auth_types::SecretValue;
use crate::{KeyError, hash_token};
use crate::{Error, hash_token};
#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct KeyHash(String);
@ -24,7 +24,7 @@ pub enum KeyStatus {
}
pub type LookupFuture<'a> =
Pin<Box<dyn Future<Output = Result<Option<KeyStatus>, KeyError>> + Send + 'a>>;
Pin<Box<dyn Future<Output = Result<Option<KeyStatus>, Error>> + Send + 'a>>;
pub trait KeyLookup: Send + Sync {
fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a>;
@ -34,15 +34,15 @@ pub async fn verify_key(
lookup: &dyn KeyLookup,
token: &SecretValue,
now: impl FnOnce() -> SystemTime,
) -> Result<KeyHash, KeyError> {
) -> Result<KeyHash, Error> {
let hash = KeyHash::from_token(token);
match lookup.lookup(&hash).await? {
Some(KeyStatus::Active { expires_at }) => {
if expires_at.is_some_and(|expiry| expiry <= now()) {
return Err(KeyError::Expired);
return Err(Error::Expired);
}
Ok(hash)
}
Some(KeyStatus::Revoked) | None => Err(KeyError::Invalid),
Some(KeyStatus::Revoked) | None => Err(Error::InvalidToken),
}
}

2
litellm-rust/crates/gateway-auth/src/lib.rs
View file

@ -16,7 +16,7 @@ pub use authorization::{
AccessRequest, AuthenticatedRequest, AuthorizedOperation, Authorizer, McpAction,
NoAdditionalPolicy, Permissions, UiAction,
};
pub use error::{Error, KeyError, UiAuthError};
pub use error::Error;
pub use http::{Bearer, CredentialExtractor, RequireMasterKey, authenticate};
pub use identity::{
AuthenticatedCaller, Authentication, AuthenticationMethod, Principal, PrincipalKind,

29
litellm-rust/crates/gateway-auth/src/ui.rs
View file

@ -11,9 +11,9 @@ use tower_sessions::Session;
use veil::Redact;
use crate::{
AccessRequest, AuthenticatedRequest, Authentication, AuthenticationMethod, LocalAdministrator,
NoAdditionalPolicy, Permissions, Principal, PrincipalKind, SystemClock, UiAction, UiAuthError,
VerifiedIdentity,
AccessRequest, AuthenticatedRequest, Authentication, AuthenticationMethod, Error,
LocalAdministrator, NoAdditionalPolicy, Permissions, Principal, PrincipalKind, SystemClock,
UiAction, VerifiedIdentity,
};
pub const UI_CSRF_KEY: &str = "litellm.ui.csrf";
@ -50,9 +50,9 @@ pub struct UiBackend {
}
impl UiBackend {
pub fn new(username: String, password: SecretValue) -> Result<Self, UiAuthError> {
pub fn new(username: String, password: SecretValue) -> Result<Self, Error> {
if username.trim().is_empty() || password.expose().trim().is_empty() {
return Err(UiAuthError::Unconfigured);
return Err(Error::UiUnconfigured);
}
Ok(Self {
user: UiUser {
@ -89,28 +89,28 @@ pub struct UiSession {
}
impl<S: Send + Sync> FromRequestParts<S> for UiSession {
type Rejection = UiAuthError;
type Rejection = Error;
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
let auth = UiAuthSession::from_request_parts(parts, state)
.await
.map_err(|_| UiAuthError::Unavailable)?;
let user = auth.user.ok_or(UiAuthError::Unauthorized)?;
.map_err(|_| Error::MissingSessionLayer)?;
let user = auth.user.ok_or(Error::InvalidSession)?;
let session = Session::from_request_parts(parts, state)
.await
.map_err(|_| UiAuthError::Unavailable)?;
.map_err(|_| Error::MissingSessionLayer)?;
let expected = session
.get::<String>(UI_CSRF_KEY)
.await?
.ok_or(UiAuthError::Unauthorized)?;
.ok_or(Error::InvalidSession)?;
let provided = parts
.headers
.get(AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.strip_prefix("Bearer "))
.ok_or(UiAuthError::Unauthorized)?;
.ok_or(Error::InvalidSession)?;
if !bool::from(Sha256::digest(expected).ct_eq(&Sha256::digest(provided))) {
return Err(UiAuthError::Unauthorized);
return Err(Error::InvalidSession);
}
let identity = session_identity(&user).await?;
parts.extensions.insert(identity.clone());
@ -118,7 +118,7 @@ impl<S: Send + Sync> FromRequestParts<S> for UiSession {
}
}
async fn session_identity(user: &UiUser) -> Result<AuthenticatedRequest, UiAuthError> {
async fn session_identity(user: &UiUser) -> Result<AuthenticatedRequest, Error> {
crate::authentication::resolve(
VerifiedIdentity {
principal: Principal::new(
@ -142,7 +142,6 @@ async fn session_identity(user: &UiUser) -> Result<AuthenticatedRequest, UiAuthE
Arc::new(SystemClock),
)
.await
.map_err(|_| UiAuthError::Unauthorized)
}
#[cfg(test)]
@ -180,7 +179,7 @@ mod tests {
deployment: "provider/model".into(),
})
.await,
Err(crate::Error::Forbidden)
Err(Error::Forbidden)
));
}
}

6
litellm-rust/crates/gateway-auth/tests/auth.rs
View file

@ -132,7 +132,11 @@ impl Authenticator for Verifier {
Box::pin(async move {
match self.failure {
Some(401) => return Err(Error::InvalidToken),
Some(_) => return Err(Error::Unavailable),
Some(_) => {
return Err(Error::Unavailable(litellm_error::Source::new(
std::io::Error::other("verifier offline"),
)));
}
None => (),
}
let litellm_gateway_auth::Credential::Token(credential) = credential else {

22
litellm-rust/crates/gateway-auth/tests/keys.rs
View file

@ -2,7 +2,7 @@ use std::time::{Duration, SystemTime};
use litellm_auth_types::SecretValue;
use litellm_gateway_auth::{
KeyError,
Error,
keys::{KeyHash, KeyLookup, KeyStatus, LookupFuture, verify_key},
};
use rstest::{fixture, rstest};
@ -59,7 +59,7 @@ async fn checks_expiration(
match valid {
true => assert_eq!(result.unwrap(), hash),
false => assert!(matches!(result, Err(KeyError::Expired))),
false => assert!(matches!(result, Err(Error::Expired))),
}
}
@ -79,7 +79,7 @@ async fn rejects_unusable_keys(
assert!(matches!(
verify_key(&lookup, &token, || now).await,
Err(KeyError::Invalid)
Err(Error::InvalidToken)
));
}
@ -93,7 +93,7 @@ async fn rejects_a_different_token(token: SecretValue, now: SystemTime) {
assert!(matches!(
verify_key(&lookup, &SecretValue::new("sk-other-credential"), || now).await,
Err(KeyError::Invalid)
Err(Error::InvalidToken)
));
}
@ -102,9 +102,9 @@ struct Unavailable;
impl KeyLookup for Unavailable {
fn lookup<'a>(&'a self, _: &'a KeyHash) -> LookupFuture<'a> {
Box::pin(async {
Err(KeyError::Lookup(Box::new(std::io::Error::other(
"backend detail",
))))
Err(Error::Unavailable(litellm_error::Source::new(
std::io::Error::other("backend detail"),
)))
})
}
}
@ -114,10 +114,8 @@ impl KeyLookup for Unavailable {
async fn lookup_failure_preserves_source_without_exposing_it(token: SecretValue, now: SystemTime) {
let error = verify_key(&Unavailable, &token, || now).await.unwrap_err();
assert!(matches!(&error, KeyError::Lookup(_)));
assert_eq!(
std::error::Error::source(&error).unwrap().to_string(),
"backend detail"
);
assert_eq!(error.status(), axum::http::StatusCode::SERVICE_UNAVAILABLE);
let cause = litellm_error::find::<std::io::Error>(&error).expect("lookup cause");
assert_eq!(cause.to_string(), "backend detail");
assert!(!error.to_string().contains("backend detail"));
}

1
litellm-rust/crates/gateway-inference/Cargo.toml
View file

@ -11,6 +11,7 @@ axum = { workspace = true, features = ["json", "multipart", "original-uri"] }
base64.workspace = true
bytes.workspace = true
litellm-auth.workspace = true
litellm-error.workspace = true
litellm-gateway-auth.workspace = true
litellm-inference.workspace = true
litellm-inference-transcription.workspace = true

12
litellm-rust/crates/gateway-inference/src/error.rs
View file

@ -61,7 +61,7 @@ impl Error {
let status = self.status();
let message = match &self {
Self::UnknownModel(model) => format!("Invalid model name passed in model={model}"),
_ => self.to_string(),
_ => litellm_error::Report::new(&self).to_string(),
};
(
status,
@ -82,7 +82,7 @@ impl Error {
Self::Route(error) if let Some((_, body)) = error.upstream_response() => {
body.to_owned()
}
other => other.to_string(),
other => litellm_error::Report::new(other).to_string(),
};
let parsed = serde_json::from_str::<Value>(&raw).ok();
let envelope = match parsed {
@ -235,7 +235,13 @@ mod tests {
#[rstest]
#[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)]
#[case::rejected_request(Error::Route(InferenceError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)]
#[case::rejected_request(
Error::Route(InferenceError::InvalidRequest(litellm_llms::ErrorDetail::Expected {
subject: "top_k",
expected: "an integer",
})),
StatusCode::BAD_REQUEST
)]
#[case::missing_key(
Error::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey {
provider: "Anthropic",

16
litellm-rust/crates/gateway-inference/tests/routes.rs
View file

@ -125,7 +125,10 @@ async fn chat_errors_come_from_core(
.await;
assert_eq!(response.status(), status);
let body = support::json(response).await;
assert_eq!(body["error"]["message"], error.to_string());
assert_eq!(
body["error"]["message"],
litellm_error::Report::new(&error).to_string()
);
assert_eq!(body["error"]["code"], status);
assert!(upstream.received_requests().await.unwrap().is_empty());
}
@ -193,17 +196,20 @@ async fn transcription_aliases_reach_core_validation(#[case] path: &str) {
}
#[rstest]
#[case::no_extension("test")]
#[case::unsupported_extension("test.invalid")]
#[case::no_extension("test", json!(null))]
#[case::unsupported_extension("test.invalid", json!("invalid"))]
#[tokio::test]
async fn upload_audio_format_validation_matches_core(#[case] filename: &str) {
async fn upload_audio_format_validation_matches_core(
#[case] filename: &str,
#[case] format: Value,
) {
let upstream = MockServer::start().await;
let app = support::app("bedrock/test-model", &upstream.uri());
let path = "/v1/audio/transcriptions";
let expected = support::post(
app.clone(),
path,
json!({"model": "public/model", "audio": {"data": "YWJj", "format": null}}),
json!({"model": "public/model", "audio": {"data": "YWJj", "format": format}}),
)
.await;
let payload = format!(

1
litellm-rust/crates/gateway-management/Cargo.toml
View file

@ -8,6 +8,7 @@ repository.workspace = true
[dependencies]
base64.workspace = true
litellm-auth-types.workspace = true
litellm-error.workspace = true
litellm-gateway-auth.workspace = true
rand.workspace = true
thiserror.workspace = true

2
litellm-rust/crates/gateway-management/src/error.rs
View file

@ -5,7 +5,7 @@ pub enum Error {
#[error("key already exists")]
AlreadyExists,
#[error("key storage unavailable")]
Storage(#[source] Box<dyn std::error::Error + Send + Sync>),
Storage(#[source] litellm_error::Source),
#[error("key generation unavailable")]
Entropy(#[source] rand::Error),
}

10
litellm-rust/crates/gateway-management/tests/keys.rs
View file

@ -9,7 +9,7 @@ use std::{
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use litellm_auth_types::SecretValue;
use litellm_gateway_auth::{
KeyError, hash_token,
hash_token,
keys::{KeyHash, KeyLookup, KeyStatus, verify_key},
};
use litellm_gateway_management::{
@ -73,7 +73,7 @@ async fn generated_key_can_be_retrieved_verified_and_revoked(
);
assert!(matches!(
verify_key(lookup.as_ref(), &generated.token, || now).await,
Err(KeyError::Invalid)
Err(litellm_gateway_auth::Error::InvalidToken)
));
assert_eq!(keys.revoke(hash).await.unwrap(), Revocation::Revoked);
}
@ -137,9 +137,7 @@ async fn storage_failures_are_returned_for_all_operations(now: SystemTime) {
assert!(matches!(keys.get(&hash).await, Err(Error::Storage(_))));
let error = keys.revoke(&hash).await.unwrap_err();
assert!(matches!(&error, Error::Storage(_)));
assert_eq!(
std::error::Error::source(&error).unwrap().to_string(),
"backend detail"
);
let cause = litellm_error::find::<std::io::Error>(&error).expect("storage cause");
assert_eq!(cause.to_string(), "backend detail");
assert!(!error.to_string().contains("backend detail"));
}

15
litellm-rust/crates/gateway-management/tests/keys/support.rs
View file

@ -1,10 +1,7 @@
use std::{collections::BTreeMap, sync::Mutex};
use litellm_gateway_auth::keys::KeyStatus;
use litellm_gateway_auth::{
KeyError,
keys::{KeyHash, KeyLookup, LookupFuture},
};
use litellm_gateway_auth::keys::{KeyHash, KeyLookup, LookupFuture};
use litellm_gateway_management::{
Error,
keys::{KeyRecord, KeyStore, NewKey, Revocation, StoreFuture},
@ -56,7 +53,9 @@ impl KeyLookup for MemoryStore {
self.get(hash)
.await
.map(|record| record.map(|record| record.status))
.map_err(|error| KeyError::Lookup(Box::new(error)))
.map_err(|error| {
litellm_gateway_auth::Error::Unavailable(litellm_error::Source::new(error))
})
})
}
}
@ -64,9 +63,9 @@ impl KeyLookup for MemoryStore {
pub struct Unavailable;
fn unavailable<T>() -> Result<T, Error> {
Err(Error::Storage(Box::new(std::io::Error::other(
"backend detail",
))))
Err(Error::Storage(litellm_error::Source::new(
std::io::Error::other("backend detail"),
)))
}
impl KeyStore for Unavailable {

2
litellm-rust/crates/gateway-mcp/AGENTS.md
View file

@ -6,6 +6,8 @@ The HTTP router serves `/mcp`, `/{server}/mcp`, `/mcp/{server}`, `/mcp/{server}/
`NativeGateway` aggregates upstream tools, prompts, resources, and resource templates. Tool and prompt calls resolve registered server prefixes, including prefixes containing hyphens. Resource names receive prefixes while their URIs remain unchanged, and reads require exactly one selected server. Tool allowlists apply to discovery and execution. Upstream pagination is drained with cycle detection. Aggregate tool listings retain healthy servers and expose sanitized outcomes under `_meta["litellm.ai/server_outcomes"]`, while scoped REST listings preserve upstream failures
Every HTTP failure, including extractor rejections and the legacy SSE and session transport checks, is an `Error` with its status from `status()` and the `{"detail"}` envelope. The detail is the `litellm_error::Report` of the error, except for upstream and authorization failures, whose causes can carry upstream URLs and backend detail. Startup failures are `ConnectError`, whose `Debug` redacts the SDK's initialization cause for the same reason
## Integration
Construct `Server` entries from connected SDK peers and pass a `Registry` to `NativeGateway::new`. Keep their `RunningService` handles alive in the host and cancel them during shutdown. SDK peers can use stdio, Streamable HTTP, or another SDK transport. Connection establishment, outbound credentials, HTTP client pools, and reconnection remain owned by the host

2
litellm-rust/crates/gateway-mcp/Cargo.toml
View file

@ -12,6 +12,7 @@ futures-util.workspace = true
litellm-config.workspace = true
litellm-secrets.workspace = true
litellm-auth-types.workspace = true
litellm-error.workspace = true
base64.workspace = true
http.workspace = true
uuid.workspace = true
@ -25,6 +26,7 @@ serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true
tokio-util = "0.7"
veil.workspace = true
[dev-dependencies]
rstest.workspace = true

120
litellm-rust/crates/gateway-mcp/src/configured.rs
View file

@ -22,25 +22,25 @@ use rmcp::{
use sha2::{Digest, Sha256};
use tokio_util::sync::CancellationToken;
use crate::{ConnectError, Limits, NativeGateway, Registry, Server, ServerInfo};
use crate::{ConfigFault, ConnectError, Limits, NativeGateway, Registry, Server, ServerInfo};
pub struct ConfiguredGateway {
pub operations: Arc<NativeGateway>,
services: Vec<RunningService<RoleClient, ()>>,
}
fn invalid(server: &str, message: impl Into<String>) -> ConnectError {
fn invalid(server: &str, fault: ConfigFault) -> ConnectError {
ConnectError::Configuration {
server: server.into(),
message: message.into(),
fault,
}
}
fn timeout(name: &str, config: &McpServer) -> Result<Duration, ConnectError> {
let duration = Duration::try_from_secs_f64(config.timeout.unwrap_or(30.0))
.map_err(|_| invalid(name, "timeout must be finite and positive"))?;
.map_err(|source| invalid(name, ConfigFault::InvalidTimeout(source)))?;
if duration.is_zero() {
return Err(invalid(name, "timeout must be positive"));
return Err(invalid(name, ConfigFault::ZeroTimeout));
}
Ok(duration)
}
@ -75,27 +75,24 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> {
.as_ref()
.is_some_and(|alias| alias.trim().is_empty())
{
return Err(invalid(name, "server identifiers must not be blank"));
return Err(invalid(name, ConfigFault::BlankIdentifier));
}
if let Some(field) = config.unsupported.keys().next() {
return Err(invalid(
name,
format!("setting '{field}' is not supported by the Rust MCP host"),
ConfigFault::UnsupportedSetting(field.clone()),
));
}
if config
.max_concurrent_requests
.is_some_and(|limit| limit == 0 || limit > tokio::sync::Semaphore::MAX_PERMITS)
{
return Err(invalid(name, "max_concurrent_requests is out of range"));
return Err(invalid(name, ConfigFault::ConcurrencyOutOfRange));
}
if config.mcp_info.contains_key("protocol_version")
|| config.mcp_info.contains_key("mcp_server_cost_info")
{
return Err(invalid(
name,
"mcp_info protocol and cost settings are not supported",
));
return Err(invalid(name, ConfigFault::UnsupportedMcpInfo));
}
if !matches!(
config.auth_type,
@ -108,17 +105,11 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> {
| McpAuth::Token
)
) {
return Err(invalid(
name,
"configured upstream auth mode is not supported by the Rust MCP host",
));
return Err(invalid(name, ConfigFault::UnsupportedAuth));
}
match config.transport {
McpTransport::Sse => {
return Err(invalid(
name,
"legacy SSE upstreams are not supported by the pinned SDK; use transport: http",
));
return Err(invalid(name, ConfigFault::SseTransport));
}
McpTransport::Http => {
if config.url.is_none()
@ -126,10 +117,7 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> {
|| !config.args.is_empty()
|| !config.env.is_empty()
{
return Err(invalid(
name,
"HTTP transport requires url and does not accept command, args or env",
));
return Err(invalid(name, ConfigFault::HttpShape));
}
}
McpTransport::Stdio => {
@ -139,53 +127,49 @@ fn validate(name: &str, config: &McpServer) -> Result<(), ConnectError> {
.is_none_or(|command| command.trim().is_empty())
|| config.url.is_some()
{
return Err(invalid(
name,
"stdio transport requires command and does not accept url",
));
return Err(invalid(name, ConfigFault::StdioShape));
}
if config.authentication_token.is_some()
|| !config.static_headers.is_empty()
|| config.upstream_token_header.is_some()
|| !matches!(config.auth_type, None | Some(McpAuth::None))
{
return Err(invalid(
name,
"stdio credentials must be supplied through env",
));
return Err(invalid(name, ConfigFault::StdioCredentials));
}
}
}
if matches!(config.auth_type, None | Some(McpAuth::None)) {
if config.authentication_token.is_some() || config.upstream_token_header.is_some() {
return Err(invalid(
name,
"authentication_token and upstream_token_header require an auth_type",
));
return Err(invalid(name, ConfigFault::TokenWithoutAuth));
}
} else if config.authentication_token.is_none() {
return Err(invalid(
name,
"configured auth_type requires authentication_token",
));
return Err(invalid(name, ConfigFault::AuthWithoutToken));
}
Ok(())
}
async fn resolve(value: &SecretValue, secrets: &dyn SecretSource) -> Result<String, ConnectError> {
async fn resolve(
name: &str,
value: &SecretValue,
secrets: &dyn SecretSource,
) -> Result<String, ConnectError> {
match value.expose().strip_prefix("os.environ/") {
Some(key) => secrets
.get_secret_str(key)
.await?
.map(|secret| secret.expose().to_owned())
.ok_or_else(|| invalid("configuration", "referenced secret is unavailable")),
.ok_or_else(|| invalid(name, ConfigFault::MissingSecret)),
None => Ok(value.expose().to_owned()),
}
}
fn header(name: &str, value: &str) -> Result<(HeaderName, HeaderValue), ConnectError> {
fn header(
server: &str,
name: &str,
value: &str,
) -> Result<(HeaderName, HeaderValue), ConnectError> {
let key = HeaderName::try_from(name)
.map_err(|_| invalid("configuration", "invalid upstream header name"))?;
.map_err(|source| invalid(server, ConfigFault::HeaderName(source)))?;
if matches!(
key.as_str(),
"host"
@ -198,13 +182,10 @@ fn header(name: &str, value: &str) -> Result<(HeaderName, HeaderValue), ConnectE
| "mcp-protocol-version"
| "last-event-id"
) {
return Err(invalid(
"configuration",
"upstream header conflicts with HTTP framing",
));
return Err(invalid(server, ConfigFault::FramingHeader));
}
let mut value = HeaderValue::try_from(value)
.map_err(|_| invalid("configuration", "invalid upstream header value"))?;
.map_err(|source| invalid(server, ConfigFault::HeaderValue(source)))?;
value.set_sensitive(true);
Ok((key, value))
}
@ -231,21 +212,19 @@ fn basic(value: &str) -> String {
}
async fn headers(
name: &str,
config: &McpServer,
secrets: &dyn SecretSource,
) -> Result<HashMap<HeaderName, HeaderValue>, ConnectError> {
let entries = try_join_all(
config
.static_headers
.iter()
.map(|(key, value)| async move { header(key, &resolve(value, secrets).await?) }),
)
let entries = try_join_all(config.static_headers.iter().map(|(key, value)| async move {
header(name, key, &resolve(name, value, secrets).await?)
}))
.await?;
let auth = match (&config.authentication_token, config.auth_type) {
(Some(token), Some(kind)) => {
let token = resolve(token, secrets).await?;
let token = resolve(name, token, secrets).await?;
if token.trim().is_empty() {
return Err(invalid("configuration", "upstream credential is empty"));
return Err(invalid(name, ConfigFault::EmptyCredential));
}
let (slot, value) = match kind {
McpAuth::ApiKey => ("x-api-key", token),
@ -253,9 +232,10 @@ async fn headers(
McpAuth::Token => ("authorization", scheme(&token, "token")),
McpAuth::Basic => ("authorization", basic(&token)),
McpAuth::Authorization => ("authorization", token),
_ => return Err(invalid("configuration", "unsupported credential mode")),
_ => return Err(invalid(name, ConfigFault::UnsupportedAuth)),
};
Some(header(
name,
config.upstream_token_header.as_deref().unwrap_or(slot),
&value,
)?)
@ -349,36 +329,34 @@ async fn connect_service<C: StreamableHttpClient>(
let configured = config
.url
.as_ref()
.ok_or_else(|| invalid(name, "HTTP transport requires url"))?;
let url = resolve(configured, secrets).await?;
let parsed = url::Url::parse(&url).map_err(|_| invalid(name, "invalid HTTP URL"))?;
.ok_or_else(|| invalid(name, ConfigFault::HttpShape))?;
let url = resolve(name, configured, secrets).await?;
let parsed = url::Url::parse(&url)
.map_err(|source| invalid(name, ConfigFault::InvalidUrl(source)))?;
if !matches!(parsed.scheme(), "http" | "https")
|| parsed.host_str().is_none()
|| !parsed.username().is_empty()
|| parsed.password().is_some()
|| parsed.fragment().is_some()
{
return Err(invalid(
name,
"HTTP URL must use http or https without user credentials or fragments",
));
return Err(invalid(name, ConfigFault::UrlShape));
}
let transport = StreamableHttpClientTransport::with_client(
client,
StreamableHttpClientTransportConfig::with_uri(url)
.custom_headers(headers(config, secrets).await?),
.custom_headers(headers(name, config, secrets).await?),
);
().serve_with_ct(transport, shutdown).await
}
McpTransport::Stdio => {
let environment = try_join_all(config.env.iter().map(|(key, value)| async move {
Ok::<_, ConnectError>((key, resolve(value, secrets).await?))
Ok::<_, ConnectError>((key, resolve(name, value, secrets).await?))
}))
.await?;
let program = config
.command
.as_deref()
.ok_or_else(|| invalid(name, "stdio requires command"))?;
.ok_or_else(|| invalid(name, ConfigFault::StdioShape))?;
let mut command = tokio::process::Command::new(program);
command
.args(&config.args)
@ -387,7 +365,7 @@ async fn connect_service<C: StreamableHttpClient>(
().serve_with_ct(TokioChildProcess::new(command)?, shutdown)
.await
}
McpTransport::Sse => return Err(invalid(name, "unsupported transport")),
McpTransport::Sse => return Err(invalid(name, ConfigFault::SseTransport)),
};
initialized.map_err(|source| ConnectError::Initialize {
server: name.into(),
@ -452,7 +430,7 @@ mod tests {
"mcp_servers:\n test:\n auth_type: {mode}\n authentication_token: '{token}'\n"
))
.unwrap();
let built = headers(&config.mcp_servers["test"], &Secrets)
let built = headers("test", &config.mcp_servers["test"], &Secrets)
.await
.unwrap();
let key = HeaderName::try_from(slot).unwrap();
@ -464,7 +442,7 @@ mod tests {
#[tokio::test]
async fn explicit_credential_slot_overrides_static_header() {
let config = litellm_config::Config::from_yaml("mcp_servers: {test: {auth_type: api_key, authentication_token: os.environ/KEY, upstream_token_header: x-custom, static_headers: {X-Custom: old, x-extra: extra}}}").unwrap();
let built = headers(&config.mcp_servers["test"], &Secrets)
let built = headers("test", &config.mcp_servers["test"], &Secrets)
.await
.unwrap();
assert_eq!(

269
litellm-rust/crates/gateway-mcp/src/error.rs
View file

@ -1,58 +1,162 @@
use std::time::TryFromFloatSecsError;
use axum::{
Json,
extract::rejection::{JsonRejection, PathRejection, QueryRejection},
http::StatusCode,
response::{IntoResponse, Response},
};
use http::header::{InvalidHeaderName, InvalidHeaderValue, ToStrError};
use litellm_error::{Report, Source};
use rmcp::{
model::{ErrorCode, ErrorData},
service::ServiceError,
service::{ClientInitializeError, ServiceError},
};
use serde_json::json;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("{0}")]
InvalidRequest(String),
#[error("{0} is required in request body")]
MissingField(&'static str),
#[error("Gateway listings do not accept an upstream cursor")]
UpstreamCursor,
#[error("read_resource requires exactly one allowed MCP server")]
AmbiguousResourceServer,
#[error("Invalid x-mcp-servers header")]
InvalidServersHeader(#[source] ToStrError),
#[error("invalid JSON body")]
Json(#[source] JsonRejection),
#[error("invalid query string")]
Query(#[source] QueryRejection),
#[error("invalid path")]
Path(#[source] PathRejection),
#[error("request body is too large")]
BodyTooLarge(#[source] axum::Error),
#[error("User not allowed to access this MCP server or operation.")]
Forbidden,
#[error("{0}")]
Configuration(String),
#[error("request host is not allowed")]
HostNotAllowed,
#[error("request origin is not allowed")]
OriginNotAllowed,
#[error("MCP session belongs to another caller")]
ForeignSession,
#[error("MCP session not found")]
UnknownSession,
#[error("MCP session is busy")]
SessionBusy,
#[error("MCP upstream request failed")]
Upstream(#[from] ServiceError),
#[error("MCP upstream returned invalid pagination")]
InvalidPagination,
#[error("MCP operation returned an unexpected result")]
UnexpectedResult,
#[error("MCP authorization unavailable")]
AuthorizationUnavailable(#[source] Source),
#[error("MCP request cancelled")]
Cancelled,
}
impl Error {
pub fn into_mcp(self) -> ErrorData {
pub fn status(&self) -> StatusCode {
match self {
Self::Upstream(ServiceError::McpError(error)) => error,
Self::InvalidRequest(message) => ErrorData::invalid_params(message, None),
Self::Forbidden => ErrorData::new(ErrorCode(-32003), self.to_string(), None),
_ => ErrorData::internal_error(self.to_string(), None),
Self::MissingField(_)
| Self::UpstreamCursor
| Self::AmbiguousResourceServer
| Self::InvalidServersHeader(_) => StatusCode::BAD_REQUEST,
Self::Json(rejection) => rejection.status(),
Self::Query(rejection) => rejection.status(),
Self::Path(rejection) => rejection.status(),
Self::BodyTooLarge(_) => StatusCode::PAYLOAD_TOO_LARGE,
Self::Forbidden
| Self::HostNotAllowed
| Self::OriginNotAllowed
| Self::ForeignSession => StatusCode::FORBIDDEN,
Self::UnknownSession => StatusCode::NOT_FOUND,
Self::SessionBusy => StatusCode::TOO_MANY_REQUESTS,
Self::Upstream(_) | Self::InvalidPagination => StatusCode::BAD_GATEWAY,
Self::Cancelled => StatusCode::REQUEST_TIMEOUT,
Self::UnexpectedResult | Self::AuthorizationUnavailable(_) => {
StatusCode::INTERNAL_SERVER_ERROR
}
}
}
fn is_request(&self) -> bool {
match self {
Self::MissingField(_)
| Self::UpstreamCursor
| Self::AmbiguousResourceServer
| Self::InvalidServersHeader(_)
| Self::Json(_)
| Self::Query(_)
| Self::Path(_)
| Self::BodyTooLarge(_) => true,
Self::Forbidden
| Self::HostNotAllowed
| Self::OriginNotAllowed
| Self::ForeignSession
| Self::UnknownSession
| Self::SessionBusy
| Self::Upstream(_)
| Self::InvalidPagination
| Self::UnexpectedResult
| Self::AuthorizationUnavailable(_)
| Self::Cancelled => false,
}
}
fn detail(&self) -> String {
match self {
Self::Upstream(_) | Self::AuthorizationUnavailable(_) => self.to_string(),
Self::MissingField(_)
| Self::UpstreamCursor
| Self::AmbiguousResourceServer
| Self::InvalidServersHeader(_)
| Self::Json(_)
| Self::Query(_)
| Self::Path(_)
| Self::BodyTooLarge(_)
| Self::Forbidden
| Self::HostNotAllowed
| Self::OriginNotAllowed
| Self::ForeignSession
| Self::UnknownSession
| Self::SessionBusy
| Self::InvalidPagination
| Self::UnexpectedResult
| Self::Cancelled => Report::new(self).to_string(),
}
}
pub fn into_mcp(self) -> ErrorData {
if let Self::Upstream(ServiceError::McpError(error)) = self {
return error;
}
let message = self.detail();
if self.status() == StatusCode::FORBIDDEN {
ErrorData::new(ErrorCode(-32003), message, None)
} else if self.is_request() {
ErrorData::invalid_params(message, None)
} else {
ErrorData::internal_error(message, None)
}
}
}
impl IntoResponse for Error {
fn into_response(self) -> Response {
let status = match &self {
Self::InvalidRequest(_) => StatusCode::BAD_REQUEST,
Self::Forbidden => StatusCode::FORBIDDEN,
Self::Upstream(_) => StatusCode::BAD_GATEWAY,
Self::Cancelled => StatusCode::REQUEST_TIMEOUT,
Self::Configuration(_) | Self::UnexpectedResult => StatusCode::INTERNAL_SERVER_ERROR,
};
(status, Json(json!({"detail": self.to_string()}))).into_response()
(self.status(), Json(json!({"detail": self.detail()}))).into_response()
}
}
#[derive(thiserror::Error)]
#[derive(thiserror::Error, veil::Redact)]
pub enum ConnectError {
#[error("MCP server {server}: {message}")]
Configuration { server: String, message: String },
#[error("MCP server {server} is misconfigured")]
Configuration {
server: String,
#[source]
fault: ConfigFault,
},
#[error("could not resolve MCP credentials")]
Secret(#[from] litellm_secrets::Error),
#[error("could not start MCP child process")]
@ -61,38 +165,131 @@ pub enum ConnectError {
Initialize {
server: String,
#[source]
source: Box<rmcp::service::ClientInitializeError>,
#[redact]
source: Box<ClientInitializeError>,
},
#[error("MCP upstream {0} initialization timed out")]
Timeout(String),
#[error("MCP upstream startup cancelled")]
Cancelled,
#[error(transparent)]
Registry(#[from] Error),
#[error("MCP server identifiers must not be empty")]
EmptyIdentifier,
#[error("MCP server identifiers and prefixes must be unique")]
DuplicateIdentifier,
}
impl std::fmt::Debug for ConnectError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
std::fmt::Display::fmt(self, formatter)
}
#[derive(Debug, thiserror::Error)]
pub enum ConfigFault {
#[error("timeout must be finite and positive")]
InvalidTimeout(#[source] TryFromFloatSecsError),
#[error("timeout must be positive")]
ZeroTimeout,
#[error("server identifiers must not be blank")]
BlankIdentifier,
#[error("setting '{0}' is not supported by the Rust MCP host")]
UnsupportedSetting(String),
#[error("max_concurrent_requests is out of range")]
ConcurrencyOutOfRange,
#[error("mcp_info protocol and cost settings are not supported")]
UnsupportedMcpInfo,
#[error("configured upstream auth mode is not supported by the Rust MCP host")]
UnsupportedAuth,
#[error("legacy SSE upstreams are not supported by the pinned SDK; use transport: http")]
SseTransport,
#[error("HTTP transport requires url and does not accept command, args or env")]
HttpShape,
#[error("stdio transport requires command and does not accept url")]
StdioShape,
#[error("stdio credentials must be supplied through env")]
StdioCredentials,
#[error("authentication_token and upstream_token_header require an auth_type")]
TokenWithoutAuth,
#[error("configured auth_type requires authentication_token")]
AuthWithoutToken,
#[error("referenced secret is unavailable")]
MissingSecret,
#[error("upstream credential is empty")]
EmptyCredential,
#[error("invalid upstream header name")]
HeaderName(#[source] InvalidHeaderName),
#[error("invalid upstream header value")]
HeaderValue(#[source] InvalidHeaderValue),
#[error("upstream header conflicts with HTTP framing")]
FramingHeader,
#[error("invalid HTTP URL")]
InvalidUrl(#[source] url::ParseError),
#[error("HTTP URL must use http or https without user credentials or fragments")]
UrlShape,
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest::rstest]
const SECRET_URL: &str = "https://example.test/mcp?token=private-secret";
#[rstest]
fn startup_diagnostics_do_not_expose_upstream_details() {
let error = ConnectError::Initialize {
server: "docs".into(),
source: Box::new(rmcp::service::ClientInitializeError::ConnectionClosed(
"https://example.test/mcp?token=private-secret".into(),
)),
source: Box::new(ClientInitializeError::ConnectionClosed(SECRET_URL.into())),
};
assert_eq!(
format!("{error:?}"),
"could not initialize MCP upstream docs"
);
let debug = format!("{error:?}");
assert!(debug.contains("docs"));
assert!(!debug.contains("private-secret"));
assert!(std::error::Error::source(&error).is_some());
}
#[rstest]
fn configuration_faults_reach_the_startup_diagnostic() {
let error = ConnectError::Configuration {
server: "docs".into(),
fault: ConfigFault::ZeroTimeout,
};
assert!(format!("{error:?}").contains("ZeroTimeout"));
assert_eq!(
Report::new(&error).to_string(),
"MCP server docs is misconfigured: timeout must be positive"
);
}
#[rstest]
#[case::missing_field(Error::MissingField("server_id"), StatusCode::BAD_REQUEST)]
#[case::upstream_cursor(Error::UpstreamCursor, StatusCode::BAD_REQUEST)]
#[case::forbidden(Error::Forbidden, StatusCode::FORBIDDEN)]
#[case::foreign_session(Error::ForeignSession, StatusCode::FORBIDDEN)]
#[case::unknown_session(Error::UnknownSession, StatusCode::NOT_FOUND)]
#[case::busy_session(Error::SessionBusy, StatusCode::TOO_MANY_REQUESTS)]
#[case::invalid_pagination(Error::InvalidPagination, StatusCode::BAD_GATEWAY)]
#[case::upstream(
Error::Upstream(ServiceError::TransportClosed),
StatusCode::BAD_GATEWAY
)]
#[case::cancelled(Error::Cancelled, StatusCode::REQUEST_TIMEOUT)]
fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) {
assert_eq!(error.status(), status);
}
#[rstest]
#[case::request(Error::MissingField("name"), ErrorCode::INVALID_PARAMS)]
#[case::forbidden(Error::Forbidden, ErrorCode(-32003))]
#[case::upstream_fault(Error::InvalidPagination, ErrorCode::INTERNAL_ERROR)]
fn mcp_error_code_follows_the_fault(#[case] error: Error, #[case] code: ErrorCode) {
assert_eq!(error.into_mcp().code, code);
}
#[rstest]
#[case::upstream(Error::Upstream(ServiceError::Cancelled {
reason: Some(SECRET_URL.into())
}))]
#[case::authorization(Error::AuthorizationUnavailable(Source::new(std::io::Error::other(
SECRET_URL
))))]
fn private_causes_stay_out_of_client_messages(#[case] error: Error) {
assert!(Report::new(&error).to_string().contains("private-secret"));
assert!(!error.detail().contains("private-secret"));
assert!(!error.into_mcp().message.contains("private-secret"));
}
}

88
litellm-rust/crates/gateway-mcp/src/legacy.rs
View file

@ -1,8 +1,8 @@
use std::{sync::Arc, time::Duration};
use axum::{
Json, Router,
extract::{Query, Request, State},
Router,
extract::{FromRequest, Request, State},
http::{StatusCode, request::Parts},
response::{
IntoResponse, Response, Sse,
@ -18,7 +18,11 @@ use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
use crate::{Context, HttpConfig, McpServer, Operations, sessions::fingerprint};
use crate::{
Context, Error, HttpConfig, McpServer, Operations,
request::{Json, Query},
sessions::fingerprint,
};
#[derive(Clone)]
struct Session {
@ -63,13 +67,13 @@ pub(crate) fn router(
}
impl Legacy {
fn validate(&self, parts: &Parts) -> Result<(), StatusCode> {
fn validate(&self, parts: &Parts) -> Result<(), Error> {
let host = parts
.headers
.get("host")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<http::uri::Authority>().ok())
.ok_or(StatusCode::FORBIDDEN)?;
.ok_or(Error::HostNotAllowed)?;
let allowed = self.hosts.iter().any(|allowed| {
host.as_str().eq_ignore_ascii_case(allowed)
|| host
@ -78,17 +82,16 @@ impl Legacy {
.eq_ignore_ascii_case(allowed)
});
if !allowed {
return Err(StatusCode::FORBIDDEN);
return Err(Error::HostNotAllowed);
}
if let Some(origin) = parts.headers.get("origin") {
let origin = origin.to_str().map_err(|_| StatusCode::FORBIDDEN)?;
if !self
.origins
.iter()
.any(|allowed| same_origin(origin, allowed))
{
return Err(StatusCode::FORBIDDEN);
}
if let Some(origin) = parts.headers.get("origin")
&& !origin.to_str().is_ok_and(|origin| {
self.origins
.iter()
.any(|allowed| same_origin(origin, allowed))
})
{
return Err(Error::OriginNotAllowed);
}
Ok(())
}
@ -111,27 +114,22 @@ fn same_origin(value: &str, allowed: &str) -> bool {
}
}
async fn connect(State(state): State<Legacy>, request: Request) -> Response {
async fn connect(State(state): State<Legacy>, request: Request) -> Result<Response, Error> {
let owner = fingerprint(&request);
let (parts, _) = request.into_parts();
if let Err(status) = state.validate(&parts) {
return status.into_response();
}
if let Err(error) = state
state.validate(&parts)?;
state
.operations
.authorize(Context {
parts,
server: None,
mcp: None,
})
.await
{
return error.into_response();
}
.await?;
let connection = start_session(state, owner).await;
Sse::new(connection.events())
Ok(Sse::new(connection.events())
.keep_alive(KeepAlive::default())
.into_response()
.into_response())
}
struct Connection {
@ -216,29 +214,21 @@ async fn message(
State(state): State<Legacy>,
Query(query): Query<SessionQuery>,
request: Request,
) -> Response {
) -> Result<Response, Error> {
let owner = fingerprint(&request);
let (parts, body) = request.into_parts();
if let Err(status) = state.validate(&parts) {
return status.into_response();
}
let Some(session) = state.sessions.get(&query.session_id).await else {
return StatusCode::NOT_FOUND.into_response();
};
if session.owner != owner {
return StatusCode::FORBIDDEN.into_response();
}
let request = Request::from_parts(parts.clone(), body);
let Json(mut message) =
match <Json<ClientJsonRpcMessage> as axum::extract::FromRequest<()>>::from_request(
request,
&(),
)
state.validate(&parts)?;
let session = state
.sessions
.get(&query.session_id)
.await
{
Ok(message) => message,
Err(error) => return error.into_response(),
};
.ok_or(Error::UnknownSession)?;
if session.owner != owner {
return Err(Error::ForeignSession);
}
let Json(mut message) =
Json::<ClientJsonRpcMessage>::from_request(Request::from_parts(parts.clone(), body), &())
.await?;
match &mut message {
ClientJsonRpcMessage::Request(request) => {
request.request.extensions_mut().insert(parts);
@ -249,8 +239,8 @@ async fn message(
_ => (),
}
match session.sender.try_send(message) {
Ok(()) => (StatusCode::ACCEPTED, "Accepted").into_response(),
Err(mpsc::error::TrySendError::Full(_)) => StatusCode::TOO_MANY_REQUESTS.into_response(),
Err(mpsc::error::TrySendError::Closed(_)) => StatusCode::NOT_FOUND.into_response(),
Ok(()) => Ok((StatusCode::ACCEPTED, "Accepted").into_response()),
Err(mpsc::error::TrySendError::Full(_)) => Err(Error::SessionBusy),
Err(mpsc::error::TrySendError::Closed(_)) => Err(Error::UnknownSession),
}
}

3
litellm-rust/crates/gateway-mcp/src/lib.rs
View file

@ -3,6 +3,7 @@ mod error;
mod legacy;
mod native;
mod relay;
mod request;
mod routes;
mod server;
mod sessions;
@ -13,7 +14,7 @@ use http::request::Parts;
use rmcp::{RoleServer, model::*, service::RequestContext};
pub use configured::ConfiguredGateway;
pub use error::{ConnectError, Error};
pub use error::{ConfigFault, ConnectError, Error};
pub use native::{Limits, NativeGateway, Registry, Server, ServerResolver};
pub use relay::RelayClient;
pub use rmcp;

8
litellm-rust/crates/gateway-mcp/src/native.rs
View file

@ -321,9 +321,7 @@ impl NativeGateway {
return Err(Error::Forbidden);
}
let [server] = servers.as_ref() else {
return Err(Error::InvalidRequest(
"read_resource requires exactly one allowed MCP server".into(),
));
return Err(Error::AmbiguousResourceServer);
};
self.send(server, ReadResourceRequest::new(request).into(), &context)
.await
@ -363,9 +361,7 @@ impl Operations for NativeGateway {
fn reject_cursor(params: &Option<PaginatedRequestParams>) -> Result<(), Error> {
match params.as_ref().and_then(|params| params.cursor.as_ref()) {
Some(_) => Err(Error::InvalidRequest(
"Gateway listings do not accept an upstream cursor".into(),
)),
Some(_) => Err(Error::UpstreamCursor),
None => Ok(()),
}
}

4
litellm-rust/crates/gateway-mcp/src/native/catalog.rs
View file

@ -28,9 +28,7 @@ impl NativeGateway {
let state = match next {
None => None,
Some(cursor) if seen.contains(&cursor) || seen.len() >= 1000 => {
return Err(Error::InvalidRequest(
"MCP upstream returned invalid pagination".into(),
));
return Err(Error::InvalidPagination);
}
Some(cursor) => Some((
Some(cursor.clone()),

18
litellm-rust/crates/gateway-mcp/src/native/registry.rs
View file

@ -1,5 +1,5 @@
use super::NativeGateway;
use crate::{Context, Error, GatewayFuture, ServerInfo};
use crate::{ConnectError, Context, Error, GatewayFuture, ServerInfo};
use rmcp::{Peer, RoleClient};
use std::sync::Arc;
@ -39,14 +39,12 @@ pub trait ServerResolver: Send + Sync + 'static {
pub struct Registry(Arc<[Server]>);
impl Registry {
pub fn new(servers: impl Into<Arc<[Server]>>) -> Result<Self, Error> {
pub fn new(servers: impl Into<Arc<[Server]>>) -> Result<Self, ConnectError> {
let servers = servers.into();
for (index, server) in servers.iter().enumerate() {
let identifiers = [&server.info.server_id, &server.info.server_name];
if identifiers.iter().any(|value| value.is_empty()) || server.prefix().is_empty() {
return Err(Error::Configuration(
"MCP server identifiers must not be empty".into(),
));
return Err(ConnectError::EmptyIdentifier);
}
if servers[..index].iter().any(|previous| {
previous.prefix() == server.prefix()
@ -57,9 +55,7 @@ impl Registry {
.as_deref()
.is_some_and(|alias| previous.matches(alias))
}) {
return Err(Error::Configuration(
"MCP server identifiers and prefixes must be unique".into(),
));
return Err(ConnectError::DuplicateIdentifier);
}
}
Ok(Self(servers))
@ -79,11 +75,7 @@ impl NativeGateway {
.parts
.headers
.get("x-mcp-servers")
.map(|value| {
value
.to_str()
.map_err(|_| Error::InvalidRequest("Invalid x-mcp-servers header".into()))
})
.map(|value| value.to_str().map_err(Error::InvalidServersHeader))
.transpose()?;
let scope = context.server.as_deref().or(header);
let Some(scope) = scope else {

44
litellm-rust/crates/gateway-mcp/src/request.rs Normal file
View file

@ -0,0 +1,44 @@
use axum::extract::{FromRequest, FromRequestParts, Request};
use http::request::Parts;
use serde::de::DeserializeOwned;
use crate::Error;
pub(crate) struct Json<T>(pub T);
impl<T: DeserializeOwned, S: Send + Sync> FromRequest<S> for Json<T> {
type Rejection = Error;
async fn from_request(request: Request, state: &S) -> Result<Self, Self::Rejection> {
axum::Json::<T>::from_request(request, state)
.await
.map(|axum::Json(value)| Self(value))
.map_err(Error::Json)
}
}
pub(crate) struct Query<T>(pub T);
impl<T: DeserializeOwned, S: Send + Sync> FromRequestParts<S> for Query<T> {
type Rejection = Error;
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
axum::extract::Query::<T>::from_request_parts(parts, state)
.await
.map(|axum::extract::Query(value)| Self(value))
.map_err(Error::Query)
}
}
pub(crate) struct Path<T>(pub T);
impl<T: DeserializeOwned + Send, S: Send + Sync> FromRequestParts<S> for Path<T> {
type Rejection = Error;
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
axum::extract::Path::<T>::from_request_parts(parts, state)
.await
.map(|axum::extract::Path(value)| Self(value))
.map_err(Error::Path)
}
}

14
litellm-rust/crates/gateway-mcp/src/routes.rs
View file

@ -2,7 +2,7 @@ use std::sync::Arc;
use axum::{
Json, Router,
extract::{Path, Query, Request, State},
extract::{Request, State},
http::request::Parts,
middleware::{self, Next},
response::{IntoResponse, Response},
@ -12,7 +12,11 @@ use rmcp::{model::*, transport::streamable_http_server::StreamableHttpServerConf
use serde::{Deserialize, Serialize};
use tokio_util::sync::CancellationToken;
use crate::{Context, Error, McpServer, Operation, Operations, ToolCatalog, server::ServerScope};
use crate::{
Context, Error, McpServer, Operation, Operations, ToolCatalog,
request::{self, Path, Query},
server::ServerScope,
};
pub struct HttpConfig {
pub allowed_hosts: Vec<String>,
@ -132,16 +136,16 @@ struct ToolCall {
async fn rest_call(
State(operations): State<Arc<dyn Operations>>,
parts: Parts,
Json(call): Json<ToolCall>,
request::Json(call): request::Json<ToolCall>,
) -> Result<impl IntoResponse, Error> {
let server = call
.server_id
.filter(|id| !id.is_empty())
.ok_or_else(|| Error::InvalidRequest("server_id is required in request body".into()))?;
.ok_or(Error::MissingField("server_id"))?;
let name = call
.name
.filter(|name| !name.is_empty())
.ok_or_else(|| Error::InvalidRequest("name is required in request body".into()))?;
.ok_or(Error::MissingField("name"))?;
let params =
CallToolRequestParams::new(name).with_arguments(call.arguments.unwrap_or_default());
let context = Context {

12
litellm-rust/crates/gateway-mcp/src/sessions.rs
View file

@ -4,7 +4,7 @@ use axum::{
Router,
body::{Body, to_bytes},
extract::{Request, State},
http::{Method, StatusCode},
http::Method,
response::IntoResponse,
routing::any,
};
@ -19,7 +19,7 @@ use rmcp::{
use sha2::{Digest, Sha256};
use tower::ServiceExt;
use crate::{McpServer, server::ServerScope};
use crate::{Error, McpServer, server::ServerScope};
#[derive(Clone)]
pub struct SessionOwner(pub String);
@ -65,7 +65,7 @@ pub(crate) fn transport(server: McpServer, config: StreamableHttpServerConfig) -
async fn dispatch(
State(transport): State<Transport>,
request: Request,
) -> Result<impl IntoResponse, StatusCode> {
) -> Result<impl IntoResponse, Error> {
let owner = fingerprint(&request);
let session = request
.headers()
@ -75,8 +75,8 @@ async fn dispatch(
if let Some(session) = &session {
match transport.owners.get(session).await {
Some(expected) if expected == owner => (),
Some(_) => return Err(StatusCode::FORBIDDEN),
None => return Err(StatusCode::NOT_FOUND),
Some(_) => return Err(Error::ForeignSession),
None => return Err(Error::UnknownSession),
}
}
let method = request.method().clone();
@ -84,7 +84,7 @@ async fn dispatch(
let (parts, body) = request.into_parts();
let bytes = to_bytes(body, 4 * 1024 * 1024)
.await
.map_err(|_| StatusCode::PAYLOAD_TOO_LARGE)?;
.map_err(Error::BodyTooLarge)?;
let initialize = matches!(serde_json::from_slice::<ClientJsonRpcMessage>(&bytes),
Ok(ClientJsonRpcMessage::Request(request)) if matches!(request.request, ClientRequest::InitializeRequest(_)));
(Request::from_parts(parts, Body::from(bytes)), initialize)

115
litellm-rust/crates/gateway-mcp/tests/gateway/http.rs
View file

@ -171,43 +171,44 @@ async fn resolver_receives_each_requests_identity(
Duration::from_secs(2),
));
let app = router(operations, HttpConfig::default());
for token in ["allowed", "denied", "allowed"] {
let response = app
.clone()
.oneshot(
Request::builder()
.method(method)
.uri(path)
.header("host", "localhost")
.header("authorization", format!("Bearer {token}"))
.header("content-type", "application/json")
.header("accept", "application/json, text/event-stream")
.header("mcp-protocol-version", "2025-11-25")
.body(Body::from(
json!({"jsonrpc":"2.0","id":1,"method":"tools/list"}).to_string(),
))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = String::from_utf8(
to_bytes(response.into_body(), 65536)
.await
.unwrap()
.to_vec(),
let (status, body) = list_as(&app, method, path, "allowed").await;
assert_eq!(status, StatusCode::OK);
assert!(body.contains("echo-value"), "{body}");
let (status, body) = list_as(&app, method, path, "denied").await;
assert_eq!(status == StatusCode::FORBIDDEN, method == "GET");
assert!(body.contains("not allowed"), "{body}");
let (status, body) = list_as(&app, method, path, "allowed").await;
assert_eq!(status, StatusCode::OK);
assert!(body.contains("echo-value"), "{body}");
}
async fn list_as(
app: &axum::Router,
method: &str,
path: &str,
token: &str,
) -> (StatusCode, String) {
let response = app
.clone()
.oneshot(
Request::builder()
.method(method)
.uri(path)
.header("host", "localhost")
.header("authorization", format!("Bearer {token}"))
.header("content-type", "application/json")
.header("accept", "application/json, text/event-stream")
.header("mcp-protocol-version", "2025-11-25")
.body(Body::from(
json!({"jsonrpc":"2.0","id":1,"method":"tools/list"}).to_string(),
))
.unwrap(),
)
.await
.unwrap();
if token == "denied" {
if method == "GET" {
assert_eq!(status, StatusCode::FORBIDDEN);
}
assert!(body.contains("not allowed"), "{body}");
} else {
assert_eq!(status, StatusCode::OK);
assert!(body.contains("echo-value"), "{body}");
}
}
let status = response.status();
let body = to_bytes(response.into_body(), 65536).await.unwrap();
(status, String::from_utf8(body.to_vec()).unwrap())
}
fn rpc_request(path: &str, token: &str, session: Option<&str>, body: Value) -> Request<Body> {
@ -352,10 +353,14 @@ async fn header_scope_cannot_broaden_path_scope(
}
#[rstest]
#[case::missing_server(json!({"name":"echo-value"}))]
#[case::missing_name(json!({"server_id":"alpha"}))]
#[case::missing_server(json!({"name":"echo-value"}), "server_id")]
#[case::missing_name(json!({"server_id":"alpha"}), "name")]
#[tokio::test]
async fn rest_requires_server_and_tool(#[future] harness: Harness, #[case] body: Value) {
async fn rest_requires_server_and_tool(
#[future] harness: Harness,
#[case] body: Value,
#[case] field: &str,
) {
let harness = harness.await;
let response = harness
.app
@ -368,6 +373,40 @@ async fn rest_requires_server_and_tool(#[future] harness: Harness, #[case] body:
.await
.unwrap();
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let body: Value =
serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap();
assert!(body["detail"].as_str().unwrap().contains(field), "{body}");
}
#[rstest]
#[case::malformed_call("/mcp-rest/tools/call", "POST", "{")]
#[case::unknown_list_parameter("/mcp-rest/tools/list?unknown=1", "GET", "")]
#[case::legacy_message_without_session("/mcp/sse/messages", "POST", "{}")]
#[tokio::test]
async fn rejected_requests_use_the_detail_envelope(
#[future] harness: Harness,
#[case] path: &str,
#[case] method: &str,
#[case] body: &'static str,
) {
let harness = harness.await;
let response = harness
.app
.oneshot(
Request::builder()
.method(method)
.uri(path)
.header("host", "localhost")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert!(response.status().is_client_error(), "{}", response.status());
let body: Value =
serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap();
assert!(body["detail"].is_string(), "{body}");
}
#[rstest]

11
litellm-rust/crates/gateway-mcp/tests/gateway/native.rs
View file

@ -148,10 +148,13 @@ async fn routes_prompts_and_resources_without_rewriting_uris(#[future] harness:
"fixture://{file}"
);
let read = Operation::ReadResource(ReadResourceRequestParams::new("fixture://document"));
assert!(matches!(
harness.gateway.execute(read.clone(), context(None)).await,
Err(Error::InvalidRequest(_))
));
let ambiguous = harness
.gateway
.execute(read.clone(), context(None))
.await
.err()
.expect("read without a server scope");
assert_eq!(ambiguous.status(), axum::http::StatusCode::BAD_REQUEST);
let result = harness
.gateway
.execute(read, context(Some("alpha")))

2
litellm-rust/crates/gateway-ui/Cargo.toml
View file

@ -11,6 +11,7 @@ axum-login.workspace = true
base64.workspace = true
governor = { version = "0.10.4", default-features = false, features = ["std"] }
jsonwebtoken = { workspace = true, features = ["rust_crypto"] }
litellm-error.workspace = true
litellm-gateway-auth.workspace = true
rand.workspace = true
serde.workspace = true
@ -22,6 +23,7 @@ tower-http = { version = "0.6.11", features = ["fs", "set-header"] }
tower-sessions.workspace = true
[dev-dependencies]
futures-util.workspace = true
serde_json.workspace = true
tower = { version = "0.5", features = ["util"] }
litellm-auth-types.workspace = true

16
litellm-rust/crates/gateway-ui/src/dashboard.rs
View file

@ -2,7 +2,8 @@ use std::path::Path;
use axum::{
Router,
extract::{Query, Request},
extract::{FromRequestParts, Query, Request},
http::request::Parts,
routing::get,
};
use serde::{Deserialize, Serialize};
@ -34,6 +35,17 @@ struct LogoQuery {
variant: Option<Variant>,
}
impl<S: Send + Sync> FromRequestParts<S> for LogoQuery {
type Rejection = crate::Error;
async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
Query::<Self>::from_request_parts(parts, state)
.await
.map(|Query(query)| query)
.map_err(crate::Error::Query)
}
}
fn logo_file(query: &LogoQuery) -> &'static str {
match (query.variant, query.theme) {
(Some(Variant::Monogram), Some(Theme::Dark)) => "assets/logos/litellm_monogram_dark.svg",
@ -59,7 +71,7 @@ pub fn dashboard_assets(directory: impl AsRef<Path>) -> Router {
)
.route(
"/get_image",
get(move |Query(query): Query<LogoQuery>, request: Request| {
get(move |query: LogoQuery, request: Request| {
ServeFile::new(logos.join(logo_file(&query))).oneshot(request)
}),
)

93
litellm-rust/crates/gateway-ui/src/error.rs
View file

@ -1,9 +1,9 @@
use axum::{
Json,
extract::rejection::{JsonRejection, QueryRejection},
http::{StatusCode, header},
response::{IntoResponse, Response},
};
use litellm_gateway_auth::UiAuthError;
#[derive(serde::Serialize)]
struct ErrorBody {
@ -21,39 +21,82 @@ pub enum Error {
InvalidCredentials,
#[error("Too many login attempts, try again in one minute")]
RateLimited,
#[error("UI authentication unavailable")]
Auth(#[from] axum_login::Error<litellm_gateway_auth::UiBackend>),
#[error("invalid login request")]
Body(#[source] JsonRejection),
#[error("invalid query string")]
Query(#[source] QueryRejection),
#[error("UI login unavailable")]
Login(#[from] axum_login::Error<litellm_gateway_auth::UiBackend>),
#[error("UI session unavailable")]
Session(#[from] tower_sessions::session::Error),
#[error("UI token unavailable")]
Token(#[from] jsonwebtoken::errors::Error),
#[error("Invalid UI session")]
Unauthorized(#[from] UiAuthError),
#[error(transparent)]
Auth(#[from] litellm_gateway_auth::Error),
}
impl Error {
pub fn status(&self) -> StatusCode {
match self {
Self::InvalidCredentials => StatusCode::UNAUTHORIZED,
Self::RateLimited => StatusCode::TOO_MANY_REQUESTS,
Self::Body(rejection) => rejection.status(),
Self::Query(rejection) => rejection.status(),
Self::Auth(error) => error.status(),
Self::Login(_) | Self::Session(_) | Self::Token(_) => StatusCode::INTERNAL_SERVER_ERROR,
}
}
}
impl IntoResponse for Error {
fn into_response(self) -> Response {
let status = match &self {
Self::InvalidCredentials | Self::Unauthorized(UiAuthError::Unauthorized) => {
StatusCode::UNAUTHORIZED
let status = self.status();
let body = Json(ErrorBody {
error: ErrorMessage {
message: litellm_error::Report::new(&self).to_string(),
},
});
match status {
StatusCode::TOO_MANY_REQUESTS => {
(status, [(header::RETRY_AFTER, "60")], body).into_response()
}
Self::RateLimited => StatusCode::TOO_MANY_REQUESTS,
_ => StatusCode::INTERNAL_SERVER_ERROR,
};
let response = (
status,
Json(ErrorBody {
error: ErrorMessage {
message: self.to_string(),
},
}),
)
.into_response();
if matches!(self, Self::RateLimited) {
return ([(header::RETRY_AFTER, "60")], response).into_response();
_ => (status, body).into_response(),
}
response
}
}
#[cfg(test)]
mod tests {
use litellm_gateway_auth::Error as AuthError;
use rstest::rstest;
use super::*;
#[rstest]
#[case::rejected_session(Error::Auth(AuthError::InvalidSession), StatusCode::UNAUTHORIZED)]
#[case::session_layer_missing(
Error::Auth(AuthError::MissingSessionLayer),
StatusCode::INTERNAL_SERVER_ERROR
)]
#[case::operation_outside_session_scope(
Error::Auth(AuthError::Forbidden),
StatusCode::FORBIDDEN
)]
#[case::rate_limited(Error::RateLimited, StatusCode::TOO_MANY_REQUESTS)]
fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) {
assert_eq!(error.status(), status);
}
#[rstest]
#[case::rate_limited(Error::RateLimited, true)]
#[case::invalid_credentials(Error::InvalidCredentials, false)]
fn only_rate_limits_ask_to_retry_later(#[case] error: Error, #[case] retry_after: bool) {
assert_eq!(
error
.into_response()
.headers()
.contains_key(header::RETRY_AFTER),
retry_after
);
}
}

35
litellm-rust/crates/gateway-ui/src/session.rs
View file

@ -4,14 +4,18 @@ use serde::Serialize;
use std::sync::Arc;
use time::{Duration, OffsetDateTime};
use axum::{Json, extract::State as ExtractState, response::IntoResponse};
use axum::{
Json,
extract::{FromRequest, Request, State as ExtractState},
response::IntoResponse,
};
use governor::DefaultDirectRateLimiter;
use tower_cookies::{Cookie, Cookies};
use tower_sessions::{Expiry, Session, cookie::SameSite};
use crate::Error;
use litellm_gateway_auth::{
AccessRequest, UI_CSRF_KEY, UiAction, UiAuthError, UiAuthSession, UiCredentials, UiSession,
AccessRequest, UI_CSRF_KEY, UiAction, UiAuthSession, UiCredentials, UiSession,
};
pub struct State {
@ -60,12 +64,25 @@ struct Discovery {
is_control_plane: bool,
}
pub struct Credentials(UiCredentials);
impl<S: Send + Sync> FromRequest<S> for Credentials {
type Rejection = Error;
async fn from_request(request: Request, state: &S) -> Result<Self, Self::Rejection> {
Json::<UiCredentials>::from_request(request, state)
.await
.map(|Json(credentials)| Self(credentials))
.map_err(Error::Body)
}
}
pub async fn login(
ExtractState(state): ExtractState<Arc<State>>,
mut auth: UiAuthSession,
session: Session,
cookies: Cookies,
Json(credentials): Json<UiCredentials>,
Credentials(credentials): Credentials,
) -> Result<impl IntoResponse, Error> {
state.login_limit.check().map_err(|_| Error::RateLimited)?;
let user = auth
@ -107,13 +124,14 @@ pub async fn login(
}))
}
pub async fn info(session: Result<UiSession, UiAuthError>) -> Result<impl IntoResponse, Error> {
pub async fn info(
session: Result<UiSession, litellm_gateway_auth::Error>,
) -> Result<impl IntoResponse, Error> {
let session = session?;
session
.identity
.authorize(AccessRequest::Ui(UiAction::SessionInfo))
.await
.map_err(|_| UiAuthError::Unauthorized)?;
.await?;
Ok(Json(SessionInfo {
user_id: session.user.username,
user_role: "proxy_admin",
@ -121,7 +139,7 @@ pub async fn info(session: Result<UiSession, UiAuthError>) -> Result<impl IntoRe
}
pub async fn logout(
validated: Result<UiSession, UiAuthError>,
validated: Result<UiSession, litellm_gateway_auth::Error>,
mut auth: UiAuthSession,
cookies: Cookies,
) -> Result<impl IntoResponse, Error> {
@ -129,8 +147,7 @@ pub async fn logout(
validated
.identity
.authorize(AccessRequest::Ui(UiAction::Logout))
.await
.map_err(|_| UiAuthError::Unauthorized)?;
.await?;
auth.logout().await?;
cookies.add(
Cookie::build(("token", ""))

18
litellm-rust/crates/gateway-ui/tests/assets.rs
View file

@ -200,6 +200,24 @@ async fn logo_discovery_points_to_served_image(dashboard: App) {
);
}
#[rstest]
#[tokio::test]
async fn unknown_logo_theme_is_rejected_in_the_ui_envelope(dashboard: App) {
let response = dashboard
.router
.oneshot(
Request::get("/get_image?theme=purple")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let body: Value =
serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap();
assert!(body["error"]["message"].as_str().unwrap().contains("theme"));
}
#[rstest]
#[case::logo("/get_image")]
#[case::logo_dark("/get_image?theme=dark")]

53
litellm-rust/crates/gateway-ui/tests/routes.rs
View file

@ -330,24 +330,45 @@ async fn reauthentication_rotates_and_revokes_the_previous_session(app: App) {
);
}
async fn failed_login(app: &App, username: &str) -> Response {
app.router
.clone()
.oneshot(
Request::post("/v2/login")
.header(header::CONTENT_TYPE, "application/json")
.body(Body::from(
json!({"username": username, "password": "wrong-password"}).to_string(),
))
.unwrap(),
)
.await
.unwrap()
}
#[rstest]
#[tokio::test]
async fn failed_logins_are_generic_and_rate_limited(app: App) {
for attempt in 0..11 {
let response = app.router.clone().oneshot(Request::post("/v2/login")
.header(header::CONTENT_TYPE, "application/json")
.body(Body::from(json!({"username": format!("unknown-{attempt}"), "password": "wrong-password"}).to_string())).unwrap()).await.unwrap();
if attempt < 10 {
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(
body(response).await["error"]["message"],
"Invalid username or password"
);
} else {
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
assert!(response.headers().contains_key(header::RETRY_AFTER));
}
}
let first = failed_login(&app, "unknown-0").await;
assert_eq!(first.status(), StatusCode::UNAUTHORIZED);
assert_eq!(
body(first).await["error"]["message"],
"Invalid username or password"
);
let usernames: Vec<String> = (1..10)
.map(|attempt| format!("unknown-{attempt}"))
.collect();
let burst = futures_util::future::join_all(
usernames
.iter()
.map(|username| failed_login(&app, username)),
)
.await;
let statuses: Vec<StatusCode> = burst.iter().map(Response::status).collect();
assert_eq!(statuses, vec![StatusCode::UNAUTHORIZED; 9]);
let limited = failed_login(&app, "unknown-10").await;
assert_eq!(limited.status(), StatusCode::TOO_MANY_REQUESTS);
assert!(limited.headers().contains_key(header::RETRY_AFTER));
assert!(body(limited).await["error"]["message"].is_string());
}
#[rstest]
@ -365,6 +386,7 @@ async fn login_rejects_browser_form_posts(app: App) {
.unwrap();
assert_eq!(response.status(), StatusCode::UNSUPPORTED_MEDIA_TYPE);
assert!(!response.headers().contains_key(header::SET_COOKIE));
assert!(body(response).await["error"]["message"].is_string());
}
#[rstest]
@ -468,6 +490,7 @@ async fn invalid_login_payloads_do_not_create_sessions(
.unwrap();
assert_eq!(response.status(), status);
assert!(!response.headers().contains_key(header::SET_COOKIE));
assert!(body(response).await["error"]["message"].is_string());
}
#[rstest]

1
litellm-rust/crates/gateway/Cargo.toml
View file

@ -19,6 +19,7 @@ litellm-gateway-auth.workspace = true
litellm-gateway-ui.workspace = true
litellm-auth-types.workspace = true
litellm-config.workspace = true
litellm-error.workspace = true
litellm-http.workspace = true
litellm-llms.workspace = true
litellm-secrets.workspace = true

8
litellm-rust/crates/gateway/src/auth.rs
View file

@ -64,10 +64,8 @@ impl OperationAuthorizer for McpPolicy {
}
fn mcp_error(error: Error) -> litellm_gateway_mcp::Error {
match error {
Error::InvalidToken | Error::Expired | Error::Forbidden => {
litellm_gateway_mcp::Error::Forbidden
}
_ => litellm_gateway_mcp::Error::Configuration("MCP authorization unavailable".into()),
if error.is_denied() {
return litellm_gateway_mcp::Error::Forbidden;
}
litellm_gateway_mcp::Error::AuthorizationUnavailable(litellm_error::Source::new(error))
}

4
litellm-rust/crates/gateway/src/error.rs
View file

@ -6,8 +6,10 @@ pub enum Error {
Environment,
#[error("MCP host does not yet support configured setting {0}")]
McpSetting(String),
#[error("MCP host does not yet support guardrails or policies")]
McpGuardrails,
#[error(transparent)]
Mcp(#[from] litellm_gateway_mcp::ConnectError),
#[error("MCP requires a configured master key")]
#[error("MCP requires working gateway authentication")]
Auth(#[from] litellm_gateway_auth::Error),
}

2
litellm-rust/crates/gateway/src/lib.rs
View file

@ -96,7 +96,7 @@ pub async fn build_mcp(
|| !config.policies.is_empty()
|| !config.policy_attachments.is_empty()
{
return Err(Error::McpSetting("guardrails or policies".into()));
return Err(Error::McpGuardrails);
}
Auth::from_config(config, secrets.clone())
.validate()

119
litellm-rust/crates/gateway/tests/mcp.rs
View file

@ -95,11 +95,29 @@ async fn upstream() -> Fixture {
}
}
#[rstest]
#[tokio::test]
async fn serves_configured_upstream_with_separate_admission_and_outbound_credentials(
#[future] upstream: Fixture,
) {
struct Served {
app: Router,
mcp: litellm_gateway_mcp::ConfiguredGateway,
shutdown: CancellationToken,
upstream: Fixture,
}
impl Served {
async fn close(self) {
self.shutdown.cancel();
tokio::time::timeout(Duration::from_secs(5), self.mcp.close())
.await
.unwrap();
self.upstream.shutdown.cancel();
tokio::time::timeout(Duration::from_secs(5), self.upstream.task)
.await
.unwrap()
.unwrap();
}
}
#[fixture]
async fn served(#[future] upstream: Fixture) -> Served {
let upstream = upstream.await;
let config = Config::from_yaml(&format!(
r#"
@ -146,30 +164,53 @@ mcp_servers:
},
);
let app = litellm_gateway::router(inference, &config, None, Some(mcp_router));
for (path, method) in [
("/mcp", "POST"),
("/mcp", "GET"),
("/mcp", "DELETE"),
("/mcp/sse", "GET"),
("/mcp/sse/messages", "POST"),
("/mcp-rest/tools/list", "GET"),
("/mcp/enabled", "GET"),
] {
let response = app
.clone()
.oneshot(
Request::builder()
.method(method)
.uri(path)
.header("host", "localhost")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), 401, "{method} {path}");
Served {
app,
mcp,
shutdown,
upstream,
}
}
#[rstest]
#[case::streamable_post("/mcp", "POST")]
#[case::streamable_get("/mcp", "GET")]
#[case::streamable_delete("/mcp", "DELETE")]
#[case::legacy_sse("/mcp/sse", "GET")]
#[case::legacy_messages("/mcp/sse/messages", "POST")]
#[case::rest_list("/mcp-rest/tools/list", "GET")]
#[case::enabled("/mcp/enabled", "GET")]
#[tokio::test]
async fn every_mcp_route_requires_admission_credentials(
#[future] served: Served,
#[case] path: &str,
#[case] method: &str,
) {
let served = served.await;
let response = served
.app
.clone()
.oneshot(
Request::builder()
.method(method)
.uri(path)
.header("host", "localhost")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), 401);
served.close().await;
}
#[rstest]
#[tokio::test]
async fn serves_configured_upstream_with_separate_admission_and_outbound_credentials(
#[future] served: Served,
) {
let served = served.await;
let app = served.app.clone();
let list = app
.clone()
.oneshot(
@ -211,15 +252,7 @@ mcp_servers:
.unwrap();
assert_eq!(blocked.status(), 403);
shutdown.cancel();
tokio::time::timeout(Duration::from_secs(5), mcp.close())
.await
.unwrap();
upstream.shutdown.cancel();
tokio::time::timeout(Duration::from_secs(5), upstream.task)
.await
.unwrap()
.unwrap();
served.close().await;
}
#[rstest]
@ -252,9 +285,10 @@ async fn rejects_unsupported_or_invalid_settings_before_serving(
&inference.http,
)
.await;
let error = result.err().expect("startup must fail").to_string();
assert!(error.contains(expected), "{error}");
assert!(!error.contains("secret-value"));
let error = result.err().expect("startup must fail");
let report = litellm_error::Report::new(&error).to_string();
assert!(report.contains(expected), "{report}");
assert!(!report.contains("secret-value"));
}
#[rstest]
@ -341,5 +375,8 @@ async fn refuses_unimplemented_global_mcp_policy(#[case] settings: &str) {
.await
.err()
.unwrap();
assert!(matches!(error, litellm_gateway::Error::McpSetting(_)));
assert!(matches!(
error,
litellm_gateway::Error::McpSetting(_) | litellm_gateway::Error::McpGuardrails
));
}

1
litellm-rust/crates/host/Cargo.toml
View file

@ -10,6 +10,7 @@ futures-util.workspace = true
litellm-auth.workspace = true
litellm-coroutine.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["sync"] }
[dev-dependencies]

6
litellm-rust/crates/host/src/machine/coroutine.rs
View file

@ -10,12 +10,14 @@ use super::{
use crate::protocol::{HostRequest, Protocol};
/// The machine's own failures, distinct from anything the provider call reports.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum MachineFault {
/// The host dropped an op's reply unanswered, or went away while the call waited.
#[error("the host abandoned the call")]
Abandoned,
/// The host resumed the call out of turn.
Protocol(ResumeError),
#[error("the host resumed the call out of turn")]
Protocol(#[source] ResumeError),
}
pub type ExecuteFuture<R, C = <R as Protocol>::Response> =

7
litellm-rust/crates/inference-chat/src/common_utils.rs
View file

@ -56,9 +56,10 @@ pub(super) fn resolve_request(
))
})?;
if messages.is_empty() {
return Err(Error::InvalidRequest(
"chat completions requires at least one message".into(),
));
return Err(Error::InvalidRequest(litellm_llms::ErrorDetail::Expected {
subject: "chat completions messages",
expected: "non-empty",
}));
}
if let Some(reason) = config.unsupported_reason(&messages, &call.optional_params) {
return Err(Error::Unsupported(reason.0));

10
litellm-rust/crates/inference-chat/src/handler.rs
View file

@ -147,7 +147,10 @@ pub(super) fn as_response_error(err: Error) -> Error {
match err {
already @ Error::InvalidResponse(_) => already,
upstream if upstream.upstream_response().is_some() => upstream,
other => Error::InvalidResponse(other.to_string().into()),
other => Error::InvalidResponse(litellm_llms::ErrorDetail::invalid(
"chat completions response",
other,
)),
}
}
@ -315,7 +318,10 @@ mod tests {
for original in [
Error::MissingField("usage"),
Error::Unsupported("non-text response content block"),
Error::InvalidRequest("whatever".to_string().into()),
Error::InvalidRequest(litellm_llms::ErrorDetail::Expected {
subject: "whatever",
expected: "anything else",
}),
Error::Auth(litellm_auth::Error::InvalidHeader),
] {
let label = format!("{original:?}");

4
litellm-rust/crates/inference-chat/tests/chat_completions.rs
View file

@ -352,7 +352,7 @@ async fn a_post_call_hook_failure_never_looks_safe_to_retry(
Ok(wire)
}
async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> {
Err(Error::InvalidRequest("callback rejected".into()))
Err(Error::MissingField("callback rejected"))
}
}
let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await;
@ -371,7 +371,7 @@ async fn a_post_call_hook_failure_never_looks_safe_to_retry(
let Error::PostCallHook(source) = error else {
panic!("expected retained callback error")
};
assert_eq!(*source, Error::InvalidRequest("callback rejected".into()));
assert_eq!(*source, Error::MissingField("callback rejected"));
assert_eq!(received(&upstream).await.len(), 1);
}

12
litellm-rust/crates/inference-chat/tests/request.rs
View file

@ -246,12 +246,18 @@ async fn rejects_a_model_with_no_resolvable_provider() {
}
#[rstest]
#[case::empty(json!([]), Some("chat completions requires at least one message"))]
#[case::empty(
json!([]),
Some(litellm_llms::ErrorDetail::Expected {
subject: "chat completions messages",
expected: "non-empty",
})
)]
#[case::malformed(json!("not a list"), None)]
#[tokio::test]
async fn rejects_an_empty_or_malformed_message_list(
#[case] messages: Value,
#[case] expected: Option<&str>,
#[case] expected: Option<litellm_llms::ErrorDetail>,
) {
let error = complete(request(
"anthropic/claude-sonnet-4-5",
@ -262,7 +268,7 @@ async fn rejects_an_empty_or_malformed_message_list(
.await
.unwrap_err();
if let Some(expected) = expected {
assert_eq!(error, Error::InvalidRequest(expected.to_string().into()));
assert_eq!(error, Error::InvalidRequest(expected));
} else {
assert!(matches!(error, Error::InvalidRequest(_)));
}

1
litellm-rust/crates/inference-messages/Cargo.toml
View file

@ -26,6 +26,7 @@ serde_json = { workspace = true, features = ["preserve_order"] }
tracing.workspace = true
[dev-dependencies]
litellm-error.workspace = true
litellm-framer = { workspace = true, features = ["sse"] }
litellm-cache-memory.workspace = true
litellm-host-native.workspace = true

4
litellm-rust/crates/inference-messages/tests/messages/host.rs
View file

@ -322,14 +322,14 @@ async fn a_before_send_failure_never_sends(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let host = RecordingHost::new(
authenticated(call, upstream.uri()),
Box::new(|_| Err(Error::InvalidRequest("vetoed by the host".into()))),
Box::new(|_| Err(Error::MissingField("vetoed by the host"))),
);
let error = run_through(&host)
.await
.expect_err("the host failure fails the call");
assert_eq!(error, Error::InvalidRequest("vetoed by the host".into()));
assert_eq!(error, Error::MissingField("vetoed by the host"));
assert!(received(&upstream).await.is_empty());
assert!(host.raw_responses().is_empty());
}

2
litellm-rust/crates/inference-messages/tests/messages/main.rs
View file

@ -144,7 +144,7 @@ impl LocalMessagesHost {
.lock()
.unwrap_or_else(|error| error.into_inner())
.take()
.ok_or_else(|| Error::InvalidRequest("messages request was already projected".into()))
.ok_or(Error::MissingField("messages request"))
}
pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, (), Self, ()> {
litellm_host_native::in_process::Host {

7
litellm-rust/crates/inference-messages/tests/messages/request.rs
View file

@ -957,11 +957,10 @@ async fn provider_validation_runs_before_caller_parameter_removal(
))
.await;
let error = result.expect_err("metadata is validated before removal");
let report = litellm_error::Report::new(&error).to_string();
assert!(
error
.to_string()
.contains("metadata.user_id must be a string"),
"{error}"
report.contains("metadata.user_id must be a string"),
"{report}"
);
assert!(received(&upstream).await.is_empty());
}

5
litellm-rust/crates/inference-messages/tests/messages/response.rs
View file

@ -312,8 +312,11 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes
fn a_body_that_does_not_parse_is_an_invalid_request(#[case] raw: Value) {
let error = messages_body(object(raw)).expect_err("the body is rejected");
assert!(matches!(error, Error::InvalidRequest(_)), "{error:?}");
assert!(
matches!(&error, Error::InvalidRequest(message) if message.to_string().starts_with("invalid Anthropic messages request: ")),
litellm_error::Report::new(&error)
.to_string()
.contains("invalid Anthropic messages request"),
"{error:?}"
);
}

4
litellm-rust/crates/inference-ocr/src/types.rs
View file

@ -67,6 +67,7 @@ impl OcrConnectionInputs {
.map(|value| (name.clone(), value.to_string()))
.ok_or_else(|| Error::RequestField {
path: format!("extra_headers.{name}"),
cause: None,
})
})
.collect()
@ -105,6 +106,7 @@ impl LiteLLMOcrRequest {
.filter(|value| *value > 0 && *value <= default_transport.max_response_bytes)
.ok_or_else(|| Error::RequestField {
path: "max_response_bytes".into(),
cause: None,
})
})
.transpose()?
@ -313,7 +315,7 @@ mod tests {
};
assert!(matches!(
error,
Error::RequestField { ref path } if path == "extra_headers.x-a"
Error::RequestField { ref path, .. } if path == "extra_headers.x-a"
));
}
}

9
litellm-rust/crates/inference-ocr/src/wire.rs
View file

@ -73,9 +73,8 @@ pub fn decode_request_input<D: Into<OcrDocumentInput>>(
let timeout = wire
.timeout_seconds
.map(|seconds| {
Duration::try_from_secs_f64(seconds).map_err(|_| Error::RequestField {
path: "timeout_seconds".into(),
})
Duration::try_from_secs_f64(seconds)
.map_err(|error| Error::request_field("timeout_seconds", error))
})
.transpose()?;
LiteLLMOcrRequest::from_inputs(
@ -123,8 +122,8 @@ mod tests {
#[case::non_object(json!([]), "document")]
#[case::missing_type(json!({"document_url":"https://example.com/a.pdf"}), "document")]
#[case::unsupported_type(json!({"type":"text"}), "type")]
#[case::missing_document_url(json!({"type":"document_url"}), "Document URL")]
#[case::missing_image_url(json!({"type":"image_url"}), "Document URL")]
#[case::missing_document_url(json!({"type":"document_url"}), "document URL")]
#[case::missing_image_url(json!({"type":"image_url"}), "document URL")]
fn ocr_contract_malformed_document_is_bad_request(
#[case] document: Value,
#[case] field: &str,

2
litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs
View file

@ -239,7 +239,7 @@ async fn credential_precedence(
true,
json!({"azure_ad_token": "static-token"}),
|_| String::new(),
|error: &Error| matches!(error, Error::MissingAzureAiCredentials),
|error: &Error| error.is_missing_credential() && error.to_string().contains("Azure AI credentials"),
1
)]
#[tokio::test]

16
litellm-rust/crates/inference-ocr/tests/ocr/azure_document_intelligence.rs
View file

@ -5,7 +5,7 @@ use std::{
use litellm_host::lifecycle::CallEvent;
use litellm_host::lifecycle::ExecutionEvent;
use litellm_llms::base_llm::ocr::settings::OcrSettings;
use litellm_llms::base_llm::ocr::{error::PagesFault, settings::OcrSettings};
use rstest::rstest;
@ -64,10 +64,10 @@ async fn pages_features_and_extra_options_map_to_the_analyze_call() {
}
#[rstest]
#[case(json!({"pages": [true]}), Error::Pages("expected only integers or only strings".into()))]
#[case(json!({"pages": [1, "2"]}), Error::Pages("expected only integers or only strings".into()))]
#[case(json!({"pages": [-1]}), Error::Pages("negative page index".into()))]
#[case(json!({"pages": "1&&features=bad"}), Error::Pages("invalid native page range".into()))]
#[case(json!({"pages": [true]}), Error::Pages(PagesFault::Mixed))]
#[case(json!({"pages": [1, "2"]}), Error::Pages(PagesFault::Mixed))]
#[case(json!({"pages": [-1]}), Error::Pages(PagesFault::Negative))]
#[case(json!({"pages": "1&&features=bad"}), Error::Pages(PagesFault::Range))]
#[case(json!({"features": "languages&pages=1"}), Error::Features)]
#[case(json!({"req_format": "azure"}), Error::RequestFormat)]
#[tokio::test]
@ -315,7 +315,11 @@ async fn polling_does_not_follow_redirects() {
.await
.unwrap_err();
assert!(error.to_string().contains("status 302"), "{error}");
assert_eq!(
error.upstream_response().map(|(status, ..)| status),
Some(302),
"{error}"
);
assert_eq!(received(&upstream).await.len(), 2);
}

11
litellm-rust/crates/inference-ocr/tests/ocr/cohere.rs
View file

@ -37,6 +37,15 @@ async fn a_non_image_document_is_rejected_before_sending(
.await
.unwrap_err();
assert!(matches!(error, Error::CohereImageOnly), "{error:?}");
assert!(
matches!(
error,
Error::UnsupportedDocument {
provider: "Cohere Parse",
..
}
),
"{error:?}"
);
assert!(received(&upstream).await.is_empty());
}

7
litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs
View file

@ -38,7 +38,7 @@ fn recording_host(
.with_before_send(move |wire, _| {
interceptions.fetch_add(1, Ordering::SeqCst);
match block {
true => Err(Error::InvalidRequest("blocked".into())),
true => Err(Error::MissingField("blocked")),
false => Ok(wire),
}
})
@ -95,10 +95,7 @@ async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() {
.await
.unwrap_err();
assert!(
matches!(&error, Error::InvalidRequest(message) if message == "blocked"),
"{error:?}"
);
assert!(matches!(error, Error::MissingField("blocked")), "{error:?}");
assert_eq!(interceptions.load(Ordering::SeqCst), 1);
assert_eq!(*events.lock().unwrap(), ["started", "failure"]);
assert!(received(&upstream).await.is_empty());

35
litellm-rust/crates/inference-ocr/tests/ocr/machine.rs
View file

@ -117,10 +117,14 @@ async fn a_hand_driven_machine_performs_the_same_call() {
assert_eq!(outcome.unwrap().pages[0].markdown, "native");
assert_eq!(received(&upstream).await.len(), 1);
assert_eq!(ops, ["BeforeSend", "response"]);
assert!(matches!(
machine.resume().await,
Err(Error::InvalidRequest(_))
));
let Err(error) = machine.resume().await else {
panic!("resuming a finished call is a host protocol fault");
};
assert!(!error.is_request(), "{error:?}");
assert!(
litellm_error::find::<litellm_host::machine::MachineFault>(&error).is_some(),
"{error:?}"
);
}
#[rstest::rstest]
@ -153,8 +157,11 @@ async fn a_path_document_is_read_by_core_without_a_host_operation() {
}
#[rstest]
#[case::failed(HostFailure::Error(Error::InvalidRequest("before_provider_request failed".into())), "before_provider_request failed")]
#[case::cancelled(HostFailure::Cancelled(Error::InvalidRequest("cancelled".into())), "cancelled")]
#[case::failed(
HostFailure::Error(Error::MissingField("before_provider_request failed")),
"before_provider_request failed"
)]
#[case::cancelled(HostFailure::Cancelled(Error::MissingField("cancelled")), "cancelled")]
#[tokio::test]
async fn a_before_send_failure_ends_the_call_without_reaching_transport(
#[case] failure: HostFailure<Error>,
@ -174,7 +181,7 @@ async fn a_before_send_failure_ends_the_call_without_reaching_transport(
.await;
assert!(
matches!(&outcome, Err(Error::InvalidRequest(actual)) if actual == message),
matches!(&outcome, Err(Error::MissingField(actual)) if *actual == message),
"{outcome:?}"
);
assert_eq!(ops, ["BeforeSend"]);
@ -255,17 +262,17 @@ async fn interrupt_drops_provider_captures_before_returning() {
drive_until_notified(&mut machine, &host, &entered).await;
assert!(!dropped.load(Ordering::SeqCst));
let acknowledgement = machine.interrupt(HostFailure::Cancelled(Error::InvalidRequest(
"cancelled".into(),
)));
let acknowledgement =
machine.interrupt(HostFailure::Cancelled(Error::MissingField("cancelled")));
assert!(
dropped.load(Ordering::SeqCst),
"interrupt returned while provider captures were still alive"
);
assert!(
matches!(acknowledgement.await, Err(Error::InvalidRequest(message)) if message == "cancelled")
);
assert!(matches!(
acknowledgement.await,
Err(Error::MissingField("cancelled"))
));
}
#[rstest::rstest]
@ -290,7 +297,7 @@ async fn interrupting_an_in_flight_provider_request_closes_its_connection() {
let mut machine = ocr_route().machine(host.request().unwrap(), None);
drive_until_notified(&mut machine, &host, &received).await;
let cancelled = Error::InvalidRequest("cancelled".into());
let cancelled = Error::MissingField("cancelled");
assert!(
machine

2
litellm-rust/crates/inference-ocr/tests/ocr/main.rs
View file

@ -209,7 +209,7 @@ impl LocalOcrHost {
request,
caller_token: false,
})
.ok_or_else(|| Error::InvalidRequest("OCR request was already projected".into()))
.ok_or(Error::MissingField("OCR request"))
}
pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, Self, Self, ()> {
litellm_host_native::in_process::Host {

25
litellm-rust/crates/inference-ocr/tests/ocr/mistral.rs
View file

@ -92,28 +92,13 @@ async fn an_upstream_error_keeps_its_status_whole_body_and_headers(
.unwrap_err();
assert_eq!(received(&upstream).await.len(), 1);
let Error::Provider {
status,
body,
headers,
} = error
else {
panic!("expected provider error, got {error:?}");
};
let (status, body, headers) = error.upstream_response().expect("the provider answered");
assert_eq!(status, 422);
for (name, value) in [
("retry-after", "17"),
("x-request-id", "request-123"),
("x-future-header", "retained"),
] {
assert!(
headers
.iter()
.any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value),
"{name} missing from {headers:?}"
);
}
assert_eq!(body, expected_body);
let header = |name: &str| litellm_http::request::header_value(headers, name);
assert_eq!(header("retry-after"), Some("17"));
assert_eq!(header("x-request-id"), Some("request-123"));
assert_eq!(header("x-future-header"), Some("retained"));
}
#[rstest]

8
litellm-rust/crates/inference-ocr/tests/ocr/reducto.rs
View file

@ -175,8 +175,11 @@ async fn a_failed_upload_stops_before_parse(#[case] upload: ResponseTemplate) {
}
#[rstest]
#[case::remote_url("https://example.com/a.pdf", Error::ReductoSource)]
#[case::empty_file_id("reducto://", Error::RequestField { path: "document file id".into() })]
#[case::remote_url(
"https://example.com/a.pdf",
Error::UnsupportedDocument { provider: "Reducto", guidance: "" }
)]
#[case::empty_file_id("reducto://", Error::RequestField { path: "document file id".into(), cause: None })]
#[case::data_uri_without_payload("data:application/pdf;base64", Error::InvalidDataUri)]
#[case::invalid_base64("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)]
#[tokio::test]
@ -202,7 +205,6 @@ async fn invalid_document_sources_are_rejected_before_sending(
std::mem::discriminant(&expected)
);
assert_eq!(error.http_status_code(), Some(400));
assert_eq!(error.to_string(), expected.to_string());
}
#[tokio::test]

15
litellm-rust/crates/inference-responses/src/handler.rs
View file

@ -72,7 +72,12 @@ pub(super) async fn execute(
let stream = match wire.body.get("stream") {
None => false,
Some(serde_json::Value::Bool(value)) => *value,
Some(_) => return Err(Error::InvalidRequest("stream must be a boolean".into())),
Some(_) => {
return Err(Error::InvalidRequest(litellm_llms::ErrorDetail::Expected {
subject: "stream",
expected: "a boolean",
}));
}
};
let outbound = litellm_inference::outbound::outbound_request(
Authenticated {
@ -117,8 +122,12 @@ pub(super) async fn execute(
.after_provider_response(raw)
.await
.map_err(Error::post_call)?;
let value = serde_json::from_str(&body)
.map_err(|error| Error::InvalidResponse(error.to_string().into()))?;
let value = serde_json::from_str(&body).map_err(|error| {
Error::InvalidResponse(litellm_llms::ErrorDetail::invalid(
"Responses API response JSON",
error,
))
})?;
config
.transform_response_api_response(value)
.map(ResponsesOutput::Complete)

22
litellm-rust/crates/inference-responses/src/websocket.rs
View file

@ -11,6 +11,8 @@ use tokio_tungstenite::tungstenite::{
http::{HeaderMap, HeaderName, HeaderValue},
};
use litellm_llms::ErrorDetail;
use super::Error;
#[derive(Clone)]
@ -36,10 +38,18 @@ impl ResponsesWebSocketConnection {
.iter()
.map(|(name, value)| {
Ok((
name.parse::<HeaderName>()
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?,
HeaderValue::from_str(value)
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?,
name.parse::<HeaderName>().map_err(|error| {
Error::InvalidRequest(ErrorDetail::invalid(
"WebSocket header name",
error,
))
})?,
HeaderValue::from_str(value).map_err(|error| {
Error::InvalidRequest(ErrorDetail::invalid(
"WebSocket header value",
error,
))
})?,
))
})
.collect::<Result<HeaderMap, Error>>()?;
@ -87,7 +97,9 @@ impl ResponsesWebSocketConnection {
Some(Ok(Message::Text(text))) => Ok(Some(text)),
Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
.map(Some)
.map_err(|error| Error::InvalidResponse(error.to_string().into())),
.map_err(|error| {
Error::InvalidResponse(ErrorDetail::invalid("WebSocket text frame", error))
}),
Some(Ok(Message::Close(_))) | None => Ok(None),
Some(Ok(_)) => Ok(None),
Some(Err(error)) => Err(litellm_http::Error::from(error).into()),

1
litellm-rust/crates/inference/Cargo.toml
View file

@ -29,6 +29,7 @@ time.workspace = true
url.workspace = true
[dev-dependencies]
litellm-error.workspace = true
litellm-cache-memory.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-llms = { workspace = true, features = ["test-support"] }

38
litellm-rust/crates/inference/src/error.rs
View file

@ -20,10 +20,10 @@ pub enum Error {
MissingField(&'static str),
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("invalid request: {0}")]
InvalidRequest(litellm_llms::ErrorDetail),
#[error("invalid response: {0}")]
InvalidResponse(litellm_llms::ErrorDetail),
#[error("invalid request")]
InvalidRequest(#[source] litellm_llms::ErrorDetail),
#[error("invalid response")]
InvalidResponse(#[source] litellm_llms::ErrorDetail),
#[error("unsupported by the rust path: {0}")]
Unsupported(&'static str),
#[error(transparent)]
@ -40,11 +40,7 @@ pub enum Error {
impl From<litellm_host::machine::MachineFault> for Error {
fn from(fault: litellm_host::machine::MachineFault) -> Self {
use litellm_host::machine::MachineFault;
Self::InvalidRequest(match fault {
MachineFault::Abandoned => "host driver was abandoned".into(),
MachineFault::Protocol(message) => format!("host {message}").into(),
})
Self::InvalidRequest(litellm_llms::ErrorDetail::failed("host driver", fault))
}
}
@ -107,6 +103,7 @@ impl From<LlmError> for Error {
LlmError::InvalidResponse(message) => Self::InvalidResponse(message),
LlmError::Unsupported(reason) => Self::Unsupported(reason),
LlmError::Auth(error) => Self::Auth(error),
LlmError::Http(error) => Self::from(error),
}
}
}
@ -151,8 +148,20 @@ mod tests {
.is_request()
);
assert!(Error::Auth(litellm_auth::Error::InvalidHeader).is_request());
assert!(Error::InvalidRequest("top_k".into()).is_request());
assert!(!Error::InvalidResponse("bad json".into()).is_request());
assert!(
Error::InvalidRequest(ErrorDetail::Expected {
subject: "top_k",
expected: "an integer",
})
.is_request()
);
assert!(
!Error::InvalidResponse(ErrorDetail::MissingField {
subject: "response",
field: "usage",
})
.is_request()
);
}
#[test]
fn a_header_the_signer_computes_is_left_to_the_host_on_every_route() {
@ -176,12 +185,9 @@ mod tests {
assert_eq!(error.is_request(), request);
let category = if request { "request" } else { "response" };
assert_eq!(
error.to_string(),
litellm_error::Report::new(&error).to_string(),
format!("invalid {category}: invalid test payload: {source_message}")
);
assert!(
std::error::Error::source(&error).is_none(),
"a cause rendered in the message is not also exposed as source()"
);
assert!(litellm_error::find::<serde_json::Error>(&error).is_some());
}
}

2
litellm-rust/crates/llms/Cargo.toml
View file

@ -11,6 +11,7 @@ test-support = ["litellm-http/test-support"]
[dependencies]
litellm-llms-types.workspace = true
litellm-core-utils.workspace = true
litellm-error.workspace = true
litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] }
litellm-auth-aws.workspace = true
litellm-auth-azure.workspace = true
@ -37,7 +38,6 @@ tokio = { workspace = true, features = ["sync"] }
url.workspace = true
[dev-dependencies]
litellm-error.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
aws-smithy-eventstream = "=0.61.4"
aws-smithy-types = "1.6.1"

45
litellm-rust/crates/llms/src/anthropic/batches/transformation.rs
View file

@ -149,7 +149,12 @@ impl AnthropicBatchesConfig for AnthropicBatchesTransformation {
}
let mut url = batches_base_url(api_base, env_lookup)?;
url.path_segments_mut()
.map_err(|_| Error::InvalidRequest("Anthropic API base cannot be a base URL".into()))?
.map_err(|()| {
Error::InvalidRequest(crate::ErrorDetail::Expected {
subject: "Anthropic API base",
expected: "a base URL",
})
})?
.push(batch_id);
Ok(url.into())
}
@ -217,11 +222,11 @@ impl AnthropicBatchesConfig for AnthropicBatchesTransformation {
.map(|(index, line)| {
let record: BatchResultRecord =
serde_json::from_str(line.trim()).map_err(|error| {
Error::InvalidResponse(crate::ErrorDetail::InvalidLine {
subject: "Anthropic batch result",
line: index + 1,
cause: crate::ErrorSource::new(error),
})
Error::InvalidResponse(crate::ErrorDetail::invalid_line(
"Anthropic batch result",
index + 1,
error,
))
})?;
match record.result {
BatchResult::Succeeded { message } => Ok(*message),
@ -314,18 +319,22 @@ mod tests {
assert_eq!(messages[0].id, "msg_1");
}
#[test]
fn reports_malformed_and_unsuccessful_batch_results() {
assert!(matches!(
ANTHROPIC_BATCHES_TRANSFORMATION.transform_batch_results("not-json"),
Err(Error::InvalidResponse(message)) if message.to_string().contains("line 1")
));
assert!(matches!(
ANTHROPIC_BATCHES_TRANSFORMATION.transform_batch_results(
r#"{"result":{"type":"errored","error":{"type":"invalid_request_error"}}}"#
),
Err(Error::InvalidResponse(message)) if message.to_string().contains("request failed")
));
#[rstest::rstest]
#[case::malformed("not-json", "line 1")]
#[case::errored(
r#"{"result":{"type":"errored","error":{"type":"invalid_request_error"}}}"#,
"request failed"
)]
fn reports_malformed_and_unsuccessful_batch_results(#[case] body: &str, #[case] fact: &str) {
let error = ANTHROPIC_BATCHES_TRANSFORMATION
.transform_batch_results(body)
.unwrap_err();
assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}");
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains(fact)
);
}
#[test]

163
litellm-rust/crates/llms/src/anthropic/chat/handler.rs
View file

@ -1,163 +0,0 @@
use std::collections::HashMap;
use litellm_llms_types::formats::{
chat_completions::{
ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
ChatCompletionsUsage,
},
messages::streaming::{
MessagesContentBlock, MessagesContentBlockDelta, MessagesStreamEvent, MessagesStreamUsage,
},
};
use serde_json::Value;
use crate::{
Error,
base_llm::{base_model_iterator::StreamTransformer, chat::streaming::StreamShape},
};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum AnthropicJsonChunkType {
ValidJson,
AccumulatedJson,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum AnthropicContentBlockType {
Text,
ToolUse,
ServerToolUse,
Thinking,
RedactedThinking,
Compaction,
ToolResult(String),
Other(String),
}
#[derive(Clone, Debug, PartialEq)]
pub struct AnthropicContentBlockDeltaEvent {
pub index: u64,
pub delta: MessagesContentBlockDelta,
}
pub struct ModelResponseIterator {
pub content_blocks: Vec<AnthropicContentBlockDeltaEvent>,
pub tool_index: i64,
pub json_mode: bool,
pub speed: Option<String>,
pub tool_name_reverse_map: HashMap<String, String>,
pub response_id: String,
pub served_model: Option<String>,
pub is_response_format_tool: bool,
pub converted_response_format_tool: bool,
pub accumulated_json: String,
pub chunk_type: AnthropicJsonChunkType,
pub current_content_block_type: Option<AnthropicContentBlockType>,
pub web_search_results: Vec<Value>,
pub web_search_calls: HashMap<String, Value>,
pub compaction_blocks: Vec<Value>,
pub reasoning_content_chunks: Vec<String>,
pub server_tool_inputs: HashMap<String, Value>,
pub tool_results: Vec<Value>,
pub current_server_tool_id: Option<String>,
pub container_id: Option<String>,
}
impl ModelResponseIterator {
pub fn new(_shape: StreamShape) -> Self {
todo!()
}
pub fn check_empty_tool_call_args(&self) -> bool {
todo!()
}
pub fn handle_usage(&mut self, _usage: MessagesStreamUsage) -> ChatCompletionsUsage {
todo!()
}
pub fn handle_content_block_delta(
&mut self,
_index: u64,
_delta: MessagesContentBlockDelta,
) -> (
String,
Option<ChatCompletionToolCallChunk>,
Vec<ChatCompletionThinkingBlock>,
Option<Value>,
Option<String>,
) {
todo!()
}
pub fn handle_content_block_start(
&mut self,
_index: u64,
_content_block: MessagesContentBlock,
) -> Result<ChatCompletionChunk, Error> {
todo!()
}
pub fn handle_json_mode_chunk(
&mut self,
_text: String,
_tool_use: Option<ChatCompletionToolCallChunk>,
) -> (String, Option<ChatCompletionToolCallChunk>) {
todo!()
}
pub fn handle_accumulated_json_chunk(
&mut self,
_data: &str,
_is_final: bool,
) -> Result<Option<ChatCompletionChunk>, Error> {
todo!()
}
pub fn handle_redacted_thinking_content(
&mut self,
_content_block: &MessagesContentBlock,
) -> Vec<ChatCompletionThinkingBlock> {
todo!()
}
pub fn web_search_call_snapshot(&self) -> HashMap<String, Value> {
todo!()
}
pub fn complete_web_search_call(&mut self, _result: Value) {
todo!()
}
pub fn build_code_interpreter_results(&self) -> Vec<Value> {
todo!()
}
pub fn handle_message_delta(
&mut self,
_event: MessagesStreamEvent,
) -> (Option<String>, Option<ChatCompletionsUsage>, Option<Value>) {
todo!()
}
pub fn chunk_parser(
&mut self,
_event: MessagesStreamEvent,
) -> Result<ChatCompletionChunk, Error> {
todo!()
}
}
impl StreamTransformer for ModelResponseIterator {
type Input = MessagesStreamEvent;
type Output = ChatCompletionChunk;
type Error = Error;
fn transform(&mut self, _input: Self::Input) -> Result<Vec<Self::Output>, Self::Error> {
todo!()
}
fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error> {
todo!()
}
}

1
litellm-rust/crates/llms/src/anthropic/chat/mod.rs
View file

@ -1,2 +1 @@
pub mod handler;
pub mod transformation;

35
litellm-rust/crates/llms/src/anthropic/chat/transformation.rs
View file

@ -14,23 +14,19 @@ use serde_json::{Map, Value, json};
use crate::{
Error,
anthropic::{
chat::handler::ModelResponseIterator,
common_utils::{
API_KEY_PLACEMENT, complete_anthropic_url, forwarded_oauth_bearer,
resolve_anthropic_api_key,
},
anthropic::common_utils::{
API_KEY_PLACEMENT, complete_anthropic_url, forwarded_oauth_bearer,
resolve_anthropic_api_key,
},
base_llm::{
auth::AuthScheme,
chat::{
streaming::{ChatStream, StreamShape},
streaming::StreamShape,
transformation::{
BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData,
Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param,
},
},
messages::streaming::event_stream as messages_event_stream,
},
};
@ -194,13 +190,6 @@ impl BaseConfig for AnthropicConfig {
Ok(ValidatedEnvironment { headers, auth })
}
fn model_response_iterator(&self, shape: StreamShape) -> Option<ChatStream> {
Some(ChatStream::new(
messages_event_stream,
ModelResponseIterator::new(shape),
))
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
DEFAULT_HEADERS
}
@ -267,7 +256,7 @@ fn anthropic_body(
#[cfg(test)]
mod tests {
use crate::{
Error, ErrorDetail,
Error,
anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG,
base_llm::{
auth::AuthScheme,
@ -676,15 +665,11 @@ mod tests {
#[case::missing_model(json!({"content": [], "usage": {"input_tokens": 1, "output_tokens": 1}}))]
fn errors_on_a_response_missing_required_fields(#[case] body: Value) {
let error = transform_response(body).expect_err("invalid response");
let Error::InvalidResponse(ErrorDetail::Invalid { subject, cause }) = &error else {
panic!("decoding failures keep their JSON cause: {error:?}");
};
assert_eq!(*subject, "messages response");
assert!(cause.downcast_ref::<serde_json::Error>().is_some());
assert_eq!(
error.to_string(),
format!("invalid response: invalid messages response: {cause}")
);
assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}");
let cause = litellm_error::find::<serde_json::Error>(&error).expect("JSON cause");
let report = litellm_error::Report::new(&error).to_string();
assert!(report.contains("invalid messages response"), "{report}");
assert_eq!(report.matches(&cause.to_string()).count(), 1, "{report}");
}
#[test]

19
litellm-rust/crates/llms/src/anthropic/messages/handler.rs
View file

@ -208,24 +208,29 @@ mod tests {
#[case::empty(json!({}), Ok(json!({})))]
#[case::numeric_user_id(
json!({"user_id": 123}),
Err(Error::InvalidRequest("metadata.user_id must be a string, got 123".to_string().into())),
Err(invalid_value("metadata.user_id", "a string", json!(123))),
)]
#[case::boolean_user_id(
json!({"user_id": true}),
Err(Error::InvalidRequest("metadata.user_id must be a string, got true".to_string().into())),
Err(invalid_value("metadata.user_id", "a string", json!(true))),
)]
#[case::not_an_object(
json!(["u-1"]),
Err(Error::InvalidRequest(r#"metadata must be an object, got ["u-1"]"#.to_string().into())),
Err(invalid_value("metadata", "an object", json!(["u-1"]))),
)]
fn validate_anthropic_api_metadata_passes_only_a_string_user_id(
#[case] metadata: Value,
#[case] expected: Result<Value, Error>,
) {
assert_eq!(
validate_anthropic_api_metadata(&metadata).map_err(|error| error.to_string()),
expected.map_err(|error| error.to_string()),
);
assert_eq!(validate_anthropic_api_metadata(&metadata), expected);
}
fn invalid_value(field: &'static str, expected: &'static str, actual: Value) -> Error {
Error::InvalidRequest(crate::ErrorDetail::InvalidValue {
field,
expected,
actual,
})
}
#[rstest]

6
litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs
View file

@ -120,7 +120,8 @@ impl BaseOcrConfig for AzureAICohereParseConfig {
impl AzureAICohereParseConfig {
fn get_complete_url(&self, base: &str) -> Result<String, Error> {
let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?;
let mut url =
reqwest::Url::parse(base).map_err(|error| Error::request_field("api_base", error))?;
if !matches!(url.scheme(), "http" | "https") {
return Err(invalid_api_base());
}
@ -133,13 +134,14 @@ impl AzureAICohereParseConfig {
ApiUrl::parse(url.as_str())
.and_then(|url| url.complete_path(&AZURE_COHERE_PARSE_PATH))
.map(|url| url.into_string())
.map_err(|_| invalid_api_base())
.map_err(|error| Error::request_field("api_base", error))
}
}
fn invalid_api_base() -> Error {
Error::RequestField {
path: "api_base".into(),
cause: None,
}
}

27
litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs
View file

@ -13,7 +13,7 @@ use tokio::time::Instant;
use crate::base_llm::ocr::{
document::InlineDocument,
error::Error,
error::{Error, PagesFault},
handler::{CallHooks, OcrClient, read_json_response},
settings::OcrSettings,
transformation::{
@ -259,14 +259,12 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result<Option<String>, Error>
Some(Value::Array(pages)) if pages.iter().all(Value::is_number) => pages
.iter()
.map(|page| {
let page = page
.as_i64()
.ok_or_else(|| Error::Pages("page index is out of range".into()))?;
let page = page.as_i64().ok_or(Error::Pages(PagesFault::OutOfRange))?;
if page < 0 {
return Err(Error::Pages("negative page index".into()));
return Err(Error::Pages(PagesFault::Negative));
}
page.checked_add(1)
.ok_or_else(|| Error::Pages("page index is out of range".into()))
.ok_or(Error::Pages(PagesFault::OutOfRange))
})
.collect::<Result<BTreeSet<_>, _>>()?
.into_iter()
@ -279,7 +277,7 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result<Option<String>, Error>
token
.as_str()
.map(str::trim)
.ok_or_else(|| Error::Pages("expected only integers or only strings".into()))
.ok_or(Error::Pages(PagesFault::Mixed))
})
.collect::<Result<Vec<_>, _>>()?
.join(","),
@ -289,13 +287,11 @@ fn normalize_pages_param(pages: Option<&Value>) -> Result<Option<String>, Error>
.collect::<Vec<_>>()
.join(","),
Some(_) => {
return Err(Error::Pages(
"expected an array of integers or strings, or a native page range".into(),
));
return Err(Error::Pages(PagesFault::Shape));
}
};
if !normalized.split(',').all(valid_page_token) {
return Err(Error::Pages("invalid native page range".into()));
return Err(Error::Pages(PagesFault::Range));
}
Ok(Some(normalized))
}
@ -581,9 +577,7 @@ impl AzureDocumentIntelligenceOcrConfig {
)
.into_string()
})
.map_err(|_| Error::RequestField {
path: "api_base".into(),
})
.map_err(|error| Error::request_field("api_base", error))
}
async fn resolve_headers(
@ -626,7 +620,10 @@ impl AzureDocumentIntelligenceOcrConfig {
}
let token = super::super::common_utils::resolve_entra(auth, config, env_lookup)
.await?
.ok_or(Error::MissingAzureDocumentIntelligenceCredentials)?;
.ok_or(Error::MissingCredential {
credential: "Azure Document Intelligence credentials",
guidance: "set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID",
})?;
super::super::common_utils::validate_destination(connection, token.source())?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {}", token.value())))

9
litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs
View file

@ -169,7 +169,10 @@ impl AzureAiOcrConfig {
}
let key = super::common_utils::resolve_entra(auth, config, env_lookup)
.await?
.ok_or(Error::MissingAzureAiCredentials)?;
.ok_or(Error::MissingCredential {
credential: "Azure AI credentials",
guidance: "set AZURE_AI_API_KEY or configure Entra ID",
})?;
super::common_utils::validate_destination(connection, key.source())?;
Ok(bearer_headers(connection, key.value()))
}
@ -183,9 +186,7 @@ impl AzureAiOcrConfig {
ApiUrl::parse(&base)
.and_then(|url| url.complete_path(&AZURE_AI_OCR_PATH))
.map(|url| url.into_string())
.map_err(|_| Error::RequestField {
path: "api_base".into(),
})
.map_err(|error| Error::request_field("api_base", error))
}
}

9
litellm-rust/crates/llms/src/base_llm/base_model_iterator.rs
View file

@ -5,6 +5,15 @@ use futures_util::{Stream, StreamExt, stream, stream::BoxStream};
pub type ByteStream = BoxStream<'static, Result<Bytes, io::Error>>;
/// The transport error a body stream carried as an `io::Error`, however many times framing
/// wrapped it on the way.
pub(crate) fn tunnelled_http(error: &io::Error) -> Option<&litellm_http::Error> {
let inner = error.get_ref()?;
inner
.downcast_ref::<litellm_http::Error>()
.or_else(|| tunnelled_http(inner.downcast_ref::<io::Error>()?))
}
pub trait StreamTransformer {
type Input;
type Output;

44
litellm-rust/crates/llms/src/base_llm/ocr/document.rs
View file

@ -58,9 +58,7 @@ pub async fn inline_remote_document(
validate_inline_document(&document)?;
return Ok(document);
}
let url = Url::parse(source).map_err(|_| Error::RequestField {
path: "document URL".into(),
})?;
let url = Url::parse(source).map_err(|error| Error::request_field("document URL", error))?;
let downloaded = fetcher
.fetch(
url,
@ -93,20 +91,17 @@ mod tests {
}
}
#[test]
fn decodes_data_urls_and_limits_decoded_size() {
for (source, expected) in [
("data:application/pdf;base64,YWJj", b"abc".as_slice()),
("DATA:application/pdf;BASE64,YWI", b"ab".as_slice()),
("data:,a%20b%00%FF", b"a b\0\xff".as_slice()),
] {
let inline = InlineDocument::parse(source).unwrap().unwrap();
assert_eq!(inline.decode(expected.len()).unwrap(), expected);
assert!(matches!(
inline.decode(expected.len() - 1),
Err(Error::InlineDocumentTooLarge)
));
}
#[rstest::rstest]
#[case::base64("data:application/pdf;base64,YWJj", b"abc")]
#[case::uppercase_unpadded("DATA:application/pdf;BASE64,YWI", b"ab")]
#[case::percent_encoded("data:,a%20b%00%FF", b"a b\0\xff")]
fn decodes_data_urls_and_limits_decoded_size(#[case] source: &str, #[case] expected: &[u8]) {
let inline = InlineDocument::parse(source).unwrap().unwrap();
assert_eq!(inline.decode(expected.len()).unwrap(), expected);
assert!(matches!(
inline.decode(expected.len() - 1),
Err(Error::InlineDocumentTooLarge)
));
}
#[test]
@ -124,15 +119,12 @@ mod tests {
);
}
#[test]
fn rejects_invalid_inline_documents() {
for source in [
"https://example.com/document.pdf",
"data:application/pdf;base64",
"data:application/pdf;base64,INVALID!",
] {
assert!(validate_inline_document(&document(source)).is_err());
}
#[rstest::rstest]
#[case::remote("https://example.com/document.pdf")]
#[case::no_payload("data:application/pdf;base64")]
#[case::invalid_base64("data:application/pdf;base64,INVALID!")]
fn rejects_invalid_inline_documents(#[case] source: &str) {
assert!(validate_inline_document(&document(source)).is_err());
}
#[tokio::test]

315
litellm-rust/crates/llms/src/base_llm/ocr/error.rs
View file

@ -1,3 +1,8 @@
use std::{io, path::PathBuf, sync::Arc};
use litellm_core_utils::call_arguments::ArgumentError;
use litellm_error::Source;
#[derive(Clone, Debug, thiserror::Error)]
pub enum Error {
#[error("upstream OCR error ({status}): {body}")]
@ -8,37 +13,41 @@ pub enum Error {
},
#[error("File is empty or could not be read")]
EmptyFile,
#[error("Failed to read OCR file {}: {cause}", path.display())]
#[error("could not read OCR file {}", path.display())]
FileRead {
path: std::path::PathBuf,
cause: std::sync::Arc<std::io::Error>,
path: PathBuf,
#[source]
cause: Arc<io::Error>,
},
#[error("OCR document preparation task failed: {0}")]
DocumentTask(std::sync::Arc<tokio::task::JoinError>),
#[error("Invalid MIME type: {0}")]
#[error("OCR document preparation task failed")]
DocumentTask(#[source] Arc<tokio::task::JoinError>),
#[error("invalid MIME type {0:?}")]
InvalidMimeType(String),
#[error(
"Cohere Parse only accepts `image_url` documents; document_url and PDF inputs are not supported"
)]
CohereImageOnly,
#[error("{provider} cannot read this document - {guidance}")]
UnsupportedDocument {
provider: &'static str,
guidance: &'static str,
},
#[error("Invalid `req_format`. Expected 'native' or 'litellm'.")]
RequestFormat,
#[error("invalid OCR request field: {path}")]
RequestField { path: String },
RequestField {
path: String,
#[source]
cause: Option<Source>,
},
#[error("invalid OCR optional params")]
Arguments(#[from] ArgumentError),
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("Document URL is required")]
#[error("document URL is required")]
MissingDocumentUrl,
#[error("invalid OCR document data URI")]
InvalidDataUri,
#[error(
"Reducto requires a reducto:// id or a data URI; plain HTTP URLs are not supported, upload the file first"
)]
ReductoSource,
#[error("inline OCR document exceeds the size limit")]
InlineDocumentTooLarge,
#[error("invalid OCR pages: {0}")]
Pages(String),
Pages(PagesFault),
#[error("invalid OCR features")]
Features,
#[error("OCR model cannot be a dot segment")]
@ -46,7 +55,11 @@ pub enum Error {
#[error("OCR response exceeds the size limit of {limit} bytes")]
TooLarge { limit: usize },
#[error("invalid OCR response field: {path}")]
ResponseField { path: String },
ResponseField {
path: String,
#[source]
cause: Option<Source>,
},
#[error("OCR response is missing non-empty content")]
EmptyContent,
#[error("OCR operation ended with status {0}")]
@ -69,24 +82,17 @@ pub enum Error {
model: String,
supported: &'static [&'static str],
},
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error(
"invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID"
)]
MissingAzureAiCredentials,
#[error(
"invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID"
)]
MissingAzureDocumentIntelligenceCredentials,
#[error(
"Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()"
)]
MissingReductoApiKey,
#[error("secret resolution failed: {0}")]
Secret(std::sync::Arc<litellm_secrets::Error>),
#[error("Missing {credential} - {guidance}")]
MissingCredential {
credential: &'static str,
guidance: &'static str,
},
#[error("secret resolution failed")]
Secret(#[source] Arc<litellm_secrets::Error>),
#[error("the OCR host driver failed")]
Host(#[source] litellm_host::machine::MachineFault),
#[error("the OCR host call raised an exception")]
HostCall,
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
@ -97,71 +103,222 @@ pub enum Error {
Http(#[from] litellm_http::Error),
}
impl From<litellm_host::machine::MachineFault> for Error {
fn from(fault: litellm_host::machine::MachineFault) -> Self {
use litellm_host::machine::MachineFault;
Self::InvalidRequest(match fault {
MachineFault::Abandoned => "OCR host driver was abandoned".into(),
MachineFault::Protocol(message) => format!("OCR {message}"),
})
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum PagesFault {
#[error("page index is out of range")]
OutOfRange,
#[error("negative page index")]
Negative,
#[error("expected only integers or only strings")]
Mixed,
#[error("expected an array of integers or strings, or a native page range")]
Shape,
#[error("invalid native page range")]
Range,
}
impl From<litellm_core_utils::call_arguments::ArgumentError> for Error {
fn from(error: litellm_core_utils::call_arguments::ArgumentError) -> Self {
Self::RequestField {
path: format!("optional_params.{}", error.path),
}
impl From<litellm_host::machine::MachineFault> for Error {
fn from(fault: litellm_host::machine::MachineFault) -> Self {
Self::Host(fault)
}
}
impl Error {
pub fn request_field(
path: impl Into<String>,
cause: impl std::error::Error + Send + Sync + 'static,
) -> Self {
Self::RequestField {
path: path.into(),
cause: Some(Source::new(cause)),
}
}
pub fn response_field(
path: impl Into<String>,
cause: impl std::error::Error + Send + Sync + 'static,
) -> Self {
Self::ResponseField {
path: path.into(),
cause: Some(Source::new(cause)),
}
}
pub fn upstream_response(&self) -> Option<(u16, &str, &[(String, String)])> {
match self {
Self::Provider {
status,
body,
headers,
} => Some((*status, body, headers)),
Self::Http(error) => error
.upstream_response()
.map(|(status, body)| (status, body, &[][..])),
Self::EmptyFile
| Self::FileRead { .. }
| Self::DocumentTask(_)
| Self::InvalidMimeType(_)
| Self::UnsupportedDocument { .. }
| Self::RequestFormat
| Self::RequestField { .. }
| Self::Arguments(_)
| Self::MissingField(_)
| Self::MissingDocumentUrl
| Self::InvalidDataUri
| Self::InlineDocumentTooLarge
| Self::Pages(_)
| Self::Features
| Self::DotModel
| Self::TooLarge { .. }
| Self::ResponseField { .. }
| Self::EmptyContent
| Self::OperationStatus(_)
| Self::NumericRange(_)
| Self::PollLocation
| Self::PollOrigin
| Self::PollTimeout
| Self::Unsupported(_)
| Self::InvalidProvider(_)
| Self::InvalidModel { .. }
| Self::MissingCredential { .. }
| Self::Secret(_)
| Self::Host(_)
| Self::HostCall
| Self::Auth(_)
| Self::Params(_)
| Self::Headers(_) => None,
}
}
pub fn is_missing_credential(&self) -> bool {
match self {
Self::MissingAzureAiCredentials
| Self::MissingAzureDocumentIntelligenceCredentials
| Self::MissingReductoApiKey => true,
Self::MissingCredential { .. } => true,
Self::Auth(error) => error.is_missing_credential(),
_ => false,
Self::Provider { .. }
| Self::EmptyFile
| Self::FileRead { .. }
| Self::DocumentTask(_)
| Self::InvalidMimeType(_)
| Self::UnsupportedDocument { .. }
| Self::RequestFormat
| Self::RequestField { .. }
| Self::Arguments(_)
| Self::MissingField(_)
| Self::MissingDocumentUrl
| Self::InvalidDataUri
| Self::InlineDocumentTooLarge
| Self::Pages(_)
| Self::Features
| Self::DotModel
| Self::TooLarge { .. }
| Self::ResponseField { .. }
| Self::EmptyContent
| Self::OperationStatus(_)
| Self::NumericRange(_)
| Self::PollLocation
| Self::PollOrigin
| Self::PollTimeout
| Self::Unsupported(_)
| Self::InvalidProvider(_)
| Self::InvalidModel { .. }
| Self::Secret(_)
| Self::Host(_)
| Self::HostCall
| Self::Params(_)
| Self::Headers(_)
| Self::Http(_) => false,
}
}
pub fn is_unsupported(&self) -> bool {
matches!(self, Self::Unsupported(_))
match self {
Self::Unsupported(_) => true,
Self::Provider { .. }
| Self::EmptyFile
| Self::FileRead { .. }
| Self::DocumentTask(_)
| Self::InvalidMimeType(_)
| Self::UnsupportedDocument { .. }
| Self::RequestFormat
| Self::RequestField { .. }
| Self::Arguments(_)
| Self::MissingField(_)
| Self::MissingDocumentUrl
| Self::InvalidDataUri
| Self::InlineDocumentTooLarge
| Self::Pages(_)
| Self::Features
| Self::DotModel
| Self::TooLarge { .. }
| Self::ResponseField { .. }
| Self::EmptyContent
| Self::OperationStatus(_)
| Self::NumericRange(_)
| Self::PollLocation
| Self::PollOrigin
| Self::PollTimeout
| Self::InvalidProvider(_)
| Self::InvalidModel { .. }
| Self::MissingCredential { .. }
| Self::Secret(_)
| Self::Host(_)
| Self::HostCall
| Self::Auth(_)
| Self::Params(_)
| Self::Headers(_)
| Self::Http(_) => false,
}
}
pub fn http_status_code(&self) -> Option<u16> {
match self {
Self::Provider { status, .. } => Some(*status),
Self::Http(error) if error.status().is_some() => error.status(),
error if error.is_request() => Some(400),
_ => None,
if let Some((status, ..)) = self.upstream_response() {
return Some(status);
}
if let Self::Http(error) = self
&& let Some(status) = error.status()
{
return Some(status);
}
self.is_request().then_some(400)
}
pub fn is_request(&self) -> bool {
match self {
Self::EmptyFile
| Self::InvalidMimeType(_)
| Self::CohereImageOnly
| Self::UnsupportedDocument { .. }
| Self::RequestFormat
| Self::RequestField { .. }
| Self::Arguments(_)
| Self::MissingField(_)
| Self::MissingDocumentUrl
| Self::InvalidDataUri
| Self::ReductoSource
| Self::InlineDocumentTooLarge
| Self::Pages(_)
| Self::Features
| Self::DotModel
| Self::InvalidRequest(_)
| Self::InvalidProvider(_)
| Self::InvalidModel { .. }
| Self::Params(_)
| Self::Headers(_) => true,
Self::Auth(error) => error.is_missing_setting(),
Self::Auth(error) => !error.is_missing_credential(),
Self::Http(error) => error.is_request(),
_ => false,
Self::Provider { .. }
| Self::FileRead { .. }
| Self::DocumentTask(_)
| Self::TooLarge { .. }
| Self::ResponseField { .. }
| Self::EmptyContent
| Self::OperationStatus(_)
| Self::NumericRange(_)
| Self::PollLocation
| Self::PollOrigin
| Self::PollTimeout
| Self::Unsupported(_)
| Self::MissingCredential { .. }
| Self::Secret(_)
| Self::Host(_)
| Self::HostCall => false,
}
}
@ -174,10 +331,34 @@ impl Error {
| Self::NumericRange(_)
| Self::PollLocation
| Self::PollOrigin
| Self::PollTimeout
| Self::InvalidResponse(_) => true,
| Self::PollTimeout => true,
Self::Http(error) => error.is_response(),
_ => false,
Self::Provider { .. }
| Self::EmptyFile
| Self::FileRead { .. }
| Self::DocumentTask(_)
| Self::InvalidMimeType(_)
| Self::UnsupportedDocument { .. }
| Self::RequestFormat
| Self::RequestField { .. }
| Self::Arguments(_)
| Self::MissingField(_)
| Self::MissingDocumentUrl
| Self::InvalidDataUri
| Self::InlineDocumentTooLarge
| Self::Pages(_)
| Self::Features
| Self::DotModel
| Self::Unsupported(_)
| Self::InvalidProvider(_)
| Self::InvalidModel { .. }
| Self::MissingCredential { .. }
| Self::Secret(_)
| Self::Host(_)
| Self::HostCall
| Self::Auth(_)
| Self::Params(_)
| Self::Headers(_) => false,
}
}
}

68
litellm-rust/crates/llms/src/base_llm/ocr/handler.rs
View file

@ -135,15 +135,10 @@ pub async fn ocr<C: BaseOcrConfig>(
.await
.map_err(litellm_http::Error::from)?;
if !response.status().is_success() {
let status = response.status().as_u16();
let headers = header_pairs(response.headers());
return match read_response_bytes(response, request.connection.max_response_bytes).await {
Err(Error::Http(error)) => Err(match error.upstream_response() {
Some((status, body)) => config.get_error_class(body.to_owned(), status, headers),
None => Error::Http(error),
}),
Err(error) => Err(error),
Ok(_) => unreachable!("non-success response produces an HTTP error"),
};
let body = read_error_body(response, request.connection.max_response_bytes).await?;
return Err(config.get_error_class(body, status, headers));
}
let context = OcrResponseContext {
client,
@ -171,35 +166,43 @@ pub async fn read_response_bytes(
mut response: reqwest::Response,
limit: usize,
) -> Result<Bytes, Error> {
let status = response.status();
if status.is_success()
&& response
.content_length()
.is_some_and(|length| length > limit as u64)
if !response.status().is_success() {
let status = response.status().as_u16();
let headers = header_pairs(response.headers());
let body = read_error_body(response, limit).await?;
return Err(Error::Provider {
status,
body,
headers,
});
}
if response
.content_length()
.is_some_and(|length| length > limit as u64)
{
return Err(Error::TooLarge { limit });
}
let mut bytes = BytesMut::new();
while let Some(chunk) = response.chunk().await.map_err(litellm_http::Error::from)? {
let remaining = limit.saturating_sub(bytes.len());
if status.is_success() && chunk.len() > remaining {
if chunk.len() > limit.saturating_sub(bytes.len()) {
return Err(Error::TooLarge { limit });
}
bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
if !status.is_success() && bytes.len() == limit {
break;
}
}
if !status.is_success() {
return Err(litellm_http::Error::upstream(
status.as_u16(),
String::from_utf8_lossy(&bytes).into_owned(),
)
.into());
bytes.extend_from_slice(&chunk);
}
Ok(bytes.freeze())
}
async fn read_error_body(mut response: reqwest::Response, limit: usize) -> Result<String, Error> {
let mut bytes = BytesMut::new();
while bytes.len() < limit
&& let Some(chunk) = response.chunk().await.map_err(litellm_http::Error::from)?
{
let remaining = limit - bytes.len();
bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
}
Ok(String::from_utf8_lossy(&bytes).into_owned())
}
pub async fn transform_request_body<C: BaseOcrConfig, B: Serialize>(
config: &C,
request: &PreparedOcrRequest,
@ -221,6 +224,7 @@ pub async fn transform_request_body<C: BaseOcrConfig, B: Serialize>(
if !changed.body.is_object() {
return Err(Error::RequestField {
path: "guardrail.body".into(),
cause: None,
});
}
config.validate_request_body(&changed.body)?;
@ -264,9 +268,8 @@ pub async fn guardrail_document(
headers: &[(String, String)],
hooks: &dyn CallHooks<Error>,
) -> Result<(OcrDocument, Vec<(String, String)>), Error> {
let body = serde_json::to_value(&request.document).map_err(|_| Error::RequestField {
path: "document".into(),
})?;
let body = serde_json::to_value(&request.document)
.map_err(|error| Error::request_field("document", error))?;
let changed = hooks
.before_provider_request(wire_request(url, headers, body))
.await?;
@ -280,6 +283,7 @@ pub fn body_document(body: &Value) -> Result<OcrDocument, Error> {
.and_then(Value::as_object)
.ok_or_else(|| Error::RequestField {
path: "body.document".into(),
cause: None,
})?;
let source = document
.iter()
@ -388,9 +392,7 @@ mod tests {
.await
.unwrap_err();
let Error::Http(error) = error else {
panic!("unexpected error: {error}");
};
assert_eq!(error.upstream_response(), Some((429, prefix.as_str())));
let (status, body, _) = error.upstream_response().expect("the provider answered");
assert_eq!((status, body), (429, prefix.as_str()));
}
}

26
litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs
View file

@ -194,17 +194,13 @@ pub struct DecodedOcrResponse<T> {
pub fn decode_request_value<T: DeserializeOwned>(value: Value, prefix: &str) -> Result<T, Error> {
serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| {
Error::RequestField {
path: format!("{prefix}.{}", error.path()),
}
Error::request_field(format!("{prefix}.{}", error.path()), error.into_inner())
})
}
pub fn decode_response_value<T: DeserializeOwned>(value: Value, prefix: &str) -> Result<T, Error> {
serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| {
Error::ResponseField {
path: format!("{prefix}.{}", error.path()),
}
Error::response_field(format!("{prefix}.{}", error.path()), error.into_inner())
})
}
@ -213,19 +209,15 @@ pub fn decode_response<T: DeserializeOwned>(
native: bool,
) -> Result<DecodedOcrResponse<T>, Error> {
let mut deserializer = serde_json::Deserializer::from_slice(bytes);
let data = serde_path_to_error::deserialize(&mut deserializer).map_err(|error| {
Error::ResponseField {
path: error.path().to_string(),
}
})?;
deserializer.end().map_err(|_| Error::ResponseField {
path: "response".into(),
})?;
let data = serde_path_to_error::deserialize(&mut deserializer)
.map_err(|error| Error::response_field(error.path().to_string(), error.into_inner()))?;
deserializer
.end()
.map_err(|error| Error::response_field("response", error))?;
let native = if native {
Some(
serde_json::from_slice(bytes).map_err(|_| Error::ResponseField {
path: "response".into(),
})?,
serde_json::from_slice(bytes)
.map_err(|error| Error::response_field("response", error))?,
)
} else {
None

55
litellm-rust/crates/llms/src/base_llm/sse.rs
View file

@ -15,12 +15,15 @@
use bytes::Bytes;
use futures_util::{StreamExt, TryStreamExt, future::ready, stream::BoxStream};
use litellm_framer::{
frames,
SseError, frames,
sse::{SseCodec, SseEvent, decode_all},
};
use serde::{Serialize, de::DeserializeOwned};
use crate::{Error, ErrorDetail, base_llm::base_model_iterator::ByteStream};
use crate::{
Error, ErrorDetail,
base_llm::base_model_iterator::{ByteStream, tunnelled_http},
};
pub type EventStream<E> = BoxStream<'static, Result<E, Error>>;
@ -70,8 +73,13 @@ pub fn invalid(error: serde_json::Error) -> Error {
Error::InvalidResponse(ErrorDetail::invalid("stream event", error))
}
fn framing(error: litellm_framer::SseError) -> Error {
Error::InvalidResponse(ErrorDetail::failed("stream framing", error))
fn framing(error: SseError) -> Error {
match &error {
SseError::Body(body) if let Some(http) = tunnelled_http(body) => Error::Http(http.clone()),
SseError::Body(_) | SseError::InvalidUtf8(_) => {
Error::InvalidResponse(ErrorDetail::failed("stream framing", error))
}
}
}
#[cfg(test)]
@ -94,6 +102,45 @@ mod tests {
assert_eq!(is_done(&event), done);
}
#[rstest]
#[tokio::test]
async fn a_transport_failure_tunnelled_through_the_body_comes_back_as_itself() {
use futures_util::StreamExt;
let failure = litellm_http::Error::upstream(502, "gateway went away".into());
let bytes: ByteStream = futures_util::stream::iter([
Ok(Bytes::from_static(b"data: {}\n\n")),
Err(std::io::Error::other(failure.clone())),
])
.boxed();
let mut events = event_stream::<serde_json::Value>(bytes, is_done);
assert!(events.next().await.expect("first event").is_ok());
let error = events.next().await.expect("failure").unwrap_err();
assert_eq!(error, Error::Http(failure));
}
#[rstest]
#[tokio::test]
async fn a_body_failure_that_is_not_transport_is_an_invalid_response() {
use futures_util::StreamExt;
let bytes: ByteStream =
futures_util::stream::iter([Err(std::io::Error::other("disk full"))]).boxed();
let error = event_stream::<serde_json::Value>(bytes, is_done)
.next()
.await
.expect("failure")
.unwrap_err();
assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}");
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("disk full")
);
}
#[rstest]
fn the_done_frame_we_send_decodes_as_a_control_frame() {
let events: Vec<serde_json::Value> = decode(DONE_FRAME, is_done).unwrap();

33
litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs
View file

@ -14,7 +14,7 @@ use serde_json::{Map, Value, json};
use strum::IntoStaticStr;
use crate::{
Error,
Error, ErrorDetail,
base_llm::{
audio_transcription::transformation::{
AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, Headers,
@ -59,20 +59,12 @@ fn audio_fields(audio: Value) -> Result<AudioInput, Error> {
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
.ok_or(Error::MissingField("audio.data"))?;
let format = object.get("format").cloned().ok_or_else(|| {
Error::InvalidRequest(
"audio.format must be wav, mp3, flac, or ogg"
.to_string()
.into(),
)
})?;
let format: AudioFormat = serde_json::from_value(format).map_err(|_| {
Error::InvalidRequest(
"audio.format must be wav, mp3, flac, or ogg"
.to_string()
.into(),
)
})?;
let format = object
.get("format")
.cloned()
.ok_or(Error::MissingField("audio.format"))?;
let format: AudioFormat = serde_json::from_value(format)
.map_err(|error| Error::InvalidRequest(ErrorDetail::invalid("audio.format", error)))?;
Ok(AudioInput {
data: data.to_string(),
format,
@ -135,14 +127,13 @@ impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig {
) -> Result<AudioTranscriptionResponseData, Error> {
let response: ConverseResponse =
serde_json::from_value(response_json).map_err(|error| {
Error::InvalidResponse(format!("invalid Converse response: {error}").into())
Error::InvalidResponse(ErrorDetail::invalid("Converse response", error))
})?;
if response.message_content_is_non_text() {
return Err(Error::InvalidResponse(
"Bedrock response contains non-text transcript content"
.to_string()
.into(),
));
return Err(Error::InvalidResponse(ErrorDetail::Expected {
subject: "Bedrock transcript content",
expected: "text",
}));
}
Ok(AudioTranscriptionResponseData {
text: response.content_text(),

69
litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs
View file

@ -19,7 +19,7 @@ use litellm_llms_types::{
use serde_json::{Map, Value, json};
use crate::{
Error,
Error, ErrorDetail,
base_llm::{
auth::AuthScheme,
chat::{
@ -154,21 +154,16 @@ impl BaseConfig for AmazonConverseConfig {
response: ProviderChatResponseData,
) -> Result<ChatCompletionsResponse, Error> {
let body = response.body;
if !body.is_object() {
return Err(Error::InvalidResponse(
"converse response is not an object".into(),
));
}
for field in ["output", "usage"] {
if body.get(field).is_none() {
return Err(Error::InvalidResponse(
format!("invalid Converse response: missing field `{field}`").into(),
));
}
}
let has_usage = body.get("usage").is_some();
let response: ConverseResponse = serde_json::from_value(body).map_err(|error| {
Error::InvalidResponse(format!("invalid Converse response: {error}").into())
Error::InvalidResponse(ErrorDetail::invalid("Converse response", error))
})?;
if !has_usage {
return Err(Error::InvalidResponse(ErrorDetail::MissingField {
subject: "Converse response",
field: "usage",
}));
}
// The route declines tool requests, so anything other than a text block
// is something this path never asked for. Decline; the host falls back.
if response.message_content_is_non_text() {
@ -858,34 +853,24 @@ mod tests {
assert_eq!(err, Error::Unsupported("non-text response content block"));
}
#[test]
fn errors_on_a_response_missing_required_fields() {
assert_eq!(
transform_response(json!("nope")).expect_err("not an object"),
Error::InvalidResponse("converse response is not an object".to_string().into())
);
assert_eq!(
transform_response(json!({"usage": {}})).expect_err("no output"),
Error::InvalidResponse("invalid Converse response: missing field `output`".into())
);
assert_eq!(
transform_response(json!({"output": {"message": {"content": []}}}))
.expect_err("no usage"),
Error::InvalidResponse("invalid Converse response: missing field `usage`".into())
);
}
#[test]
fn rejects_malformed_text_and_token_counts() {
for response in [
json!({"output": {"message": {"content": [{"text": 123}]}}, "usage": {"inputTokens": 1, "outputTokens": 1}}),
json!({"output": {"message": {"content": [{"text": "x"}]}}, "usage": {"inputTokens": "1", "outputTokens": 1}}),
] {
assert!(matches!(
transform_response(response),
Err(Error::InvalidResponse(message)) if message.to_string().starts_with("invalid Converse response:")
));
}
#[rstest::rstest]
#[case::not_an_object(json!("nope"), "invalid type")]
#[case::no_output(json!({"usage": {"inputTokens": 1, "outputTokens": 1}}), "missing field `output`")]
#[case::no_usage(json!({"output": {"message": {"content": []}}}), "missing usage")]
#[case::non_string_text(
json!({"output": {"message": {"content": [{"text": 123}]}}, "usage": {"inputTokens": 1, "outputTokens": 1}}),
"invalid type"
)]
#[case::non_integer_tokens(
json!({"output": {"message": {"content": [{"text": "x"}]}}, "usage": {"inputTokens": "1", "outputTokens": 1}}),
"invalid type"
)]
fn rejects_a_malformed_response(#[case] response: Value, #[case] fact: &str) {
let error = transform_response(response).expect_err("malformed response");
assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}");
let report = litellm_error::Report::new(&error).to_string();
assert!(report.contains("Converse response"), "{report}");
assert!(report.contains(fact), "{report}");
}
#[test]

66
litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs
View file

@ -2,6 +2,7 @@ use base64::Engine;
use bytes::Buf;
use futures_util::{Stream, StreamExt};
use litellm_framer::{
EventStreamError,
aws_event_stream::{AwsEventStreamCodec, Message},
frames,
};
@ -11,9 +12,8 @@ use serde_json::Value;
use crate::{
Error,
anthropic::chat::handler::ModelResponseIterator,
base_llm::{
chat::streaming::{ChatStream, StreamShape},
base_model_iterator::tunnelled_http,
messages::streaming::{ByteStream, EventStream},
},
};
@ -47,11 +47,21 @@ where
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
frames(input, AwsEventStreamCodec).map(|message| {
decode_invoke_chunk(message.map_err(|error| {
frames(input, AwsEventStreamCodec).map(|message| decode_invoke_chunk(message.map_err(framing)?))
}
fn framing(error: EventStreamError) -> Error {
match &error {
EventStreamError::Body(body) if let Some(http) = tunnelled_http(body) => {
Error::Http(http.clone())
}
EventStreamError::Body(_)
| EventStreamError::InvalidLength(_)
| EventStreamError::Truncated
| EventStreamError::Malformed(_) => {
Error::InvalidResponse(crate::ErrorDetail::failed("stream framing", error))
})?)
})
}
}
}
pub fn decode_invoke_anthropic_chunk(chunk: Value) -> Result<MessagesStreamEvent, Error> {
@ -64,34 +74,6 @@ pub fn invoke_anthropic_event_stream(bytes: ByteStream) -> EventStream {
Box::pin(invoke_chunk_stream(bytes).map(|chunk| decode_invoke_anthropic_chunk(chunk?)))
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString)]
enum InvokeProvider {
#[strum(serialize = "anthropic")]
Anthropic,
#[strum(serialize = "deepseek_r1")]
DeepseekR1,
#[strum(serialize = "moonshot")]
Moonshot,
#[strum(disabled)]
Unsupported,
}
pub fn invoke_chat_stream(invoke_provider: &str, shape: StreamShape) -> Result<ChatStream, Error> {
match invoke_provider
.parse()
.unwrap_or(InvokeProvider::Unsupported)
{
InvokeProvider::Anthropic => Ok(ChatStream::new(
invoke_anthropic_event_stream,
ModelResponseIterator::new(shape),
)),
InvokeProvider::DeepseekR1 | InvokeProvider::Moonshot => Err(Error::Unsupported(
"Bedrock invoke streaming for this model family",
)),
InvokeProvider::Unsupported => Err(Error::Unsupported("Bedrock invoke streaming")),
}
}
#[cfg(test)]
mod tests {
use aws_smithy_eventstream::frame::write_message_to;
@ -112,22 +94,6 @@ mod tests {
const TEXT_DELTA: &str =
r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#;
#[rstest::rstest]
#[case::anthropic("anthropic", InvokeProvider::Anthropic)]
#[case::deepseek_r1("deepseek_r1", InvokeProvider::DeepseekR1)]
#[case::moonshot("moonshot", InvokeProvider::Moonshot)]
#[case::unknown("qwen", InvokeProvider::Unsupported)]
#[case::case_sensitive("Anthropic", InvokeProvider::Unsupported)]
fn invoke_provider_maps_each_model_family_name(
#[case] name: &str,
#[case] expected: InvokeProvider,
) {
assert_eq!(
name.parse().unwrap_or(InvokeProvider::Unsupported),
expected
);
}
fn aws_wire(chunk: &str) -> Vec<u8> {
let payload = serde_json::json!({"bytes": STANDARD.encode(chunk)});
let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header(

41
litellm-rust/crates/llms/src/cohere/ocr/transformation.rs
View file

@ -207,27 +207,28 @@ impl CohereParseConfig {
}
fn build_ocr_url(&self, api_base: &str) -> Result<String, Error> {
let parsed = reqwest::Url::parse(api_base).map_err(|_| invalid_api_base())?;
let parsed = reqwest::Url::parse(api_base)
.map_err(|error| Error::request_field("api_base", error))?;
if !matches!(parsed.scheme(), "http" | "https") {
return Err(invalid_api_base());
}
ApiUrl::parse(api_base)
.and_then(|url| url.complete_path(&["v2", "parse"]))
.map(|url| url.into_string())
.map_err(|_| invalid_api_base())
.map_err(|error| Error::request_field("api_base", error))
}
}
pub fn validate_document(document: &OcrDocument) -> Result<(), Error> {
let OcrDocument::ImageUrl { image_url, .. } = document else {
return Err(Error::CohereImageOnly);
return Err(image_only());
};
if image_url.is_empty() {
return Err(Error::CohereImageOnly);
return Err(image_only());
}
if let Some(inline) = InlineDocument::parse(image_url)? {
if !inline.mime_type().type_.eq_ignore_ascii_case("image") {
return Err(Error::CohereImageOnly);
return Err(image_only());
}
inline.decode(OCR_INLINE_MAX_BYTES)?;
}
@ -259,7 +260,7 @@ pub fn normalize_response(
fn image_url(document: OcrDocument) -> Result<String, Error> {
validate_document(&document)?;
let OcrDocument::ImageUrl { image_url, .. } = document else {
return Err(Error::CohereImageOnly);
return Err(image_only());
};
Ok(image_url)
}
@ -328,9 +329,17 @@ fn billed_pages(response: &CohereResponse) -> Option<i64> {
response.meta.as_ref()?.billed_units.as_ref()?.pages
}
fn image_only() -> Error {
Error::UnsupportedDocument {
provider: "Cohere Parse",
guidance: "it only accepts `image_url` documents, not document_url or PDF inputs",
}
}
fn invalid_api_base() -> Error {
Error::RequestField {
path: "api_base".into(),
cause: None,
}
}
@ -367,11 +376,14 @@ mod tests {
#[test]
fn options_reject_invalid_output_format() {
let invalid = serde_json::from_value(json!({"output_format":"html"})).unwrap();
assert!(matches!(
CohereParseConfig.map_ocr_params(&invalid, "parse"),
Err(Error::RequestField { path })
if path == "optional_params.output_format"
));
let Err(error) = CohereParseConfig.map_ocr_params(&invalid, "parse") else {
panic!("html is not a Cohere output format");
};
assert!(error.is_request());
let argument =
litellm_error::find::<litellm_core_utils::call_arguments::ArgumentError>(&error)
.expect("the argument failure is kept");
assert_eq!(argument.path, "output_format");
}
#[test]
@ -431,7 +443,7 @@ mod tests {
.unwrap();
assert!(matches!(
normalize_response("parse", response).unwrap_err(),
Error::ResponseField { path }
Error::ResponseField { path, .. }
if path == "pages[0].markdown.images[0].image_base64"
));
}
@ -656,7 +668,10 @@ mod tests {
fn request_requires_image(#[case] value: Value) {
assert!(matches!(
validate_document(&serde_json::from_value(value).unwrap()),
Err(Error::CohereImageOnly)
Err(Error::UnsupportedDocument {
provider: "Cohere Parse",
..
})
));
}

104
litellm-rust/crates/llms/src/error.rs
View file

@ -1,3 +1,5 @@
use litellm_error::Source;
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("expected {expected}, got {actual}")]
@ -7,29 +9,41 @@ pub enum Error {
},
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("invalid request: {0}")]
InvalidRequest(ErrorDetail),
#[error("invalid response: {0}")]
InvalidResponse(ErrorDetail),
#[error("invalid request")]
InvalidRequest(#[source] ErrorDetail),
#[error("invalid response")]
InvalidResponse(#[source] ErrorDetail),
#[error("unsupported: {0}")]
Unsupported(&'static str),
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Http(#[from] litellm_http::Error),
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum ErrorDetail {
#[error("{0}")]
Message(String),
#[error("invalid {subject}: {cause}")]
#[error("invalid {subject}")]
Invalid {
subject: &'static str,
cause: ErrorSource,
#[source]
cause: Source,
},
#[error("{operation} failed: {cause}")]
#[error("{operation} failed")]
Failed {
operation: &'static str,
cause: ErrorSource,
#[source]
cause: Source,
},
#[error("{subject} must be {expected}")]
Expected {
subject: &'static str,
expected: &'static str,
},
#[error("{subject} is missing {field}")]
MissingField {
subject: &'static str,
field: &'static str,
},
#[error("{field} must be {expected}, got {actual}")]
InvalidValue {
@ -58,11 +72,18 @@ pub enum ErrorDetail {
value: String,
hint: String,
},
#[error("invalid {subject} on line {line}: {cause}")]
#[error(
"Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({{LLM_PROVIDER}}_API_KEY) or via params"
)]
MissingApiBase,
#[error("the host call raised an exception")]
HostCall,
#[error("invalid {subject} on line {line}")]
InvalidLine {
subject: &'static str,
line: usize,
cause: ErrorSource,
#[source]
cause: Source,
},
#[error("{operation} failed: {detail}")]
RemoteFailure {
@ -78,7 +99,7 @@ impl ErrorDetail {
) -> Self {
Self::Invalid {
subject,
cause: ErrorSource::new(cause),
cause: Source::new(cause),
}
}
@ -88,50 +109,19 @@ impl ErrorDetail {
) -> Self {
Self::Failed {
operation,
cause: ErrorSource::new(cause),
cause: Source::new(cause),
}
}
pub fn invalid_line(
subject: &'static str,
line: usize,
cause: impl std::error::Error + Send + Sync + 'static,
) -> Self {
Self::InvalidLine {
subject,
line,
cause: Source::new(cause),
}
}
}
impl From<String> for ErrorDetail {
fn from(message: String) -> Self {
Self::Message(message)
}
}
impl From<&str> for ErrorDetail {
fn from(message: &str) -> Self {
Self::Message(message.into())
}
}
#[derive(Clone, Debug)]
pub struct ErrorSource(std::sync::Arc<dyn std::error::Error + Send + Sync>);
impl std::ops::Deref for ErrorSource {
type Target = dyn std::error::Error + Send + Sync;
fn deref(&self) -> &Self::Target {
self.0.as_ref()
}
}
impl std::fmt::Display for ErrorSource {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
std::fmt::Display::fmt(&self.0, formatter)
}
}
impl ErrorSource {
pub fn new(error: impl std::error::Error + Send + Sync + 'static) -> Self {
Self(std::sync::Arc::new(error))
}
}
impl PartialEq for ErrorSource {
fn eq(&self, other: &Self) -> bool {
std::sync::Arc::ptr_eq(&self.0, &other.0)
}
}
impl Eq for ErrorSource {}

2
litellm-rust/crates/llms/src/lib.rs
View file

@ -17,4 +17,4 @@ pub mod reducto;
pub mod tencent;
pub mod vertex_ai;
pub use error::{Error, ErrorDetail, ErrorSource};
pub use error::{Error, ErrorDetail};

9
litellm-rust/crates/llms/src/mistral/ocr/transformation.rs
View file

@ -171,9 +171,7 @@ impl MistralOcrConfig {
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&["v1", "ocr"]))
.map(|url| url.into_string())
.map_err(|_| Error::RequestField {
path: "api_base".into(),
})
.map_err(|error| Error::request_field("api_base", error))
}
}
@ -186,6 +184,7 @@ pub fn normalize_response(
Some(None) => {
return Err(Error::ResponseField {
path: "model".into(),
cause: None,
});
}
None => model.to_string(),
@ -231,7 +230,7 @@ mod tests {
let response = serde_json::from_value(json!({"model":null})).unwrap();
assert!(matches!(
normalize_response("fallback", response).unwrap_err(),
Error::ResponseField { path } if path == "model"
Error::ResponseField { path, .. } if path == "model"
));
}
@ -263,7 +262,7 @@ mod tests {
.unwrap_err();
assert!(matches!(
error,
Error::ResponseField { path: actual } if actual == path
Error::ResponseField { path: actual, .. } if actual == path
));
}

18
litellm-rust/crates/llms/src/openai/responses/transformation.rs
View file

@ -6,7 +6,7 @@ use serde_json::{Map, Value};
use litellm_auth::{CredentialPlacement, SecretValue};
use crate::{
Error,
Error, ErrorDetail,
base_llm::{
auth::{AuthScheme, ValidatedEnvironment},
responses::transformation::{
@ -180,15 +180,19 @@ pub fn responses_request(
params: Map<String, Value>,
) -> Result<Value, Error> {
if !input.is_string() && !input.is_array() {
return Err(Error::InvalidRequest(
"responses input must be a string or an array".into(),
));
return Err(Error::InvalidRequest(ErrorDetail::Expected {
subject: "responses input",
expected: "a string or an array",
}));
}
if params
.get("stream")
.is_some_and(|stream| !stream.is_boolean())
{
return Err(Error::InvalidRequest("stream must be a boolean".into()));
return Err(Error::InvalidRequest(ErrorDetail::Expected {
subject: "stream",
expected: "a boolean",
}));
}
Ok(Value::Object(
params
@ -202,7 +206,9 @@ pub fn responses_request(
}
pub fn responses_response(body: Value) -> Result<ResponsesApiResponse, Error> {
serde_json::from_value(body).map_err(|error| Error::InvalidResponse(error.to_string().into()))
serde_json::from_value(body).map_err(|error| {
Error::InvalidResponse(ErrorDetail::invalid("Responses API response", error))
})
}
#[cfg(test)]

19
litellm-rust/crates/llms/src/openai_like/chat/transformation.rs
View file

@ -123,9 +123,12 @@ impl BaseConfig for OpenAILikeChatConfig {
) -> Result<ChatCompletionsResponse, Error> {
let mut body = response.body;
sanitize_usage(&mut body);
let body = body
.as_object()
.ok_or_else(|| Error::InvalidResponse("chat response is not an object".into()))?;
let body =
body.as_object()
.ok_or(Error::InvalidResponse(crate::ErrorDetail::Expected {
subject: "chat response",
expected: "an object",
}))?;
let choices = body
.get("choices")
@ -414,9 +417,13 @@ mod tests {
#[rstest]
fn a_missing_api_base_is_an_error() {
assert!(matches!(
OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG
.get_complete_url(None, "my-model", &params(json!({})), &no_env),
Err(Error::InvalidRequest(message)) if message.to_string().starts_with("Missing API Base")
OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG.get_complete_url(
None,
"my-model",
&params(json!({})),
&no_env
),
Err(Error::InvalidRequest(crate::ErrorDetail::MissingApiBase))
));
}

9
litellm-rust/crates/llms/src/openai_like/common_utils.rs
View file

@ -1,7 +1,7 @@
//! Shared OpenAI-like credential and endpoint resolution, mirroring
//! `litellm/llms/openai_like/common_utils.py`.
use crate::Error;
use crate::{Error, ErrorDetail};
/// `OpenAILikeChatConfig._get_openai_compatible_provider_info`: the deployment's
/// `api_base` wins over `OPENAI_LIKE_API_BASE`, and the deployment key over
@ -31,12 +31,7 @@ pub fn complete_openai_like_url(
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
let (api_base, _) = openai_compatible_provider_info(api_base, None, env_lookup);
let api_base = api_base.ok_or_else(|| {
Error::InvalidRequest(
"Missing API Base - A call is being made to LLM Provider but no api base is set either in the environment variables ({LLM_PROVIDER}_API_KEY) or via params"
.into(),
)
})?;
let api_base = api_base.ok_or(Error::InvalidRequest(ErrorDetail::MissingApiBase))?;
if custom_endpoint {
return Ok(api_base);
}

24
litellm-rust/crates/llms/src/reducto/ocr/transformation.rs
View file

@ -299,9 +299,16 @@ async fn prepare_upload_request<C: BaseOcrConfig<Environment = Vec<(String, Stri
build_http_request(request, url, headers, &body)
}
fn unsupported_source() -> Error {
Error::UnsupportedDocument {
provider: "Reducto",
guidance: "pass a reducto:// id or a data URI, plain HTTP URLs must be uploaded first",
}
}
fn uploaded_file_id(document: OcrDocument) -> Result<ReductoFileId, Error> {
if !document.source().starts_with(REDUCTO_ID_PREFIX) {
return Err(Error::ReductoSource);
return Err(unsupported_source());
}
if document.source()[REDUCTO_ID_PREFIX.len()..]
.trim()
@ -309,6 +316,7 @@ fn uploaded_file_id(document: OcrDocument) -> Result<ReductoFileId, Error> {
{
return Err(Error::RequestField {
path: "document file id".into(),
cause: None,
});
}
Ok(ReductoFileId(document.source().into()))
@ -392,6 +400,7 @@ fn build_pages_from_reducto(chunks: Vec<ReductoChunk>) -> Result<Vec<OcrPage>, E
Some(Value::String(content)) => Ok(Some(content.as_str())),
Some(_) => Err(Error::ResponseField {
path: "result.chunks.blocks.content".into(),
cause: None,
}),
})
.collect::<Result<Vec<_>, _>>()?;
@ -436,9 +445,7 @@ fn complete_endpoint_url(api_base: Option<&str>, path: &str) -> Result<String, E
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&[path]))
.map(|url| url.into_string())
.map_err(|_| Error::RequestField {
path: "api_base".into(),
})
.map_err(|error| Error::request_field("api_base", error))
}
fn resolve_headers(
@ -459,7 +466,10 @@ fn resolve_headers(
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
})
.ok_or(Error::MissingReductoApiKey)?;
.ok_or(Error::MissingCredential {
credential: REDUCTO_API_KEY_ENV,
guidance: "set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()",
})?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {api_key}")))
.chain(connection.extra_headers.clone())
@ -494,11 +504,12 @@ async fn ensure_file_id_async(
{
return Err(Error::RequestField {
path: "document file id".into(),
cause: None,
});
}
return Ok(ReductoFileId(document.source().to_string()));
}
let inline = InlineDocument::parse(document.source())?.ok_or(Error::ReductoSource)?;
let inline = InlineDocument::parse(document.source())?.ok_or_else(unsupported_source)?;
let mime = inline.mime_type().to_string();
let bytes = inline.decode(OCR_INLINE_MAX_BYTES)?;
upload_bytes_async(bytes, &mime, headers, context).await
@ -546,6 +557,7 @@ async fn upload_bytes_async(
let Some(file_id) = file_id else {
return Err(Error::ResponseField {
path: "file_id".into(),
cause: None,
});
};
Ok(ReductoFileId(file_id.to_string()))

56
litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs
View file

@ -235,8 +235,9 @@ pub fn normalize_response(
&mut output,
PythonJsonFormatter,
))
.map_err(|_| response_field("content"))?;
String::from_utf8(output).map_err(|_| response_field("content"))?
.map_err(|error| Error::response_field(content_path("content"), error))?;
String::from_utf8(output)
.map_err(|error| Error::response_field(content_path("content"), error))?
};
(data, fallback)
}
@ -261,7 +262,12 @@ pub fn normalize_response(
})
})
.collect::<Result<Vec<_>, Error>>()?,
Some(_) => return Err(response_field("pages")),
Some(_) => {
return Err(Error::ResponseField {
path: content_path("pages"),
cause: None,
});
}
None => Vec::new(),
};
let usage = ocr_data
@ -273,7 +279,12 @@ pub fn normalize_response(
.transpose()?;
let model = match ocr_data.get("model") {
Some(Value::String(model)) => model.clone(),
Some(_) => return Err(response_field("model")),
Some(_) => {
return Err(Error::ResponseField {
path: content_path("model"),
cause: None,
});
}
None => model.to_string(),
};
Ok(LiteLLMOcrResponse {
@ -369,10 +380,8 @@ impl serde_json::ser::Formatter for PythonJsonFormatter {
}
}
fn response_field(field: &str) -> Error {
Error::ResponseField {
path: format!("choices[0].message.content.{field}"),
}
fn content_path(field: &str) -> String {
format!("choices[0].message.content.{field}")
}
pub fn provider_model(model: &str) -> Result<String, Error> {
@ -380,6 +389,7 @@ pub fn provider_model(model: &str) -> Result<String, Error> {
if local_model.is_empty() {
return Err(Error::RequestField {
path: "model".into(),
cause: None,
});
}
Ok(format!("{MODEL_PREFIX}{local_model}"))
@ -411,9 +421,7 @@ impl VertexAIDeepSeekOCRConfig {
])
})
.map(|url| url.into_string())
.map_err(|_| Error::RequestField {
path: "api_base".into(),
})
.map_err(|error| Error::request_field("api_base", error))
}
}
@ -625,20 +633,18 @@ mod tests {
assert_eq!(result["future"], "kept");
}
#[test]
fn response_transform_rejects_missing_empty_and_malformed_content() {
for value in [
json!({"choices":[]}),
json!({"choices":[{"message":{"content":{}}}]}),
json!({"choices":[{"message":{"content":""}}]}),
json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}),
json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}),
] {
let result = serde_json::from_value::<DeepSeekOcrResponse>(value)
.map_err(|_| ())
.and_then(|response| normalize_response("model", response).map_err(|_| ()));
assert!(result.is_err());
}
#[rstest::rstest]
#[case::no_choices(json!({"choices":[]}))]
#[case::empty_object(json!({"choices":[{"message":{"content":{}}}]}))]
#[case::empty_string(json!({"choices":[{"message":{"content":""}}]}))]
#[case::malformed_string(json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}))]
#[case::malformed_object(json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}))]
fn response_transform_rejects_missing_empty_and_malformed_content(#[case] value: Value) {
let rejected = match serde_json::from_value::<DeepSeekOcrResponse>(value) {
Ok(response) => normalize_response("model", response).is_err(),
Err(_) => true,
};
assert!(rejected);
}
#[test]

5
litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs
View file

@ -174,9 +174,7 @@ impl VertexAiOcrConfig {
])
})
.map(|url| url.into_string())
.map_err(|_| Error::RequestField {
path: "api_base".into(),
})
.map_err(|error| Error::request_field("api_base", error))
}
}
@ -198,6 +196,7 @@ fn validate_location(location: &str) -> Result<(), Error> {
}
Err(Error::RequestField {
path: "vertex_location".into(),
cause: None,
})
}

21
litellm-rust/crates/python-bridge/src/errors.rs
View file

@ -21,7 +21,10 @@ pyo3::create_exception!(
pub(crate) fn route_error_to_pyerr(error: Error) -> PyErr {
match error.upstream_response() {
Some((status, body)) => RustUpstreamError::new_err((status, body.to_owned())),
None => by_fault(error.is_request(), error.to_string()),
None => by_fault(
error.is_request(),
litellm_error::Report::new(&error).to_string(),
),
}
}
@ -41,12 +44,24 @@ mod tests {
#[rstest::rstest]
#[case::unsupported(Error::Unsupported("test capability"), true)]
#[case::invalid_provider(Error::InvalidProvider("unknown".into()), true)]
#[case::invalid_request(Error::InvalidRequest("empty messages".into()), true)]
#[case::invalid_request(
Error::InvalidRequest(litellm_llms::ErrorDetail::Expected {
subject: "messages",
expected: "non-empty",
}),
true
)]
#[case::lost_connection(
litellm_http::Error::from(tokio_tungstenite::tungstenite::Error::ConnectionClosed).into(),
false
)]
#[case::invalid_response(Error::InvalidResponse("missing usage".into()), false)]
#[case::invalid_response(
Error::InvalidResponse(litellm_llms::ErrorDetail::MissingField {
subject: "response",
field: "usage",
}),
false
)]
fn route_failures_are_terminal(#[case] error: Error, #[case] is_request: bool) {
Python::initialize();
Python::attach(|py| {

4
litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs
View file

@ -78,8 +78,8 @@ impl PythonBinding for ChatCompletionsPythonHost {
fn map_error(&self, py: Python<'_>, error: Error) -> PyResult<PyErr> {
self.0.error(py, error)
}
fn host_error(error: &PyErr) -> Error {
Error::InvalidRequest(error.to_string().into())
fn host_error(_: &PyErr) -> Error {
Error::InvalidRequest(litellm_llms::ErrorDetail::HostCall)
}
}

16
litellm-rust/crates/python-bridge/src/routes/messages/host.rs
View file

@ -77,7 +77,7 @@ fn native_error(py: Python<'_>, error: Error) -> PyResult<PyErr> {
}
match error {
Error::InvalidRequest(message) => {
let error = PyValueError::new_err(message.to_string());
let error = PyValueError::new_err(litellm_error::Report::new(&message).to_string());
error.value(py).setattr(REQUEST_ERROR_MARKER, true)?;
Ok(error)
}
@ -283,8 +283,8 @@ impl PythonBinding for MessagesPythonHost {
Ok(self.map_failure(py, native_error(py, error)?))
}
fn host_error(error: &PyErr) -> Error {
Error::InvalidRequest(error.to_string().into())
fn host_error(_: &PyErr) -> Error {
Error::InvalidRequest(litellm_llms::ErrorDetail::HostCall)
}
}
@ -465,7 +465,15 @@ mod tests {
}
#[rstest]
#[case::rejected_request(Error::InvalidRequest("does not support top_k=5".into()), true)]
#[case::rejected_request(
Error::InvalidRequest(litellm_llms::ErrorDetail::UnsupportedParameter {
model: "claude".into(),
param: "top_k".into(),
value: "5".into(),
hint: String::new(),
}),
true
)]
#[case::missing_field(Error::MissingField("max_tokens"), true)]
#[case::unresolvable_provider(Error::InvalidProvider("openai".into()), false)]
#[case::upstream_failure(

15
litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs
View file

@ -33,7 +33,10 @@ pub(super) fn to_pyerr(error: Error) -> PyErr {
PyFileNotFoundError::new_err(format!("File not found: {}", path.display()))
}
Error::FileRead { cause, .. } => PyOSError::new_err(cause.to_string()),
other => by_fault(is_request(&other), other.to_string()),
other => by_fault(
is_request(&other),
litellm_error::Report::new(&other).to_string(),
),
})
})
.unwrap_or_else(|error| error);
@ -46,7 +49,6 @@ fn is_request(error: &Error) -> bool {
error,
Error::Auth(_)
| Error::InvalidProvider(_)
| Error::InvalidRequest(_)
| Error::MissingField(_)
| Error::MissingDocumentUrl
)
@ -86,7 +88,10 @@ mod tests {
Python::attach(|py| {
let mapped = to_pyerr(Error::MissingDocumentUrl);
assert!(mapped.is_instance_of::<pyo3::exceptions::PyValueError>(py));
assert_eq!(mapped.value(py).to_string(), "Document URL is required");
assert_eq!(
mapped.value(py).to_string(),
Error::MissingDocumentUrl.to_string()
);
assert_eq!(
mapped
.value(py)
@ -115,7 +120,7 @@ mod tests {
.expect("OCR failures retain status and unprefixed provider message");
assert_eq!(args, (429, r#"{"message":"rate limited"}"#.to_string()));
let mapped = to_pyerr(Error::InvalidRequest("invalid format".into()));
let mapped = to_pyerr(Error::Features);
assert!(mapped.is_instance_of::<PyValueError>(py));
assert_eq!(
mapped
@ -195,7 +200,7 @@ mod tests {
#[rstest::rstest]
#[case::oversized(Error::TooLarge { limit: 7 })]
#[case::malformed_field(Error::ResponseField { path: "pages[0].index".into() })]
#[case::malformed_field(Error::ResponseField { path: "pages[0].index".into(), cause: None })]
fn response_failures_are_statusless_runtime_errors(#[case] error: Error) {
Python::initialize();
Python::attach(|py| {

4
litellm-rust/crates/python-bridge/src/routes/ocr/host.rs
View file

@ -143,8 +143,8 @@ impl PythonBinding for OcrPythonHost {
Ok(self.map_failure(py, ocr_error_to_pyerr(error)))
}
fn host_error(error: &PyErr) -> Error {
Error::InvalidRequest(error.to_string())
fn host_error(_: &PyErr) -> Error {
Error::HostCall
}
}

1
litellm-rust/crates/python-bridge/src/routes/ocr/project.rs
View file

@ -92,6 +92,7 @@ impl ProjectedDocument {
{
ocr_error_to_pyerr(Error::RequestField {
path: "document.type".into(),
cause: None,
})
} else {
error

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