Merge remote-tracking branch 'origin/main' into litellm_/attribution-investigation-a81211

This commit is contained in:
Yuneng Jiang 2026-09-16 17:08:17 -07:00
commit dcde8395ec
No known key found for this signature in database
81 changed files with 5640 additions and 570 deletions

101
litellm-rust/Cargo.lock generated
View file

@ -70,6 +70,12 @@ dependencies = [
"rustversion",
]
[[package]]
name = "arcstr"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
[[package]]
name = "async-compression"
version = "0.4.46"
@ -1837,6 +1843,12 @@ version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
[[package]]
name = "litellm-auth"
version = "0.1.0"
@ -1915,6 +1927,17 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-redis"
version = "0.1.0"
dependencies = [
"litellm-cache",
"redis",
"redis-test",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-core"
version = "0.1.0"
@ -2140,6 +2163,16 @@ dependencies = [
"minimal-lexical",
]
[[package]]
name = "num-bigint"
version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0"
dependencies = [
"num-integer",
"num-traits",
]
[[package]]
name = "num-conv"
version = "0.2.2"
@ -2656,6 +2689,36 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "redis"
version = "1.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
dependencies = [
"arcstr",
"combine",
"itoa",
"num-bigint",
"percent-encoding",
"ryu",
"sha1_smol",
"socket2 0.6.5",
"url",
"xxhash-rust",
]
[[package]]
name = "redis-test"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca"
dependencies = [
"rand 0.9.5",
"redis",
"socket2 0.6.5",
"tempfile",
]
[[package]]
name = "redox_syscall"
version = "0.5.18"
@ -2846,6 +2909,19 @@ dependencies = [
"semver",
]
[[package]]
name = "rustix"
version = "1.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d"
dependencies = [
"bitflags",
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.61.2",
]
[[package]]
name = "rustls"
version = "0.21.12"
@ -3096,6 +3172,12 @@ dependencies = [
"digest 0.10.7",
]
[[package]]
name = "sha1_smol"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d"
[[package]]
name = "sha2"
version = "0.10.9"
@ -3299,6 +3381,19 @@ version = "0.13.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
[[package]]
name = "tempfile"
version = "3.27.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.3",
"once_cell",
"rustix",
"windows-sys 0.61.2",
]
[[package]]
name = "thiserror"
version = "1.0.69"
@ -4182,6 +4277,12 @@ version = "0.13.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4"
[[package]]
name = "xxhash-rust"
version = "0.8.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6"
[[package]]
name = "yoke"
version = "0.8.3"

View file

@ -0,0 +1,15 @@
[package]
name = "litellm-cache-redis"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
redis = "1.7.0"
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
redis-test = "1.0.4"

View file

@ -0,0 +1,315 @@
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
Error,
};
use redis::Commands;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
const KEY_PREFIX: &str = "litellm-cache:";
pub struct RedisCache<C = redis::Connection> {
connection: Arc<Mutex<C>>,
default_ttl: Duration,
}
impl RedisCache<redis::Connection> {
pub fn new(url: &str, default_ttl: Option<Duration>) -> Result<Self, Error> {
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
let connection = client.get_connection().map_err(|_| Error::Unavailable)?;
Ok(Self::with_connection(connection, default_ttl))
}
}
impl<C> RedisCache<C>
where
C: redis::ConnectionLike + Send + 'static,
{
fn with_connection(connection: C, default_ttl: Option<Duration>) -> Self {
Self {
connection: Arc::new(Mutex::new(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
}
}
fn connection(&self) -> Result<MutexGuard<'_, C>, Error> {
self.connection.lock().map_err(|_| Error::Unavailable)
}
fn namespaced_key(key: &str) -> String {
format!("{KEY_PREFIX}{key}")
}
fn namespaced_pattern() -> &'static str {
const PATTERN: &str = "litellm-cache:*";
PATTERN
}
fn encode(value: &CacheEntry) -> Result<Vec<u8>, Error> {
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
}
fn decode(value: Vec<u8>) -> Result<CacheEntry, Error> {
serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry)
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
fn run_blocking<T, F>(connection: Arc<Mutex<C>>, operation: F) -> CacheFuture<'static, T>
where
T: Send + 'static,
F: FnOnce(&mut C) -> Result<T, Error> + Send + 'static,
{
Box::pin(async move {
tokio::task::spawn_blocking(move || {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
operation(&mut connection)
})
.await
.map_err(|_| Error::Unavailable)?
})
}
}
impl<C> BaseCache for RedisCache<C>
where
C: redis::ConnectionLike + Send + 'static,
{
type Value = CacheEntry;
fn default_ttl(&self) -> Duration {
self.default_ttl
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
let payload = Self::encode(&value)?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
self.connection()?
.set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl)
.map_err(|_| Error::Unavailable)
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
self.connection()?
.get::<_, Option<Vec<u8>>>(Self::namespaced_key(key))
.map_err(|_| Error::Unavailable)?
.map(Self::decode)
.transpose()
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.connection()?
.del::<_, ()>(Self::namespaced_key(key))
.map_err(|_| Error::Unavailable)
}
fn flush_cache(&self) -> Result<(), Error> {
let mut connection = self.connection()?;
let keys = connection
.scan_match(Self::namespaced_pattern())
.map_err(|_| Error::Unavailable)?
.collect::<redis::RedisResult<Vec<String>>>()
.map_err(|_| Error::Unavailable)?;
if keys.is_empty() {
return Ok(());
}
connection
.del::<_, usize>(keys)
.map(|_| ())
.map_err(|_| Error::Unavailable)
}
fn async_set_cache<'a>(
&'a self,
key: &'a str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
let payload = Self::encode(&value);
let key = Self::namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
.set_ex::<_, _, ()>(key, payload?, ttl)
.map_err(|_| Error::Unavailable)
})
}
fn async_get_cache<'a>(
&'a self,
key: &'a str,
_: &'a CacheKwargs,
) -> CacheFuture<'a, Option<Self::Value>> {
let key = Self::namespaced_key(key);
Box::pin(async move {
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
.get::<_, Option<Vec<u8>>>(key)
.map_err(|_| Error::Unavailable)
})
.await?
.map(Self::decode)
.transpose()
})
}
fn async_set_cache_pipeline<'a>(
&'a self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
let entries = cache_list
.into_iter()
.map(|(key, value)| {
Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload))
})
.collect::<Result<Vec<_>, _>>();
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
for (key, payload) in entries? {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)?;
}
Ok(())
})
}
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
let key = Self::namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
Box::pin(async move {
Self::run_blocking(Arc::clone(&self.connection), |connection| {
redis::cmd("PING")
.query::<String>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Redis cache connection test successful".into(),
error: None,
})
})
}
}
#[cfg(test)]
mod tests {
use super::RedisCache;
use litellm_cache::{BaseCache, CacheEntry, CacheKwargs};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
use std::time::Duration;
fn entry() -> CacheEntry {
CacheEntry {
timestamp: 123.0,
response: json!({"choices": [{"text": "cached"}]}),
}
}
#[test]
fn cache_entries_round_trip_through_json() {
let entry = entry();
let encoded = RedisCache::<redis::Connection>::encode(&entry).unwrap();
assert_eq!(
RedisCache::<redis::Connection>::decode(encoded).unwrap(),
entry
);
}
#[test]
fn invalid_json_is_rejected() {
assert!(RedisCache::<redis::Connection>::decode(b"not json".to_vec()).is_err());
}
#[test]
fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::ZERO),
1
);
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_millis(1500)),
2
);
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_secs(15)),
15
);
}
#[test]
fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
let value = entry();
let payload = RedisCache::<redis::Connection>::encode(&value).unwrap();
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SETEX")
.arg("litellm-cache:key")
.arg(600)
.arg(payload.clone()),
Ok("OK"),
),
MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)),
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
cache
.set_cache("key", value.clone(), CacheKwargs::default())
.unwrap();
assert_eq!(
cache.get_cache("key", &CacheKwargs::default()).unwrap(),
Some(value)
);
cache.delete_cache("key").unwrap();
}
#[test]
fn flush_scans_and_deletes_only_cache_keys() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SCAN")
.cursor_arg(0)
.arg("MATCH")
.arg("litellm-cache:*"),
Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])),
),
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
cache.flush_cache().unwrap();
}
#[tokio::test]
async fn test_connection_runs_ping_off_executor() {
let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
assert_eq!(
cache.test_connection().await.unwrap().status,
litellm_cache::CacheConnectionStatus::Success
);
}
}

View file

@ -0,0 +1,3 @@
mod cache;
pub use cache::RedisCache;

View file

@ -0,0 +1,6 @@
use litellm_cache_redis::RedisCache;
#[test]
fn constructor_rejects_invalid_urls() {
assert!(RedisCache::new("not a redis url", None).is_err());
}

View file

@ -1818,6 +1818,9 @@ if TYPE_CHECKING:
from .llms.azure.responses.o_series_transformation import (
AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig,
)
from .llms.azure_ai.responses.transformation import (
AzureAIResponsesAPIConfig as AzureAIResponsesAPIConfig,
)
from .llms.xai.responses.transformation import (
XAIResponsesAPIConfig as XAIResponsesAPIConfig,
)

View file

@ -234,6 +234,7 @@ LLM_CONFIG_NAMES: Final = (
"OpenAIResponsesAPIConfig",
"AzureOpenAIResponsesAPIConfig",
"AzureOpenAIOSeriesResponsesAPIConfig",
"AzureAIResponsesAPIConfig",
"XAIResponsesAPIConfig",
"LiteLLMProxyResponsesAPIConfig",
"HostedVLLMResponsesAPIConfig",
@ -946,6 +947,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.azure.responses.o_series_transformation",
"AzureOpenAIOSeriesResponsesAPIConfig",
),
"AzureAIResponsesAPIConfig": (
".llms.azure_ai.responses.transformation",
"AzureAIResponsesAPIConfig",
),
"XAIResponsesAPIConfig": (
".llms.xai.responses.transformation",
"XAIResponsesAPIConfig",

View file

@ -346,13 +346,17 @@ class MCPClient:
self.update_auth_value(auth_value)
async def discovery_auth_fingerprint(self) -> str:
return self._hash_discovery_auth(await self.prepare_request_auth())
async def prepare_request_auth(self) -> httpx.Request:
"""Preview the authenticated request without sending it, closing the auth flow afterwards."""
request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers())
if self._resolved_auth is None:
return self._hash_discovery_auth(request)
return request
flow: Final = self._resolved_auth.async_auth_flow(request)
try:
authenticated: Final = await flow.__anext__()
return self._hash_discovery_auth(authenticated)
return authenticated
finally:
await flow.aclose()

View file

@ -148,6 +148,7 @@ class _ToolCallChunk(TypedDict):
class _UsageBearingChunk(TypedDict, total=False):
usage: Usage | None
_hidden_params: Mapping[str, str]
choices: ReadOnly[Sequence[StreamingChoices | Mapping[str, object]]]
class _UsageSummary(TypedDict):
@ -921,21 +922,22 @@ class ChunkProcessor:
prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details)
completion_tokens = self._reset_anthropic_cursor_completion_tokens(
recovered_completion_tokens: Final = self._reset_anthropic_cursor_completion_tokens(
chunks=chunks,
completion_tokens=completion_tokens,
completion_usage_updates=completion_usage_updates,
)
cursor_was_reset: Final = recovered_completion_tokens != completion_tokens
return UsagePerChunk(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
completion_tokens=recovered_completion_tokens,
cache_creation_input_tokens=cache_creation_input_tokens,
cache_read_input_tokens=cache_read_input_tokens,
server_tool_use=server_tool_use,
web_search_requests=web_search_requests,
google_maps_grounding_requests=google_maps_grounding_requests,
completion_tokens_details=completion_tokens_details,
completion_tokens_details=None if cursor_was_reset else completion_tokens_details,
prompt_tokens_details=prompt_tokens_details,
cost=cost,
inference_geo=self._last_provider_pricing_field(chunks, "inference_geo"),
@ -960,6 +962,30 @@ class ChunkProcessor:
]
return values[-1] if values else None
@staticmethod
def _finish_reason_of_choice(choice: object) -> str | None:
match choice:
case StreamingChoices(finish_reason=reason) | Choices(finish_reason=reason):
return reason
case {"finish_reason": str() as reason}:
return reason
case _:
return None
@staticmethod
def _chunk_choices(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Sequence[object]:
if isinstance(chunk, dict):
return chunk.get("choices", ())
return getattr(chunk, "choices", ())
@staticmethod
def _saw_finish_reason(chunks: Sequence["_UsageBearingChunk | ModelResponse"]) -> bool:
return any(
ChunkProcessor._finish_reason_of_choice(choice) is not None
for chunk in chunks
for choice in ChunkProcessor._chunk_choices(chunk)
)
@staticmethod
def _reset_anthropic_cursor_completion_tokens(
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
@ -970,18 +996,18 @@ class ChunkProcessor:
See the ``completion_usage_updates`` comment in
``_calculate_usage_per_chunk``. The accumulated value is NOT a stale
cursor when either it is > 1 (definitely not a placeholder) or we saw
>= 2 completion-bearing usage events (positive evidence ``message_delta``
arrived). Otherwise the only completion update we ever saw was the
Anthropic ``message_start`` cursor (=1) reset to 0 so
``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates
from the actually-received completion text instead of trusting the
placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the
heuristic (which encodes Anthropic's specific message_start SSE shape)
does not silently affect other providers that may legitimately report
``completion_tokens=1`` from a single usage event.
cursor when we saw >= 2 completion-bearing usage events or any chunk
carried a ``finish_reason`` (positive evidence ``message_delta``
arrived). Otherwise the only completion update we ever saw was the
Anthropic ``message_start`` cursor, a small placeholder whose magnitude
varies per request (1 and 8 both observed live), so reset to 0 and let
``calculate_usage()``'s ``or token_counter(...)`` fallback estimate from
the actually-received text and reasoning instead. Gated on
``custom_llm_provider == "anthropic"`` so the heuristic (which encodes
Anthropic's specific message_start SSE shape) does not silently affect
other providers that legitimately report usage from a single event.
"""
saw_non_cursor_completion: Final = completion_tokens > 1 or completion_usage_updates >= 2
saw_non_cursor_completion: Final = completion_usage_updates >= 2 or ChunkProcessor._saw_finish_reason(chunks)
if saw_non_cursor_completion:
return completion_tokens
@ -995,7 +1021,7 @@ class ChunkProcessor:
if isinstance(hp, dict):
custom_llm_provider = hp.get("custom_llm_provider")
if custom_llm_provider == "anthropic" and completion_tokens == 1:
if custom_llm_provider == "anthropic":
return 0
return completion_tokens
@ -1039,10 +1065,13 @@ class ChunkProcessor:
returned_usage.prompt_tokens = 0
returned_usage.completion_tokens = (
completion_tokens
or token_counter(
model=model,
text=completion_output,
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
or (
token_counter(
model=model,
text=completion_output,
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
)
+ (reasoning_tokens or 0)
)
)
returned_usage.total_tokens = returned_usage.prompt_tokens + returned_usage.completion_tokens
@ -1066,15 +1095,16 @@ class ChunkProcessor:
returned_usage.completion_tokens_details = completion_tokens_details
if reasoning_tokens is not None:
capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
if returned_usage.completion_tokens_details is None:
returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper(
reasoning_tokens=reasoning_tokens
reasoning_tokens=capped_reasoning_tokens,
text_tokens=returned_usage.completion_tokens - capped_reasoning_tokens,
)
elif (
returned_usage.completion_tokens_details is not None
and returned_usage.completion_tokens_details.reasoning_tokens is None
):
capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens
if returned_usage.completion_tokens_details.text_tokens is None:
returned_usage.completion_tokens_details.text_tokens = (

View file

@ -49,12 +49,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params)
def get_stripped_model_name(self, model: str) -> str:
# if "responses/" is in the model name, remove it
if "responses/" in model:
model = model.replace("responses/", "")
if "o_series" in model:
model = model.replace("o_series/", "")
return model
return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "")
def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]:
"""

View file

@ -9,6 +9,7 @@ from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"]
AZURE_OPENAI_V1_HOST_SUFFIXES: Final = (".services.ai.azure.com", ".openai.azure.com")
def is_foundry_model_inference_base(api_base: str) -> bool:
@ -19,11 +20,13 @@ def is_foundry_model_inference_base(api_base: str) -> bool:
return "/openai/deployments" not in parsed.path
def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
def is_azure_openai_v1_host(api_base: str | None) -> bool:
host: Final = urlparse(api_base).hostname if api_base else None
if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")):
return "api-key"
return "Authorization"
return host is not None and host.endswith(AZURE_OPENAI_V1_HOST_SUFFIXES)
def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
return "api-key" if is_azure_openai_v1_host(api_base) else "Authorization"
def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None:
@ -70,6 +73,17 @@ def get_azure_ai_auth_headers(
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model"
def azure_ai_supports_native_responses(model: str | None, api_base: str | None) -> bool:
resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
if resolved_base is not None and not is_azure_openai_v1_host(resolved_base):
return False
if model is None:
return True
if "claude" in model.lower():
return False
return AzureFoundryModelInfo.get_azure_ai_route(model) == "default"
class AzureFoundryModelInfo(BaseLLMModelInfo):
"""Model info for Azure AI / Azure Foundry models."""

View file

@ -0,0 +1,53 @@
from typing import Final
import httpx
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
from litellm.llms.azure_ai.common_utils import (
AzureFoundryModelInfo,
api_key_header_for_base,
get_azure_ai_auth_headers,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
_PROJECT_PATH_PREFIX: Final = ("api", "projects")
_RESPONSES_PATH: Final = ("openai", "v1", "responses")
def _responses_url(api_base: str) -> str:
base_url: Final = httpx.URL(api_base)
segments: Final = tuple(segment for segment in base_url.path.split("/") if segment)
project_root: Final = segments[:3] if segments[:2] == _PROJECT_PATH_PREFIX else ()
return str(base_url.copy_with(path="/" + "/".join((*project_root, *_RESPONSES_PATH)), query=None))
class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.AZURE_AI
def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict:
params: Final = litellm_params or GenericLiteLLMParams()
auth_headers: Final = get_azure_ai_auth_headers(
api_key=AzureFoundryModelInfo.get_api_key(params.api_key),
litellm_params=params.model_dump(),
api_key_header=api_key_header_for_base(AzureFoundryModelInfo.get_api_base(params.api_base)),
)
return { # mutable-ok: the handler updates the returned headers in place per the dict contract
**headers,
**auth_headers,
"Content-Type": "application/json",
}
def supports_native_websocket(self) -> bool:
return False
def get_complete_url(self, api_base: str | None, litellm_params: dict) -> str:
resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
if resolved_base is None:
raise ValueError(
"api_base is required for the Azure AI Foundry Responses API. "
"Set the api_base parameter or the AZURE_AI_API_BASE environment variable."
)
return _responses_url(resolved_base)

View file

@ -12,6 +12,12 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig
class DashScopeChatConfig(OpenAIGPTConfig):
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: base class contract returns a list
return [ # mutable-ok: base class contract returns a list
*super().get_supported_openai_params(model=model),
"reasoning_effort",
]
def remove_cache_control_flag_from_messages_and_tools(
self,
model: str,

View file

@ -4,9 +4,9 @@ from typing import Final
import litellm
from litellm.utils import (
_is_explicitly_disabled_factory,
_supports_factory,
declared_value_factory,
is_explicitly_disabled_factory,
)
from .gpt_transformation import OpenAIGPTConfig
@ -192,7 +192,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
Use this for opt-out checks where unknown models should be allowed through.
"""
return _is_explicitly_disabled_factory(
return is_explicitly_disabled_factory(
model=cls._model_map_lookup_name(model),
custom_llm_provider=None,
key=f"supports_{level}_reasoning_effort",

View file

@ -79,6 +79,7 @@ from litellm.utils import (
CustomStreamWrapper,
ModelResponse,
is_base64_encoded,
is_explicitly_disabled_factory,
supports_reasoning,
)
@ -866,6 +867,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
else:
raise _unsupported_reasoning_effort(reasoning_effort)
@staticmethod
def _supports_minimal_thinking_level(model: str) -> bool:
lowered: Final = model.lower()
is_gemini3flash: Final = "gemini-3" in lowered and "flash" in lowered
return is_gemini3flash and not is_explicitly_disabled_factory(
model=model, custom_llm_provider=None, key="supports_minimal_reasoning_effort"
)
@staticmethod
def _map_reasoning_effort_to_thinking_level(
reasoning_effort: str,
@ -880,13 +889,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
Returns:
GeminiThinkingConfig with thinkingLevel and includeThoughts
"""
# Check if this is gemini-3-flash which supports MINIMAL thinking level
# Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview,
# gemini-3.5-flash, and any future 3.x-flash variants.
is_gemini3flash: Final = model and ("flash" in model.lower() and "gemini-3" in model.lower())
supports_minimal: Final = bool(model) and VertexGeminiConfig._supports_minimal_thinking_level(model)
is_gemini31pro: Final = model and ("gemini-3.1-pro-preview" in model.lower())
if reasoning_effort == "minimal":
if is_gemini3flash:
if supports_minimal:
return {"thinkingLevel": "minimal", "includeThoughts": True}
else:
return {"thinkingLevel": "low", "includeThoughts": True}
@ -899,18 +906,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return {"thinkingLevel": "high", "includeThoughts": True}
elif reasoning_effort == "high":
return {"thinkingLevel": "high", "includeThoughts": True}
elif reasoning_effort == "disable":
# Gemini 3 cannot fully disable thinking, so we use "minimal" for gemini-3-flash-preview, "low" for others
if is_gemini3flash:
return {"thinkingLevel": "minimal", "includeThoughts": False}
else:
return {"thinkingLevel": "low", "includeThoughts": False}
elif reasoning_effort == "none":
# For gemini-3-flash-preview, use "minimal" instead of "low"
if is_gemini3flash:
return {"thinkingLevel": "minimal", "includeThoughts": False}
else:
return {"thinkingLevel": "low", "includeThoughts": False}
elif reasoning_effort in ("disable", "none"):
return {
"thinkingLevel": "minimal" if supports_minimal else "low",
"includeThoughts": False,
}
else:
raise _unsupported_reasoning_effort(reasoning_effort)
@ -977,8 +977,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
params["includeThoughts"] = True
# Follow provider defaults unless explicitly opted into legacy behavior.
if litellm.enable_gemini_default_thinking_level_low is True:
is_gemini3flash: Final = "gemini-3" in model.lower() and "flash" in model.lower()
params["thinkingLevel"] = "minimal" if is_gemini3flash else "low"
params["thinkingLevel"] = (
"minimal" if VertexGeminiConfig._supports_minimal_thinking_level(model) else "low"
)
else:
# Thinking disabled
params["includeThoughts"] = False

View file

@ -26105,6 +26105,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -26162,6 +26163,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28111,6 +28113,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28170,6 +28173,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28592,6 +28596,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28649,6 +28654,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,

View file

@ -102,6 +102,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import (
UpstreamCredentialProvider,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
prepare_mcp_client,
raise_public,
raise_token_exchange_challenge,
raise_user_oauth_challenge,
@ -2804,6 +2805,8 @@ class MCPServerManager:
headers=headers,
server_label=server.name or server.server_name or server.alias or server.server_id,
relays_upstream_auth=server.is_client_forwarded_token,
auth_type=server.auth_type,
upstream_token_header=server.upstream_token_header,
)
tool_func.__name__ = prefixed_tool_name
tool_func.__doc__ = description
@ -4259,15 +4262,20 @@ class MCPServerManager:
user_api_key_auth=user_api_key_auth,
extra_headers=extra_headers,
)
return MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
extra_headers=extra_headers,
resolved_auth=resolved_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
return await prepare_mcp_client(
resolved_server,
MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
timeout=(
resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT
),
extra_headers=extra_headers,
resolved_auth=resolved_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
),
)
# Create SigV4 auth if configured
@ -4297,17 +4305,20 @@ class MCPServerManager:
else AuthResolution.no_auth
)
record_auth_resolution(server.server_id, legacy_source)
return MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
auth_value=auth_value,
auth_header_name=auth_header_name,
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
extra_headers=extra_headers,
aws_auth=aws_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
return await prepare_mcp_client(
resolved_server,
MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
auth_value=auth_value,
auth_header_name=auth_header_name,
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
extra_headers=extra_headers,
aws_auth=aws_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
),
)
async def _get_tools_from_server(

View file

@ -54,7 +54,7 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.proxy._experimental.mcp_server.tool_registry import (
global_mcp_tool_registry,
)
from litellm.types.mcp import credential_redirect_hook, custom_credential_slot
from litellm.types.mcp import MCPAuthType, credential_redirect_hook, custom_credential_slot
class _OpenAPIJSONSchema(TypedDict, total=False):
@ -471,6 +471,8 @@ def create_tool_function(
headers: dict[str, str] | None = None,
server_label: str | None = None,
relays_upstream_auth: bool = False,
auth_type: MCPAuthType = None,
upstream_token_header: str | None = None,
):
"""Create a tool function for an OpenAPI operation.
@ -503,6 +505,18 @@ def create_tool_function(
by using **kwargs instead of named parameters.
"""
effective_headers: Final = _merge_openapi_tool_request_headers(headers)
if auth_type is not None:
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
raise_public,
validate_static_credential,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok
match validate_static_credential(auth_type, effective_headers, upstream_token_header):
case Error(error):
raise_public(error)
case Ok():
pass
# Build URL from base_url and path
url = base_url + path

View file

@ -13,15 +13,17 @@ from __future__ import annotations
import base64
import os
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Literal, NoReturn
from fastapi import HTTPException
from pydantic import SecretStr
from typing_extensions import assert_never
from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials
from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme, to_basic_credentials
from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError
from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
DEFAULT_CREDENTIAL_HEADER,
ApiKeyConfig,
@ -39,7 +41,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
Subject,
TokenExchangeConfig,
)
from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth
from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth, MCPAuthType, MCPTransport
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
@ -79,7 +81,7 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
BYOK is the per-user source of the ``api_key`` mode; its scheme rides on ``auth_type`` just
like a shared key, but the value is per-user and not migrated yet, so a BYOK server defers
to v1 regardless of ``auth_type`` (this guard is the seam the BYOK arm replaces later).
to v1 for its static schemes. Declared OBO always stays with the exchange arm.
Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with
an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is
@ -90,8 +92,8 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
modes ``true_passthrough`` / ``oauth_delegate`` (``PassthroughConfig``); delegated/passthrough
oauth2 and SigV4 return None and stay on v1.
"""
if server.is_byok:
return None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type)
if server.is_byok and server.auth_type != MCPAuth.oauth2_token_exchange:
return None # per-user BYOK source not migrated yet -> defer to v1
resource: Final = server.url or server.server_id
auth_type: Final = server.auth_type
match auth_type:
@ -165,21 +167,9 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec:
)
def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
"""Build a token_exchange (OBO) spec, or defer (None) when it is not OBO-configured.
An OBO server with ``client_id``/``client_secret`` is owned by the v2 arm even if the
``token_exchange_endpoint``/``token_url`` is absent: a missing endpoint then fails closed (412) at
the exchanger rather than silently deferring to v1 and connecting unauthenticated, since the
gateway must not guess the IdP or fall back to a weaker source. Without client credentials there is
nothing to own, so the server stays on v1 (parity-safe). ``profile`` selects the wire dialect
(``rfc8693`` default, ``entra_obo`` for Microsoft Entra On-Behalf-Of); an unrecognized value
normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is
forwarded only when the operator set it; a missing one is omitted, not derived.
"""
def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec:
"""Keep declared OBO owned by the resolver, including incomplete client configuration."""
endpoint: Final = server.token_exchange_endpoint or server.effective_token_url
if not server.client_id or not server.client_secret:
return None
profile: Final[Literal["rfc8693", "entra_obo"]] = (
"entra_obo" if server.token_exchange_profile == "entra_obo" else "rfc8693"
)
@ -193,7 +183,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
token_exchange_endpoint=endpoint,
audience=server.audience,
client_id=server.client_id,
client_secret=SecretStr(server.client_secret),
client_secret=SecretStr(server.client_secret) if server.client_secret else None,
token_endpoint_auth_method=server.token_endpoint_auth_method,
scopes=tuple(server.scopes or ()),
),
@ -397,3 +387,69 @@ def raise_token_exchange_challenge(
detail="Unauthorized",
headers={"WWW-Authenticate": www_authenticate},
)
_STATIC_MODES: Final = frozenset(
(MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization)
)
def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool:
if not value:
return False
if auth_type == MCPAuth.api_key and name != "authorization":
return True
if value.lower() in ("bearer", "basic", "token", "apikey"):
return False
if auth_type == MCPAuth.api_key:
api_scheme: Final = value.split(None, 1)[0]
if api_scheme.lower() in ("bearer", "token", "apikey"):
api_credential: Final = strip_auth_scheme(value, api_scheme).strip()
return api_credential.lower() != api_scheme.lower()
if auth_type in (MCPAuth.bearer_token, MCPAuth.token):
scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token"
credential: Final = strip_auth_scheme(value, scheme).strip()
return bool(credential) and credential.lower() != scheme.lower()
if auth_type == MCPAuth.basic:
parts: Final = value.split(None, 1)
if len(parts) != 2 or parts[0].lower() != "basic":
return False
try:
decoded: Final = base64.b64decode(parts[1], validate=True).strip()
return b":" in decoded
except ValueError:
return False
return True
def validate_static_credential(
auth_type: MCPAuthType,
headers: Mapping[str, str],
upstream_token_header: str | None = None,
) -> Result[None, CredError]:
if auth_type not in _STATIC_MODES:
return Ok(None)
default_slot: Final = "X-API-Key" if auth_type == MCPAuth.api_key else "Authorization"
slots: Final = frozenset(
name.lower()
for name in (
upstream_token_header or default_slot,
default_slot,
"Authorization",
)
)
values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots)
if any(_usable_credential_value(auth_type, name, value) for name, value in values):
return Ok(None)
return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential"))
async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient:
if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio:
return client
request: Final = await client.prepare_request_auth()
match validate_static_credential(server.auth_type, request.headers, server.upstream_token_header):
case Error(error):
raise_public(error)
case Ok():
return client

View file

@ -844,6 +844,9 @@ class LiteLLMRoutes(enum.Enum):
)
self_managed_routes = [
# update_team resolves proxy/org/team admin itself and filters team admins
# through the team_admin_editable_team_fields setting
"/team/update",
"/team/member_add",
"/team/member_delete",
"/management/v1/teams/{team_id}/members/bulk_delete",
@ -4467,6 +4470,29 @@ class TeamInfoMember(Member):
user_alias: str | None = None
class TeamEditUnrestricted(BaseModel):
kind: Literal["unrestricted"] = "unrestricted"
class TeamEditAsTeamAdmin(BaseModel):
kind: Literal["team_admin"] = "team_admin"
editable_fields: tuple[str, ...]
class TeamEditAsTeamAdminDisabled(BaseModel):
kind: Literal["team_admin_disabled"] = "team_admin_disabled"
class TeamEditNone(BaseModel):
kind: Literal["none"] = "none"
TeamEditAccess = Annotated[
TeamEditUnrestricted | TeamEditAsTeamAdmin | TeamEditAsTeamAdminDisabled | TeamEditNone,
Field(discriminator="kind"),
]
class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
members_with_roles: tuple[TeamInfoMember, ...] = ()
team_member_budget_table: LiteLLM_BudgetTableFull | None = None
@ -4479,6 +4505,7 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
# None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling.
organization_models: list[str] | None = None
model_max_budget_usage: Mapping[str, Mapping[str, object]] | None = None
caller_edit_access: TeamEditAccess = Field(default_factory=TeamEditNone)
class TeamInfoResponseObject(TypedDict):

View file

@ -8,7 +8,6 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.llms.anthropic.experimental_pass_through.context_management import (
@ -22,13 +21,16 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
create_response,
log_llm_api_exception,
proxy_exception_from_http_exception,
resolve_litellm_call_id,
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
openai_error_param,
openai_error_type,
with_litellm_call_id,
)
from litellm.types.utils import TokenCountResponse
@ -218,10 +220,12 @@ async def anthropic_response(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=base_llm_response_processor.data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e)
log_llm_api_exception(e, base_llm_response_processor.litellm_call_id)
if isinstance(e, ProxyException):
return _anthropic_error_json_response(e, request)
return _anthropic_error_json_response(
with_litellm_call_id(e, base_llm_response_processor.litellm_call_id), request
)
# Extract model_id from request metadata (same as success path)
litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
@ -231,7 +235,7 @@ async def anthropic_response(
# Get headers
headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=data.get("litellm_call_id", ""),
call_id=base_llm_response_processor.litellm_call_id,
model_id=model_id,
version=version,
response_cost=0,
@ -288,6 +292,7 @@ async def count_tokens(
"""
from litellm.proxy.proxy_server import token_counter as internal_token_counter
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
try:
request_data: Final = await _read_request_body(request=request)
data: Final[dict] = {**request_data}
@ -339,7 +344,7 @@ async def count_tokens(
detail=detail,
)
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - %s", e)
log_llm_api_exception(e, litellm_call_id)
raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"})

View file

@ -7,6 +7,7 @@
import asyncio
import os
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final, cast
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
@ -17,7 +18,11 @@ from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
request_litellm_call_id,
)
from litellm.proxy.common_utils.callback_utils import sanitize_openai_provider_metadata
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.openai_endpoint_utils import (
@ -383,8 +388,9 @@ async def create_batch(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
@router.get(
@ -674,8 +680,9 @@ async def retrieve_batch(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
@router.get(
@ -725,6 +732,7 @@ async def list_batches(
)
verbose_proxy_logger.debug("GET /v1/batches after=%s limit=%s", after, limit)
data: Mapping[str, object] = MappingProxyType({})
try:
if llm_router is None:
raise HTTPException(
@ -854,10 +862,11 @@ async def list_batches(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data={"after": after, "limit": limit},
request_data={**data, "after": after, "limit": limit},
)
verbose_proxy_logger.error("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
@router.post(
@ -1079,8 +1088,9 @@ async def cancel_batch(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
######################################################################

View file

@ -7,7 +7,18 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequen
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, TypeAlias, TypeVar, overload
from typing import (
TYPE_CHECKING,
Any,
Final,
Literal,
NamedTuple,
Protocol,
TypeAlias,
TypeVar,
overload,
runtime_checkable,
)
import anyio
import httpx
@ -34,7 +45,11 @@ from litellm.constants import (
UNSAFE_PROXY_RESPONSE_HEADERS,
)
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error
from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
independent_snapshot,
is_expected_client_error,
)
from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
from litellm.litellm_core_utils.get_supported_openai_params import (
get_supported_openai_params,
@ -1452,7 +1467,19 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool:
_CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request"
def _log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
@runtime_checkable
class _CarriesLitellmCallId(Protocol):
litellm_call_id: str | None
def request_litellm_call_id(data: Mapping[str, object]) -> str | None:
logging_obj: Final = data.get("litellm_logging_obj")
logged_id: Final = logging_obj.litellm_call_id if isinstance(logging_obj, _CarriesLitellmCallId) else None
call_id: Final = logged_id or data.get("litellm_call_id")
return call_id if isinstance(call_id, str) else None
def log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
if getattr(e, "status_code", None) == 499 and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL:
verbose_proxy_logger.info(
"litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, "
@ -1532,6 +1559,10 @@ class ProxyBaseLLMRequestProcessing:
def __init__(self, data: dict):
self.data = data
@property
def litellm_call_id(self) -> str | None:
return request_litellm_call_id(self.data)
@staticmethod
def _merge_passthrough_streaming_headers(
response_headers: httpx.Headers | dict | None,
@ -2062,6 +2093,13 @@ class ProxyBaseLLMRequestProcessing:
) -> tuple[dict, LiteLLMLoggingObj]:
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
configured_fallbacks: Final = (
self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict)
if llm_router is not None and not self.data.get("disable_fallbacks")
else None
)
pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None
try:
return await self.common_processing_pre_call_logic(
request=request,
@ -2080,14 +2118,19 @@ class ProxyBaseLLMRequestProcessing:
llm_router=llm_router,
)
except ProxyRateLimitError as original_exc:
original_model: Final = self.data.get("model")
if not original_model or not llm_router or self.data.get("disable_fallbacks"):
rate_limited_data: Final = self.data
original_model: Final = rate_limited_data.get("model")
if (
pristine is None
or not configured_fallbacks
or rate_limited_data.get("disable_fallbacks")
or not isinstance(original_model, str)
):
raise
fallback_models: Final = self._resolve_fallback_models(
model=original_model,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
fallbacks=configured_fallbacks,
)
if not fallback_models:
raise
@ -2102,6 +2145,7 @@ class ProxyBaseLLMRequestProcessing:
for fallback_model in fallback_models:
if fallback_model == original_model:
continue
self.data = independent_snapshot(pristine)
self.data["model"] = fallback_model
try:
return await self.common_processing_pre_call_logic(
@ -2123,39 +2167,30 @@ class ProxyBaseLLMRequestProcessing:
except ProxyRateLimitError:
continue
except BaseException:
self.data["model"] = original_model
self.data = rate_limited_data
raise
self.data["model"] = original_model
self.data = rate_limited_data
raise original_exc
def _resolve_fallback_models(
self,
model: str,
llm_router: Router,
user_api_key_dict: UserAPIKeyAuth,
) -> list | None:
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
fallbacks = None
@staticmethod
def _configured_fallbacks(llm_router: Router, user_api_key_dict: UserAPIKeyAuth) -> list | None:
key_router_settings: Final = user_api_key_dict.router_settings
if isinstance(key_router_settings, dict) and "fallbacks" in key_router_settings:
fallbacks = key_router_settings["fallbacks"]
key_fallbacks: Final = key_router_settings.get("fallbacks") if isinstance(key_router_settings, dict) else None
fallbacks: Final = key_fallbacks if key_fallbacks is not None else llm_router.fallbacks
return fallbacks if isinstance(fallbacks, list) and fallbacks else None
if fallbacks is None:
fallbacks = llm_router.fallbacks
if not fallbacks:
return None
@staticmethod
def _resolve_fallback_models(model: str, fallbacks: list) -> list | None:
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
fallback_model_group, generic_fallback_idx = get_fallback_model_group(
fallbacks=fallbacks,
model_group=model,
)
if fallback_model_group is None and generic_fallback_idx is not None:
fallback_model_group = fallbacks[generic_fallback_idx]["*"]
return fallback_model_group
if fallback_model_group is not None:
return fallback_model_group
return fallbacks[generic_fallback_idx]["*"] if generic_fallback_idx is not None else None
@staticmethod
def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str:
@ -3429,11 +3464,7 @@ class ProxyBaseLLMRequestProcessing:
version: str | None = None,
):
"""Raises ProxyException (OpenAI API compatible) if an exception is raised"""
logging_obj: Final[LiteLLMLoggingObj | None] = self.data.get("litellm_logging_obj", None)
_log_llm_api_exception(
e,
(logging_obj.litellm_call_id if logging_obj is not None else None) or self.data.get("litellm_call_id"),
)
log_llm_api_exception(e, self.litellm_call_id)
# Allow callbacks to transform the error response
transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
@ -3463,9 +3494,7 @@ class ProxyBaseLLMRequestProcessing:
custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=(
_litellm_logging_obj.litellm_call_id if _litellm_logging_obj else self.data.get("litellm_call_id")
),
call_id=self.litellm_call_id,
model_id=model_id,
version=version,
response_cost=0,

View file

@ -9,6 +9,9 @@ from typing import Final
from fastapi import status
from litellm.constants import STRINGIFIED_NONE
from litellm.proxy._types import ProxyException
LITELLM_CALL_ID_HEADER: Final = "x-litellm-call-id"
_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
{
@ -52,3 +55,23 @@ def openai_error_param(exc: object) -> str | None:
serializes as JSON ``null``."""
carried: Final = attribute_of(exc, "param")
return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None
def litellm_call_id_headers(litellm_call_id: str | None) -> dict[str, str] | None: # mutable-ok: ProxyException.headers
if litellm_call_id is None:
return None
return {LITELLM_CALL_ID_HEADER: litellm_call_id} # mutable-ok: ProxyException mutates its headers dict
def with_litellm_call_id(exc: ProxyException, litellm_call_id: str | None) -> ProxyException:
"""The same error object, answering with ``x-litellm-call-id`` when it was raised without one."""
if litellm_call_id is not None:
exc.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id)
return exc
def headers_with_litellm_call_id(headers: Mapping[str, str] | None, litellm_call_id: str) -> Mapping[str, str]:
"""``headers`` plus ``x-litellm-call-id``, keeping the value they already carry under that name."""
if headers is None:
return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id})
return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id, **headers})

View file

@ -1,6 +1,5 @@
import asyncio
import io
import traceback
from collections.abc import Sequence
from typing import Final, get_type_hints
@ -9,19 +8,23 @@ from fastapi import APIRouter, Depends, File, HTTPException, Request, Response,
from fastapi.responses import ORJSONResponse
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
)
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
resolve_litellm_call_id,
)
from litellm.proxy.common_utils.http_parsing_utils import (
coerce_numeric_form_fields,
numeric_form_fields,
)
from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
litellm_call_id_headers,
openai_error_param,
openai_error_type,
)
@ -91,11 +94,12 @@ async def image_generation(
version,
)
data = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data = {"litellm_call_id": litellm_call_id}
try:
# Use orjson to parse JSON data, orjson speeds up requests significantly
body: Final = await request.body()
data = orjson.loads(body)
data = orjson.loads(body) | data
# Include original request and headers in the data
data = await add_litellm_data_to_request(
@ -153,9 +157,7 @@ async def image_generation(
response = await llm_call
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
### CALL HOOKS ### - modify outgoing data (guardrails, otel, etc.)
response = await proxy_logging_obj.post_call_success_hook(
@ -168,7 +170,7 @@ async def image_generation(
cache_key: Final = hidden_params.get("cache_key", None) or ""
api_base: Final = hidden_params.get("api_base", None) or ""
response_cost: Final = hidden_params.get("response_cost", None) or ""
litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
response_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
fastapi_response.headers.update(
ProxyBaseLLMRequestProcessing.get_custom_headers(
@ -179,7 +181,7 @@ async def image_generation(
version=version,
response_cost=response_cost,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
call_id=litellm_call_id,
call_id=response_call_id,
request_data=data,
hidden_params=hidden_params,
)
@ -200,13 +202,13 @@ async def image_generation(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.error("litellm.proxy.proxy_server.image_generation(): Exception occured - %s", e)
verbose_proxy_logger.debug(traceback.format_exc())
log_llm_api_exception(e, litellm_call_id)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "message", str(e)),
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
param=openai_error_param(e),
headers=litellm_call_id_headers(litellm_call_id),
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
)
else:
@ -215,6 +217,7 @@ async def image_generation(
message=getattr(e, "message", error_msg),
type=openai_error_type(e, error_status_code(e, 500)),
param=openai_error_param(e),
headers=litellm_call_id_headers(litellm_call_id),
openai_code=getattr(e, "code", None),
code=error_status_code(e, 500),
)

View file

@ -0,0 +1,191 @@
"""Proxy-wide allow-list of team-settings fields a team admin may change on /team/update."""
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal, TypeAlias
from fastapi import HTTPException
from pydantic import TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm._logging import verbose_proxy_logger
from litellm.models.team import LiteLLM_TeamTable
from litellm.proxy._types import (
LiteLLM_ManagementEndpoint_MetadataFields,
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
UpdateTeamRequest,
)
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_fields"
# TODO(LIT-5722): add the remaining team settings one per PR, each with its value-diff tests and dashboard field
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset({"tpm_limit"})
_FIELD_LIST: Final = TypeAdapter(list[str])
_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
_EMPTY: Final[Mapping[str, object]] = MappingProxyType({})
_METADATA_FOLDED_FIELDS: Final[frozenset[str]] = frozenset(
(*LiteLLM_ManagementEndpoint_MetadataFields, *LiteLLM_ManagementEndpoint_MetadataFields_Premium)
)
_SYSTEM_MANAGED_METADATA_KEYS: Final[frozenset[str]] = frozenset({"team_member_budget_id"})
_NOT_COLUMNS: Final[frozenset[str]] = frozenset({"team_id", "metadata"})
_SETTINGS_LOCATION: Final = "Settings > UI > Team admin editable fields"
@dataclass(frozen=True, slots=True)
class TeamAdminEditAllowed:
request: UpdateTeamRequest
kind: Literal["allowed"] = "allowed"
@dataclass(frozen=True, slots=True)
class TeamAdminEditingDisabled:
kind: Literal["disabled"] = "disabled"
@dataclass(frozen=True, slots=True)
class TeamAdminFieldNotPermitted:
field: str
kind: Literal["field_not_permitted"] = "field_not_permitted"
TeamAdminEditVerdict: TypeAlias = TeamAdminEditAllowed | TeamAdminEditingDisabled | TeamAdminFieldNotPermitted
def resolve_team_admin_editable_fields(
general_settings: Mapping[str, object],
supported: frozenset[str],
) -> frozenset[str]:
raw: Final = general_settings.get(TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING)
if raw is None:
return frozenset()
try:
configured: Final = frozenset(_FIELD_LIST.validate_python(raw))
except ValidationError:
verbose_proxy_logger.warning(
"%s must be a list of field names; ignoring %r", TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, raw
)
return frozenset()
unsupported: Final = configured - supported
if unsupported:
verbose_proxy_logger.warning(
"%s ignores unsupported field(s) %s; supported: %s",
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
sorted(unsupported),
sorted(supported),
)
return configured & supported
def _as_object(value: object) -> Mapping[str, object]:
try:
return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value)
except ValidationError:
return _EMPTY
def _stored_metadata(existing: Mapping[str, object]) -> Mapping[str, object]:
return _as_object(existing.get("metadata"))
def _submitted_metadata(
data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object]
) -> Mapping[str, object]:
"""Metadata as it would be stored: the caller's dict (or the stored one) with top-level folded fields laid over."""
base: Final = (
_as_object(submitted.get("metadata")) if "metadata" in data.model_fields_set else _stored_metadata(existing)
)
folded: Final = data.model_fields_set & _METADATA_FOLDED_FIELDS
return MappingProxyType({key: submitted[key] if key in folded else base[key] for key in base.keys() | folded})
def _metadata_changes(
data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object]
) -> frozenset[str]:
merged: Final = _submitted_metadata(data, submitted, existing)
stored: Final = _stored_metadata(existing)
return frozenset(
key if key in _METADATA_FOLDED_FIELDS else "metadata"
for key in (merged.keys() | stored.keys()) - _SYSTEM_MANAGED_METADATA_KEYS
if merged.get(key) != stored.get(key)
)
def _stored_model_aliases(existing_row: LiteLLM_TeamTable) -> Mapping[str, object]:
table: Final = existing_row.litellm_model_table
return _as_object(_JSON_OBJECT.validate_json(table.model_dump_json()).get("model_aliases")) if table else _EMPTY
def _column_changed(
field: str, submitted: Mapping[str, object], existing: Mapping[str, object], existing_row: LiteLLM_TeamTable
) -> bool:
if field == "model_aliases":
return _as_object(submitted.get(field)) != _stored_model_aliases(existing_row)
if field in LiteLLM_TeamTable.model_fields:
return submitted.get(field) != existing.get(field)
return True
def changed_team_fields(data: UpdateTeamRequest, existing_row: LiteLLM_TeamTable) -> frozenset[str]:
"""Logical field names whose stored value the request would change.
Request and stored row are compared as JSON values so both sides share one representation. Fields the
server folds into metadata are attributed to their own name whether they arrive top-level or inside
``metadata``; anything else in ``metadata`` is attributed to ``metadata``. Fields with no stored
counterpart on the team row count as changed whenever they are sent.
"""
submitted: Final = _JSON_OBJECT.validate_json(data.model_dump_json(exclude_unset=True))
existing: Final = _JSON_OBJECT.validate_json(existing_row.model_dump_json())
column_fields: Final = frozenset(data.model_fields_set) - _NOT_COLUMNS - _METADATA_FOLDED_FIELDS
column_changes: Final = frozenset(
field for field in column_fields if _column_changed(field, submitted, existing, existing_row)
)
return column_changes | _metadata_changes(data, submitted, existing)
def _only_changes(data: UpdateTeamRequest, changed: frozenset[str]) -> UpdateTeamRequest:
"""The request without the values it resends unchanged, which would otherwise still trigger derived writes
such as a resent budget_duration pushing budget_reset_at back."""
sent: Final = frozenset(data.model_fields_set)
via_metadata: Final = frozenset({"metadata"}) if changed - sent else frozenset()
kept: Final = frozenset({"team_id"}) | (changed & sent) | via_metadata
return UpdateTeamRequest.model_validate(data.model_dump(include=MappingProxyType({field: True for field in kept})))
def team_admin_edit_verdict(
data: UpdateTeamRequest,
existing: LiteLLM_TeamTable,
permitted: frozenset[str],
) -> TeamAdminEditVerdict:
if not permitted:
return TeamAdminEditingDisabled()
changed: Final = changed_team_fields(data, existing)
blocked: Final = sorted(changed - permitted)
if blocked:
return TeamAdminFieldNotPermitted(field=blocked[0])
return TeamAdminEditAllowed(request=_only_changes(data, changed))
def team_admin_request_or_raise(verdict: TeamAdminEditVerdict) -> UpdateTeamRequest:
match verdict:
case TeamAdminEditAllowed(request=request):
return request
case TeamAdminEditingDisabled():
raise HTTPException(
status_code=403,
detail=(
"Team admins on this proxy cannot edit team settings. "
f"Ask a proxy admin to enable fields under {_SETTINGS_LOCATION}."
),
)
case TeamAdminFieldNotPermitted(field=field):
raise HTTPException(
status_code=403,
detail=(
f"Team admins on this proxy do not have permission to update '{field}'. "
f"Ask a proxy admin to add it under {_SETTINGS_LOCATION}."
),
)
case _:
assert_never(verdict)

View file

@ -18,12 +18,23 @@ from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protocol, TypeVar, cast
from typing import (
TYPE_CHECKING,
Annotated,
Final,
Literal,
NamedTuple,
NoReturn,
Protocol,
TypeAlias,
TypeVar,
cast,
)
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import BaseModel, JsonValue, ValidationError
from typing_extensions import ReadOnly, TypedDict
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
from litellm._logging import verbose_proxy_logger
@ -62,6 +73,11 @@ from litellm.proxy._types import (
SpecialProxyStrings,
TeamAccessGroupModelGrant,
TeamAddMemberResponse,
TeamEditAccess,
TeamEditAsTeamAdmin,
TeamEditAsTeamAdminDisabled,
TeamEditNone,
TeamEditUnrestricted,
TeamInfoMember,
TeamInfoResponseObject,
TeamInfoResponseObjectTeamTable,
@ -122,6 +138,12 @@ from litellm.proxy.management_endpoints.router_weights import validate_router_se
from litellm.proxy.management_endpoints.tag_management_endpoints import (
get_daily_activity,
)
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS,
resolve_team_admin_editable_fields,
team_admin_edit_verdict,
team_admin_request_or_raise,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
TEAM_ADVISORY_LOCK_SQL,
AccessGroupSyncTx,
@ -439,32 +461,70 @@ async def _refresh_cached_team(
)
async def _can_manage_team(
TeamAccessRole: TypeAlias = Literal["proxy_admin", "org_admin", "team_admin"]
def _raise_team_access_denied() -> NoReturn:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="You do not have access to this team",
)
async def _resolve_team_access(
team_obj: LiteLLM_TeamTable,
user_api_key_dict: UserAPIKeyAuth,
) -> bool:
"""True for a proxy admin, an admin of this team, or an org admin for the team's organization."""
) -> TeamAccessRole | None:
"""Strongest role the caller holds over ``team_obj``, or None when they hold none.
Org admin outranks team admin so a caller holding both keeps unrestricted edits.
"""
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return True
return "proxy_admin"
if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
return "org_admin"
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
return True
return "team_admin"
return await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj)
return None
async def _verify_team_access(
team_obj: LiteLLM_TeamTable,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""Raise HTTPException(403) unless the caller can manage the given team."""
if await _can_manage_team(team_obj=team_obj, user_api_key_dict=user_api_key_dict):
return
"""Raise 403 unless the caller is a proxy admin, an org admin for the team's org, or a team admin."""
if await _resolve_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict) is None:
_raise_team_access_denied()
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="You do not have access to this team",
)
_GENERAL_SETTINGS: Final = TypeAdapter(dict[str, object])
def _general_settings() -> Mapping[str, object]:
from litellm.proxy.proxy_server import general_settings
return _GENERAL_SETTINGS.validate_python(general_settings)
def _caller_edit_access(role: TeamAccessRole | None, general_settings: Mapping[str, object]) -> TeamEditAccess:
"""What the caller may change on /team/update, reported on /team/info so the dashboard never re-derives it."""
match role:
case "proxy_admin" | "org_admin":
return TeamEditUnrestricted()
case "team_admin":
permitted: Final = resolve_team_admin_editable_fields(
general_settings, SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
)
if not permitted:
return TeamEditAsTeamAdminDisabled()
return TeamEditAsTeamAdmin(editable_fields=tuple(sorted(permitted)))
case None:
return TeamEditNone()
case _:
assert_never(role)
class TeamMemberBudgetHandler:
@ -2144,16 +2204,29 @@ async def update_team(
)
if existing_team_row is None:
# Non-proxy-admins get the same 403 as an access denial so /team/update
# cannot be used to probe which team ids exist
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
_raise_team_access_denied()
raise HTTPException(
status_code=404,
detail={"error": f"Team not found, passed team_id={data.team_id}"},
)
# Verify caller has access to manage this team
await _verify_team_access(
team_obj=LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()),
user_api_key_dict=user_api_key_dict,
)
existing_team: Final = LiteLLM_TeamTable.model_validate(existing_team_row.model_dump())
access_role: Final = await _resolve_team_access(team_obj=existing_team, user_api_key_dict=user_api_key_dict)
if access_role is None:
_raise_team_access_denied()
if access_role == "team_admin":
data = team_admin_request_or_raise( # rebind-ok: resent values must not reach the derived writes below
team_admin_edit_verdict(
data=data,
existing=existing_team,
permitted=resolve_team_admin_editable_fields(
_general_settings(), SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
),
)
)
await validate_router_settings_weights(
data.router_settings,
@ -2257,6 +2330,7 @@ async def update_team(
org_id=org_id_to_check,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
include_budget_table=True,
)
if org_table is not None:
await _check_org_team_limits(
@ -4583,10 +4657,9 @@ async def team_info(
)
team_table: Final = LiteLLM_TeamTable.model_validate(team_info.model_dump())
await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_table)
access_role: Final = await _resolve_team_access(team_obj=team_table, user_api_key_dict=user_api_key_dict)
organization_models: Final[list[str] | None] = (
_parent_organization_models(team_info)
if await _can_manage_team(team_obj=team_table, user_api_key_dict=user_api_key_dict)
else None
_parent_organization_models(team_info) if access_role is not None else None
)
## GET ALL KEYS ##
@ -4655,6 +4728,7 @@ async def team_info(
model_max_budget=resolved_team_info.model_max_budget,
cache=model_max_budget_limiter.dual_cache,
),
"caller_edit_access": _caller_edit_access(access_role, _general_settings()),
}
)

View file

@ -1180,9 +1180,8 @@ async def bedrock_proxy_route(
endpoint_func: Final = create_pass_through_route(
endpoint=endpoint,
target=str(prepped.url),
custom_headers=prepped.headers,
custom_headers=_upstream_headers_for_bedrock_agent_runtime_route(request, user_api_key_dict, prepped.headers),
is_streaming_request=is_streaming_request,
_forward_headers=True,
) # dynamically construct pass-through endpoint based on incoming path
setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, data)
# SigV4 signs an exact payload; pass-through must send prepped.body, not json.dumps
@ -2001,6 +2000,9 @@ _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS: Final = frozenset({"authorization", "x-a
_HEADERS_NEVER_FORWARDED_TO_ANTHROPIC: Final = frozenset({"content-length", "host", "accept-encoding"}) | (
SpecialHeaders.litellm_credential_header_names() - _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS
)
_HEADERS_NEVER_FORWARDED_TO_BEDROCK: Final = (
frozenset({"content-length", "host", "accept-encoding"}) | SpecialHeaders.litellm_credential_header_names()
)
_MAPPED_ROUTE_CALLER_KEY_HEADER: Final = "litellm_user_api_key"
@ -2099,6 +2101,17 @@ def _upstream_headers_for_anthropic_route(
return MappingProxyType({**caller_headers, **(proxy_auth_header or {})})
def _upstream_headers_for_bedrock_agent_runtime_route(
request: Request, user_api_key_dict: UserAPIKeyAuth, signed_headers: Mapping[str, object]
) -> Mapping[str, object]:
caller_headers: Final = _caller_headers_without_litellm_secrets(
request,
user_api_key_dict,
_HEADERS_NEVER_FORWARDED_TO_BEDROCK | frozenset(name.lower() for name in signed_headers),
)
return MappingProxyType({**caller_headers, **signed_headers})
async def _prepare_vertex_auth_headers(
request: Request,
vertex_credentials: VertexPassThroughCredentials | None,

View file

@ -72,7 +72,9 @@ from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_end
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
open_sse_before_first_byte,
resolve_litellm_call_id,
)
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
@ -80,6 +82,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
)
from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
litellm_call_id_headers,
openai_error_param,
openai_error_type,
)
@ -196,14 +199,15 @@ async def chat_completion_pass_through_endpoint(
version,
)
data = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data = {"litellm_call_id": litellm_call_id}
try:
body: Final = await request.body()
body_str: Final = body.decode()
try:
data = ast.literal_eval(body_str)
data = ast.literal_eval(body_str) | data
except Exception:
data = json.loads(body_str)
data = json.loads(body_str) | data
data["adapter_id"] = adapter_id
@ -290,9 +294,7 @@ async def chat_completion_pass_through_endpoint(
response_cost: Final = hidden_params.get("response_cost", None) or ""
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
verbose_proxy_logger.debug("final response: %s", response)
@ -313,12 +315,13 @@ async def chat_completion_pass_through_endpoint(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - %s", e)
log_llm_api_exception(e, litellm_call_id)
error_msg: Final = f"{e}"
raise ProxyException(
message=getattr(e, "message", error_msg),
type=openai_error_type(e, error_status_code(e, 500)),
param=openai_error_param(e),
headers=litellm_call_id_headers(litellm_call_id),
code=error_status_code(e, 500),
)

View file

@ -350,7 +350,10 @@ from litellm.proxy.common_request_processing import (
_is_azure_model_router_request,
_should_return_raw_model_name,
create_response,
log_llm_api_exception,
open_sse_before_first_byte,
request_litellm_call_id,
resolve_litellm_call_id,
ttft_keepalive_interval,
)
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
@ -389,6 +392,11 @@ from litellm.proxy.common_utils.model_listing_utils import (
from litellm.proxy.common_utils.openai_endpoint_utils import (
remove_sensitive_info_from_deployment,
)
from litellm.proxy.common_utils.openai_error_payload import (
headers_with_litellm_call_id,
litellm_call_id_headers,
with_litellm_call_id,
)
from litellm.proxy.common_utils.periodic_reload_schedule import (
MODEL_COST_MAP_RELOAD_PARAM_NAME,
clear_reload_interval,
@ -681,6 +689,9 @@ from litellm.proxy.types_utils.utils import get_instance_fn
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
router as ui_crud_endpoints_router,
)
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
sync_ui_settings_to_general_settings,
)
from litellm.proxy.ui_crud_endpoints.user_banner_endpoints import (
router as user_banner_endpoints_router,
)
@ -1745,10 +1756,6 @@ class _SSOConfigRow(Protocol):
sso_settings: MutableMapping[str, object]
class _UISettingsRow(Protocol):
ui_settings: Mapping[str, object] | str | None
class _InvitationLinkRow(Protocol):
user_id: str
expires_at: datetime
@ -7404,7 +7411,12 @@ class ProxyConfig:
Returns what the reconcile saw, captured before the lock is released so a
caller's verdict cannot be corrupted by the next reconcile's own in-flight
window. See ReconcileOutcome.
Also re-reads the UI settings that back runtime flags. That runs before the lock, so a
setting written through one pod reaches the others without waiting on a model reconcile.
"""
await sync_ui_settings_to_general_settings(prisma_client)
async with MODEL_RECONCILE_LOCK:
return await self._add_deployment_locked(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
@ -9640,35 +9652,12 @@ class ProxyStartupEvent:
@classmethod
async def _sync_ui_settings_to_general_settings(cls):
"""
Load persisted UI settings from the database and sync runtime flags
into general_settings so they take effect immediately after startup.
"""
try:
import json
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
_RUNTIME_GENERAL_SETTINGS_FLAGS,
)
if prisma_client is None:
return
db_record: Final[_UISettingsRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime is a dict
"_UISettingsRow | None",
await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}),
)
if db_record and db_record.ui_settings:
raw: Final = db_record.ui_settings
ui_settings: Final = json.loads(raw) if isinstance(raw, str) else dict(raw)
flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
if flags_to_sync:
general_settings.update(flags_to_sync)
verbose_proxy_logger.info(
"Synced UI settings to general_settings on startup: %s",
list(flags_to_sync.keys()),
)
except Exception as e:
verbose_proxy_logger.debug("UI settings sync on startup skipped or failed: %s", e)
"""Apply the persisted UI settings to general_settings before this pod serves traffic."""
if prisma_client is None:
return
applied: Final = await sync_ui_settings_to_general_settings(prisma_client)
if applied:
verbose_proxy_logger.info("Synced UI settings to general_settings on startup: %s", list(applied))
@classmethod
async def _load_heuristic_v1_tuning_baselines(
@ -11302,12 +11291,14 @@ async def completion(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - %s", e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
error_msg: Final = f"{e}"
raise ProxyException(
message=getattr(e, "message", error_msg),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
openai_code=getattr(e, "code", None),
code=getattr(e, "status_code", 500),
)
@ -11464,11 +11455,12 @@ async def moderations(
```
"""
global proxy_logging_obj
data: dict = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data: dict = {"litellm_call_id": litellm_call_id}
try:
# Use orjson to parse JSON data, orjson speeds up requests significantly
body: Final = await request.body()
data = orjson.loads(body)
data = orjson.loads(body) | data
# Include original request and headers in the data
data = await add_litellm_data_to_request(
@ -11505,9 +11497,7 @@ async def moderations(
response: Final = await llm_call
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
### RESPONSE HEADERS ###
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
@ -11533,14 +11523,15 @@ async def moderations(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.moderations(): Exception occured - %s", e)
log_llm_api_exception(e, litellm_call_id)
if isinstance(e, ProxyException):
raise
raise with_litellm_call_id(e, litellm_call_id)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "message", str(e)),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
else:
@ -11549,6 +11540,7 @@ async def moderations(
message=getattr(e, "message", error_msg),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
code=getattr(e, "status_code", 500),
)
@ -11586,11 +11578,12 @@ async def audio_speech(
https://platform.openai.com/docs/api-reference/audio/createSpeech
"""
global proxy_logging_obj
data: dict = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data: dict = {"litellm_call_id": litellm_call_id}
try:
# Use orjson to parse JSON data, orjson speeds up requests significantly
body: Final = await request.body()
data = orjson.loads(body)
data = orjson.loads(body) | data
# Include original request and headers in the data
data = await add_litellm_data_to_request(
@ -11623,9 +11616,7 @@ async def audio_speech(
response: Final = await llm_call
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
### RESPONSE HEADERS ###
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
@ -11633,7 +11624,7 @@ async def audio_speech(
cache_key: Final = hidden_params.get("cache_key", None) or ""
api_base: Final = hidden_params.get("api_base", None) or ""
response_cost: Final = hidden_params.get("response_cost", None) or ""
litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
response_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
@ -11644,7 +11635,7 @@ async def audio_speech(
response_cost=response_cost,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
fastest_response_batch_completion=None,
call_id=litellm_call_id,
call_id=response_call_id,
request_data=data,
hidden_params=hidden_params,
)
@ -11680,14 +11671,20 @@ async def audio_speech(
original_exception=e,
request_data=data,
)
verbose_proxy_logger.error("litellm.proxy.proxy_server.audio_speech(): Exception occured - %s", e)
verbose_proxy_logger.debug(traceback.format_exc())
if isinstance(e, (ProxyException, HTTPException)):
raise e
log_llm_api_exception(e, litellm_call_id)
if isinstance(e, ProxyException):
raise with_litellm_call_id(e, litellm_call_id)
if isinstance(e, HTTPException):
raise HTTPException(
status_code=e.status_code,
detail=e.detail,
headers=headers_with_litellm_call_id(e.headers, litellm_call_id),
)
raise ProxyException(
message=getattr(e, "message", f"{e}"),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
openai_code=getattr(e, "code", None),
code=getattr(e, "status_code", 500),
)
@ -11715,11 +11712,12 @@ async def audio_transcriptions(
https://platform.openai.com/docs/api-reference/audio/createTranscription?lang=curl
"""
global proxy_logging_obj
data: dict = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data: dict = {"litellm_call_id": litellm_call_id}
try:
# Use orjson to parse JSON data, orjson speeds up requests significantly
form_data: Final = await get_form_data(request)
data = {key: value for key, value in form_data.items() if key != "file"}
data = {key: value for key, value in form_data.items() if key != "file"} | data
# Include original request and headers in the data
data = await add_litellm_data_to_request(
@ -11786,9 +11784,7 @@ async def audio_transcriptions(
file_object.close() # close the file read in by io library
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
### RESPONSE HEADERS ###
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
@ -11796,7 +11792,7 @@ async def audio_transcriptions(
cache_key: Final = hidden_params.get("cache_key", None) or ""
api_base: Final = hidden_params.get("api_base", None) or ""
response_cost: Final = hidden_params.get("response_cost", None) or ""
litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
response_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
additional_headers: Final[dict] = hidden_params.get("additional_headers", {}) or {}
fastapi_response.headers.update(
@ -11808,7 +11804,7 @@ async def audio_transcriptions(
version=version,
response_cost=response_cost,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
call_id=litellm_call_id,
call_id=response_call_id,
request_data=data,
hidden_params=hidden_params,
**additional_headers,
@ -11830,12 +11826,13 @@ async def audio_transcriptions(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.audio_transcription(): Exception occured - %s", e)
log_llm_api_exception(e, litellm_call_id)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "message", str(e.detail)),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
else:
@ -11844,6 +11841,7 @@ async def audio_transcriptions(
message=getattr(e, "message", error_msg),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
openai_code=getattr(e, "code", None),
code=getattr(e, "status_code", 500),
)
@ -12862,7 +12860,6 @@ from litellm.repositories.table_repositories import (
InvitationLinkRepository,
PromptRepository,
SSOConfigRepository,
UISettingsRepository,
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
@ -15540,18 +15537,34 @@ async def model_group_info(
from litellm.proxy.utils import get_available_models_for_user
# Get available models for the user
all_models_str: Final = await get_available_models_for_user(
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
general_settings=general_settings,
user_model=user_model,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
team_id=None,
include_model_access_groups=False,
only_model_access_groups=False,
return_wildcard_routes=False,
user_api_key_cache=user_api_key_cache,
is_proxy_admin: Final = user_api_key_dict.user_role in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
)
all_models_str: Final = (
get_complete_model_list(
key_models=(),
team_models=(),
proxy_model_list=llm_router.get_model_names(),
user_model=user_model,
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
return_wildcard_routes=False,
llm_router=llm_router,
)
if is_proxy_admin
else await get_available_models_for_user(
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
general_settings=general_settings,
user_model=user_model,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
team_id=None,
include_model_access_groups=False,
only_model_access_groups=False,
return_wildcard_routes=False,
user_api_key_cache=user_api_key_cache,
)
)
model_groups: list[ModelGroupInfoProxy] = _get_model_group_info(
llm_router=llm_router, all_models_str=all_models_str, model_group=model_group

View file

@ -7,12 +7,16 @@ import orjson
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
from fastapi.responses import ORJSONResponse
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
resolve_litellm_call_id,
)
from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
litellm_call_id_headers,
openai_error_param,
openai_error_type,
)
@ -54,10 +58,11 @@ async def rerank(
version,
)
data = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data = {"litellm_call_id": litellm_call_id}
try:
body: Final = await request.body()
data = orjson.loads(body)
data = orjson.loads(body) | data
# Include original request and headers in the data
data = await add_litellm_data_to_request(
@ -82,9 +87,7 @@ async def rerank(
response: Final = await llm_call
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
### RESPONSE HEADERS ###
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
@ -95,7 +98,7 @@ async def rerank(
fastapi_response.headers.update(
ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=hidden_params.get("litellm_call_id", None) or data.get("litellm_call_id", None),
call_id=hidden_params.get("litellm_call_id", None) or litellm_call_id,
model_id=model_id,
cache_key=cache_key,
api_base=api_base,
@ -113,12 +116,13 @@ async def rerank(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.error("litellm.proxy.proxy_server.rerank(): Exception occured - %s", e)
log_llm_api_exception(e, litellm_call_id)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "message", str(e)),
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
param=openai_error_param(e),
headers=litellm_call_id_headers(litellm_call_id),
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
)
else:
@ -127,5 +131,6 @@ async def rerank(
message=getattr(e, "message", error_msg),
type=openai_error_type(e, error_status_code(e, 500)),
param=openai_error_param(e),
headers=litellm_call_id_headers(litellm_call_id),
code=error_status_code(e, 500),
)

View file

@ -14,7 +14,7 @@ from typing import (
from urllib.parse import urlparse
from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile
from pydantic import ConfigDict, JsonValue, ValidationError, create_model
from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model
from pydantic.fields import FieldInfo
from typing_extensions import NotRequired, ReadOnly, TypedDict
@ -29,6 +29,10 @@ from litellm.proxy.config_resolvers.sso import (
SSO_SECRET_FIELDS,
resolve_sso_config,
)
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS,
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
from litellm.proxy.utils import invalidate_config_param
from litellm.repositories.config_repository import ConfigRepository
@ -212,6 +216,9 @@ class UIThemeSettingsResponse(SettingsResponse):
"""Response model for UI theme settings"""
_TEAM_ADMIN_FIELD_ENUM: Final = tuple(sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS))
class UISettings(BaseModel):
"""Configuration for UI-specific flags"""
@ -304,6 +311,18 @@ class UISettings(BaseModel):
description="If true, shows the Chat page in the UI sidebar, letting users chat with an LLM and connect their own MCP server credentials via OAuth.",
)
team_admin_editable_team_fields: Sequence[str] = Field(
default=(),
description=(
"Team settings fields a team admin may change on the teams they administer. "
"Empty means team admins cannot edit team settings at all. "
"Proxy admins and org admins are not affected."
),
json_schema_extra={ # mutable-ok: pydantic only merges json_schema_extra when it is a plain dict
"items": {"type": "string", "enum": [*_TEAM_ADMIN_FIELD_ENUM]}, # mutable-ok: nested in the dict above
},
)
class UISettingsResponse(SettingsResponse):
"""Response model for UI settings"""
@ -326,6 +345,7 @@ ALLOWED_UI_SETTINGS_FIELDS: Final = {
"disable_custom_api_keys",
"disable_key_generate_for_org_admin",
"enable_chat_ui",
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
}
ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: Final = "enable_ptu_cost_attribution"
@ -360,6 +380,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [
"disable_vector_stores_for_internal_users",
"allow_vector_stores_for_team_admins",
"disable_key_generate_for_org_admin",
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
]
# Extension point: packages outside OSS (e.g. litellm_enterprise) can
@ -1457,6 +1478,42 @@ async def get_ui_settings_cached() -> dict[str, JsonValue]:
return ui_settings
_UI_SETTINGS_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def apply_runtime_general_settings_flags(ui_settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
"""Copy the UI settings that gate runtime behavior into ``general_settings``. Returns what was applied."""
from litellm.proxy.proxy_server import general_settings
flags: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
if flags:
general_settings.update(flags)
return MappingProxyType(flags)
async def sync_ui_settings_to_general_settings(prisma_client: object) -> Mapping[str, JsonValue]:
"""Re-read the persisted UI settings and apply the runtime flags to ``general_settings``.
Runs on startup and on every periodic config reload: the PATCH handler only updates the pod
that served it, so every other pod needs its own read to pick up a change without a restart.
Never raises. A read that fails leaves this pod on the flags it already had.
"""
try:
db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique(
where={"id": "ui_settings"}
)
stored: Final = (db_record.ui_settings if db_record else None) or "{}"
parsed: Final = (
_UI_SETTINGS_OBJECT.validate_json(stored)
if isinstance(stored, str)
else _UI_SETTINGS_OBJECT.validate_python(stored)
)
except Exception as e:
verbose_proxy_logger.warning("Could not refresh UI settings from the database: %s", e)
return MappingProxyType({})
return apply_runtime_general_settings_flags(parsed)
@router.get(
"/get/ui_settings",
tags=["UI Settings"],
@ -1485,13 +1542,7 @@ async def get_ui_settings():
# Sanitize any unexpected keys from persisted config before returning
ui_settings: Final = {k: v for k, v in parsed.items() if k in ALLOWED_UI_SETTINGS_FIELDS}
# Sync runtime flags into general_settings so the proxy picks them up
# at runtime (covers server restart scenarios).
_flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
if _flags_to_sync:
from litellm.proxy.proxy_server import general_settings
general_settings.update(_flags_to_sync)
apply_runtime_general_settings_flags(ui_settings)
# Refresh DualCache so other code paths (e.g. /user/filter/ui) see fresh values
from litellm.proxy.proxy_server import user_api_key_cache
@ -1571,6 +1622,20 @@ async def update_ui_settings(
except ValidationError as e:
raise HTTPException(status_code=422, detail=e.errors())
unsupported_team_fields: Final = sorted(
frozenset(settings.team_admin_editable_team_fields) - SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
)
if unsupported_team_fields:
raise HTTPException(
status_code=400,
detail={ # mutable-ok: HTTPException detail must be a plain dict for FastAPI JSON serialization
"error": (
f"{TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING} does not support {unsupported_team_fields}. "
f"Supported fields: {sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS)}."
)
},
)
# Only include fields the caller actually sent (not Pydantic defaults).
settings_dict: Final[Mapping[str, JsonValue]] = settings.model_dump(exclude_unset=True)
@ -1616,13 +1681,7 @@ async def update_ui_settings(
},
)
# Sync runtime flags to general_settings so the proxy picks them up
# at runtime (general_settings is checked in pre-call utils).
_flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
if _flags_to_sync:
from litellm.proxy.proxy_server import general_settings
general_settings.update(_flags_to_sync)
apply_runtime_general_settings_flags(ui_settings)
# Invalidate + set DualCache so subsequent reads see the new values immediately
from litellm.proxy.proxy_server import user_api_key_cache

View file

@ -38,7 +38,11 @@ from litellm.proxy._types import (
SpendLogsMetadata,
SpendLogsPayload,
)
from litellm.proxy.common_utils.openai_error_payload import openai_error_param
from litellm.proxy.common_utils.openai_error_payload import (
litellm_call_id_headers,
openai_error_param,
with_litellm_call_id,
)
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.model_listing import ModelInfoResponse
@ -3050,7 +3054,7 @@ class ProxyLogging:
if litellm_logging_obj is None:
from litellm._uuid import uuid
request_data["litellm_call_id"] = str(uuid.uuid4())
request_data.setdefault("litellm_call_id", str(uuid.uuid4()))
user_api_key_logged_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
user_api_key_dict=user_api_key_dict
)
@ -7659,7 +7663,7 @@ def _recreate_writer_on_read_only_transaction(prisma_client: "PrismaClient | Non
asyncio.create_task(prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction"))
def handle_exception_on_proxy(e: Exception) -> ProxyException:
def handle_exception_on_proxy(e: Exception, litellm_call_id: str | None = None) -> ProxyException:
"""
Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible
"""
@ -7671,20 +7675,23 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException:
_recreate_writer_on_read_only_transaction(prisma_client)
headers: Final = litellm_call_id_headers(litellm_call_id)
if isinstance(e, HTTPException):
return ProxyException(
message=getattr(e, "detail", f"error({e})"),
type=ProxyErrorTypes.internal_server_error,
param=openai_error_param(e),
headers=headers,
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
)
elif isinstance(e, ProxyException):
return e
return with_litellm_call_id(e, litellm_call_id)
_status_code: Final = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
return ProxyException(
message=str(e),
type=ProxyErrorTypes.internal_server_error,
param=openai_error_param(e),
headers=headers,
code=_status_code,
)

View file

@ -482,18 +482,27 @@ class _AsyncPromptManagementOutcome:
def _resolve_responses_api_provider_config(
model: str, custom_llm_provider: str, model_info: object
model: str, custom_llm_provider: str, model_info: object, api_base: str | None
) -> BaseResponsesAPIConfig | None:
provider_config: Final = ProviderConfigManager.get_provider_responses_api_config(
model=model, provider=custom_llm_provider
model=model, provider=custom_llm_provider, api_base=api_base
)
if provider_config is not None or not _deployment_passes_through_responses(model_info):
return provider_config
return OpenAILikeResponsesConfig()
def _api_base_kwarg(kwargs: Mapping[str, object]) -> str | None:
api_base: Final = kwargs.get("api_base")
return api_base if isinstance(api_base, str) else None
def _will_bridge_to_chat_completions(
model: str, custom_llm_provider: str | None, use_chat_completions_api: bool, model_info: object
model: str,
custom_llm_provider: str | None,
use_chat_completions_api: bool,
model_info: object,
api_base: str | None,
) -> bool:
"""``_bridges_to_chat_completions`` for callers running before the provider config is resolved.
@ -507,7 +516,7 @@ def _will_bridge_to_chat_completions(
if custom_llm_provider is None:
return True
return _bridges_to_chat_completions(
_resolve_responses_api_provider_config(normalized_model[0], custom_llm_provider, model_info),
_resolve_responses_api_provider_config(normalized_model[0], custom_llm_provider, model_info, api_base),
use_chat_completions_api or normalized_model[1],
)
@ -618,6 +627,7 @@ async def aresponses(
custom_llm_provider,
bool(kwargs.get("use_chat_completions_api")),
kwargs.get("model_info"),
_api_base_kwarg(kwargs),
),
):
(
@ -783,7 +793,11 @@ def _apply_prompt_management_to_responses_call(
with _prompt_management_sees_a_provisional_message_list(
kwargs,
bridged=_will_bridge_to_chat_completions(
model, custom_llm_provider, use_chat_completions_api, kwargs.get("model_info")
model,
custom_llm_provider,
use_chat_completions_api,
kwargs.get("model_info"),
_api_base_kwarg(kwargs),
),
):
(
@ -1237,7 +1251,7 @@ def responses(
responses_api_provider_config = None
else:
responses_api_provider_config = _resolve_responses_api_provider_config(
model, custom_llm_provider, deployment_model_info
model, custom_llm_provider, deployment_model_info, litellm_params.api_base
)
if (
@ -1496,6 +1510,7 @@ def delete_responses(
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=custom_llm_provider,
api_base=litellm_params.api_base,
)
)
@ -1667,6 +1682,7 @@ def get_responses(
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=custom_llm_provider,
api_base=litellm_params.api_base,
)
)
@ -1811,6 +1827,7 @@ def list_input_items(
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=custom_llm_provider,
api_base=litellm_params.api_base,
)
)
@ -1960,6 +1977,7 @@ def cancel_responses(
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=custom_llm_provider,
api_base=litellm_params.api_base,
)
)
@ -2132,6 +2150,7 @@ def compact_responses(
ProviderConfigManager.get_provider_responses_api_config(
model=model,
provider=custom_llm_provider,
api_base=litellm_params.api_base,
)
)
@ -2270,14 +2289,15 @@ async def _aresponses_websocket(
custom_llm_provider=_custom_llm_provider,
)
resolved_api_base: Final = dynamic_api_base or litellm_params.api_base or litellm.api_base or None
responses_api_provider_config: BaseResponsesAPIConfig | None = None
if _custom_llm_provider is not None:
responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config(
model=resolved_model,
provider=litellm.LlmProviders(_custom_llm_provider),
api_base=resolved_api_base,
)
resolved_api_base: Final = dynamic_api_base or litellm_params.api_base or litellm.api_base or None
resolved_api_key: Final = (
dynamic_api_key
or litellm_params.api_key

View file

@ -2689,7 +2689,7 @@ def declared_value_factory(model: str, custom_llm_provider: str | None, key: str
"""Return a string value the model map declares for *key*, or ``None`` when it says nothing.
The string-valued sibling of :func:`_supports_factory` and
:func:`_is_explicitly_disabled_factory`, public where those two are not because it is read
:func:`is_explicitly_disabled_factory`, public like the latter because both are read
from the provider configs rather than from this module, sharing their
``get_llm_provider`` -> ``_get_model_info_helper`` chain and their unprefixed-twin
fallback (#20885), so a provider-prefixed entry that omits the key still answers
@ -2725,7 +2725,7 @@ def declared_value_factory(model: str, custom_llm_provider: str | None, key: str
return None
def _is_explicitly_disabled_factory(model: str, custom_llm_provider: str | None, key: str) -> bool:
def is_explicitly_disabled_factory(model: str, custom_llm_provider: str | None, key: str) -> bool:
"""Return True only when the model map explicitly sets *key* to ``False``.
This is the opt-out mirror of :func:`_supports_factory`. Where
@ -2844,7 +2844,7 @@ def is_vision_explicitly_disabled(model: str, custom_llm_provider: str | None =
The opt-out mirror of :func:`supports_vision`: a missing declaration reads as not
disabled, so unknown or newly added models stay eligible for image routing.
"""
return _is_explicitly_disabled_factory(model, custom_llm_provider, "supports_vision")
return is_explicitly_disabled_factory(model, custom_llm_provider, "supports_vision")
def supports_vision(model: str, custom_llm_provider: str | None = None) -> bool:
@ -8746,6 +8746,7 @@ class ProviderConfigManager:
def get_provider_responses_api_config(
provider: LlmProviders | str,
model: str | None = None,
api_base: str | None = None,
) -> BaseResponsesAPIConfig | None:
from litellm.llms.openai_like.dynamic_config import (
create_responses_config_class,
@ -8767,7 +8768,7 @@ class ProviderConfigManager:
pass
# Check Python classes first (custom overrides take priority)
result: Final = ProviderConfigManager._get_python_responses_api_config(provider_enum, model)
result: Final = ProviderConfigManager._get_python_responses_api_config(provider_enum, model, api_base)
if result is not None:
return result
@ -8783,6 +8784,7 @@ class ProviderConfigManager:
def _get_python_responses_api_config(
provider: LlmProviders | None,
model: str | None = None,
api_base: str | None = None,
) -> BaseResponsesAPIConfig | None:
"""Check for Python-class-based responses API configs (custom overrides)."""
if provider is None:
@ -8801,6 +8803,14 @@ class ProviderConfigManager:
return litellm.AzureOpenAIOSeriesResponsesAPIConfig()
else:
return litellm.AzureOpenAIResponsesAPIConfig()
elif litellm.LlmProviders.AZURE_AI == provider:
from litellm.llms.azure_ai.common_utils import (
azure_ai_supports_native_responses,
)
if azure_ai_supports_native_responses(model, api_base):
return litellm.AzureAIResponsesAPIConfig()
return None
elif litellm.LlmProviders.XAI == provider:
return litellm.XAIResponsesAPIConfig()
elif litellm.LlmProviders.GITHUB_COPILOT == provider:

View file

@ -26105,6 +26105,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -26162,6 +26163,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28111,6 +28113,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28170,6 +28173,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28592,6 +28596,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28649,6 +28654,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,

View file

@ -30,6 +30,9 @@
- {id: mgmt.key.health.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4292", rationale: "Key health endpoint"}
- {id: mgmt.key.bulk_update.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:2677", rationale: "Batch key updates"}
- {id: mgmt.team.update.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py:1582", rationale: "Metadata/budget updates persist"}
- {id: mgmt.team.update.team_admin_forbidden_until_enabled, module: mgmt, tier: P0, surface: api, assertions: [team_admin_forbidden_until_enabled], source: "team_admin_field_permissions.py:156", rationale: "With no team admin editable fields enabled, a team admin's /team/update is 403 and /team/info reports editing disabled"}
- {id: mgmt.team.update.team_admin_limited_to_enabled_fields, module: mgmt, tier: P0, surface: api, assertions: [team_admin_limited_to_enabled_fields], source: "team_admin_field_permissions.py:156", rationale: "A team admin may change only the enabled fields; a request that also changes any other field is 403 and writes nothing"}
- {id: mgmt.team.update.team_admin_resend_keeps_budget_reset, module: mgmt, tier: P1, surface: api, assertions: [team_admin_resend_keeps_budget_reset], source: "team_admin_field_permissions.py:147", fail_before_fix: proven, rationale: "A team admin resending unchanged budget settings with an enabled field must not push the team's budget reset times back"}
- {id: mgmt.team.delete.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py:1750", rationale: "Deletion prevents key access"}
- {id: mgmt.team.block.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py", rationale: "Block suspends all members"}
- {id: mgmt.team.info.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "team_endpoints.py:2244", rationale: "Metadata+members+budgets"}

View file

@ -1,5 +1,6 @@
"""Live e2e: the /team/* management routes' block, membership, and admin-only
contract.
contract, plus the team settings a team admin may change on /team/update once a
proxy admin enables them under Settings > UI > Team admin editable fields.
Each test creates its team/user/key resources under unique names (deleted on
teardown) and asserts both halves of the contract: the recorded state (the info
@ -8,21 +9,25 @@ Team writes reach the read path once their db/cache entry propagates, so the
read-backs poll to a deadline instead of asserting once.
Everything the shared harness does not already model lives here: the local
request/response models for /team/block, /team/member_update, and the
/team/info fields (blocked flag and per-member budget) these tests assert on.
request/response models for /team/block, /team/member_update, the partial
/team/update, the UI settings allow-list, and the /team/info fields (blocked
flag, limits, budgets, per-member budget, the caller's edit access) these tests
assert on.
"""
from __future__ import annotations
import time
from collections.abc import Callable
from typing import Literal
from collections.abc import Callable, Generator
from contextlib import contextmanager
from datetime import UTC, datetime, timedelta
from typing import Final, Literal
import pytest
from pydantic import BaseModel
from e2e_config import unique_marker
from e2e_http import NoBody, StreamingResponse, unwrap
from e2e_config import settle_propagation, unique_marker
from e2e_http import NoBody, PartialBody, StreamingResponse, unwrap
from lifecycle import ResourceManager
from management_client import ManagementClient
from models import (
@ -39,6 +44,8 @@ pytestmark = pytest.mark.e2e
TeamRole = Literal["admin", "user"]
_TEAM_TPM_LIMIT: Final = 1000
class TeamBlockBody(BaseModel):
team_id: str
@ -66,11 +73,37 @@ class TeamMembership(BaseModel):
litellm_budget_table: MemberBudgetTable | None = None
class TeamInfoData(BaseModel):
class CallerEditAccess(BaseModel):
kind: Literal["unrestricted", "team_admin", "team_admin_disabled", "none"]
editable_fields: list[str] = []
class BudgetWindow(BaseModel):
budget_duration: str
max_budget: float
reset_at: str | None = None
class TeamCustomMetadata(BaseModel):
cost_center: str | None = None
class TeamSettings(BaseModel):
team_alias: str | None = None
models: list[str] = []
tpm_limit: int | None = None
rpm_limit: int | None = None
max_budget: float | None = None
budget_duration: str | None = None
budget_limits: list[BudgetWindow] | None = None
metadata: TeamCustomMetadata | None = None
class TeamInfoData(TeamSettings):
blocked: bool | None = None
members_with_roles: list[MemberRoleEntry] = []
budget_reset_at: datetime | None = None
caller_edit_access: CallerEditAccess | None = None
class TeamInfoRead(BaseModel):
@ -79,6 +112,27 @@ class TeamInfoRead(BaseModel):
team_memberships: list[TeamMembership] = []
class TeamWithAdminNewBody(TeamNewBody):
tpm_limit: int
members_with_roles: list[TeamMemberEntry]
class TeamSettingsChange(PartialBody, TeamSettings):
pass
class TeamSettingsUpdate(TeamSettingsChange):
team_id: str
class TeamAdminEditableFields(BaseModel):
team_admin_editable_team_fields: list[str] = []
class UiSettingsRead(BaseModel):
values: TeamAdminEditableFields
def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T:
deadline = time.monotonic() + client.proxy.poll_timeout
while time.monotonic() < deadline:
@ -107,17 +161,27 @@ def _generate_key(client: ManagementClient, resources: ResourceManager, body: Ke
return key
def _read_team(client: ManagementClient, team_id: str) -> TeamInfoRead:
def _read_team(client: ManagementClient, team_id: str, caller_key: str | None = None) -> TeamInfoRead:
return unwrap(
client.proxy.transport.get(
"/team/info",
headers=client.proxy.transport.master,
headers=client.proxy.transport.master if caller_key is None else client.proxy.transport.bearer(caller_key),
params=TeamInfoParams(team_id=team_id),
response_type=TeamInfoRead,
)
)
def _poll_team(
client: ManagementClient, team_id: str, ready: Callable[[TeamInfoData], bool], failure: str
) -> TeamInfoData:
def read() -> TeamInfoData | None:
info = _read_team(client, team_id).team_info
return info if ready(info) else None
return _poll(client, read, failure)
def _set_blocked(client: ManagementClient, team_id: str, *, blocked: bool) -> None:
_ = unwrap(
client.proxy.transport.post(
@ -301,3 +365,218 @@ class TestTeamManagementRoutes:
client.add_team_member(team_id, member_id)
member_key = _generate_key(client, resources, KeyGenerateBody(user_id=member_id, team_id=team_id))
return member_id, other_id, member_key, team_id
def _team_admin_editable_fields(client: ManagementClient) -> list[str]:
return unwrap(
client.proxy.transport.get(
"/get/ui_settings",
headers=client.proxy.transport.master,
params=NoBody(),
response_type=UiSettingsRead,
)
).values.team_admin_editable_team_fields
def _set_team_admin_editable_fields(client: ManagementClient, fields: list[str]) -> None:
_ = unwrap(
client.proxy.transport.patch(
"/update/ui_settings",
headers=client.proxy.transport.master,
json=TeamAdminEditableFields(team_admin_editable_team_fields=fields),
response_type=NoBody,
)
)
@contextmanager
def _team_admins_may_edit(client: ManagementClient, fields: list[str]) -> Generator[None]:
"""The allow-list is proxy-wide, so restore whatever was there. Other replicas pick a change up on their
config reload, which the wait covers before any team admin call lands on one of them."""
original = _team_admin_editable_fields(client)
_set_team_admin_editable_fields(client, fields)
settle_propagation(time.monotonic())
try:
yield
finally:
_set_team_admin_editable_fields(client, original)
@pytest.fixture(scope="class")
def no_team_admin_editable_fields(client: ManagementClient) -> Generator[None]:
with _team_admins_may_edit(client, []):
yield
@pytest.fixture(scope="class")
def tpm_limit_editable_by_team_admins(client: ManagementClient) -> Generator[None]:
with _team_admins_may_edit(client, ["tpm_limit"]):
yield
def _team_with_admin(client: ManagementClient, resources: ResourceManager) -> tuple[str, str]:
"""A team with a tpm_limit, and the key of a user who is an admin of that team."""
admin_id = _create_user(client, resources, f"e2e-team-admin-{unique_marker()}@example.com")
team_id = client.create_team(
TeamWithAdminNewBody(
team_alias=f"e2e-team-admin-{unique_marker()}",
tpm_limit=_TEAM_TPM_LIMIT,
members_with_roles=[TeamMemberEntry(role="admin", user_id=admin_id)],
)
)
resources.defer(lambda: client.delete_team(team_id))
return team_id, _generate_key(client, resources, KeyGenerateBody(user_id=admin_id))
def _update_team_as(client: ManagementClient, caller_key: str, body: TeamSettingsUpdate) -> StreamingResponse:
return client.proxy.transport.send("/team/update", headers=client.proxy.transport.bearer(caller_key), json=body)
@pytest.mark.usefixtures("no_team_admin_editable_fields")
class TestTeamAdminWithNoEditableFields:
"""No proxy admin has enabled a team field for team admins, which is how every proxy starts."""
@pytest.mark.covers("mgmt.team.update.team_admin_forbidden_until_enabled")
def test_team_admin_cannot_change_any_team_setting(
self, client: ManagementClient, resources: ResourceManager
) -> None:
team_id, admin_key = _team_with_admin(client, resources)
access = _read_team(client, team_id, admin_key).team_info.caller_edit_access
assert access == CallerEditAccess(kind="team_admin_disabled"), (
f"/team/info should tell the team admin that editing is disabled, got {access}"
)
outcome = _update_team_as(client, admin_key, TeamSettingsUpdate(team_id=team_id, tpm_limit=5000))
assert outcome.status_code == 403, (
f"/team/update by a team admin must be 403 while nothing is enabled, got {outcome.status_code}: "
f"{outcome.body[:300]}"
)
assert "cannot edit team settings" in outcome.body, f"403 body should say why, got: {outcome.body[:300]}"
tpm_limit = _read_team(client, team_id).team_info.tpm_limit
assert tpm_limit == _TEAM_TPM_LIMIT, f"the refused update still changed tpm_limit to {tpm_limit}"
@pytest.mark.usefixtures("tpm_limit_editable_by_team_admins")
class TestTeamAdminWithTpmLimitEnabled:
"""A proxy admin has enabled tpm_limit, so a team admin may change that setting and no other."""
@pytest.mark.covers("mgmt.team.update.team_admin_limited_to_enabled_fields")
def test_team_admin_saves_the_settings_form_with_a_new_tpm_limit(
self, client: ManagementClient, resources: ResourceManager
) -> None:
team_id, admin_key = _team_with_admin(client, resources)
access = _read_team(client, team_id, admin_key).team_info.caller_edit_access
assert access == CallerEditAccess(kind="team_admin", editable_fields=["tpm_limit"]), (
f"/team/info should list tpm_limit as the team admin's only editable field, got {access}"
)
before = _read_team(client, team_id).team_info
outcome = _update_team_as(
client,
admin_key,
TeamSettingsUpdate(team_id=team_id, team_alias=before.team_alias, models=before.models, tpm_limit=5000),
)
assert outcome.status_code == 200, (
f"a team admin resending the form with only tpm_limit changed must succeed, got {outcome.status_code}: "
f"{outcome.body[:300]}"
)
after = _poll_team(
client, team_id, lambda info: info.tpm_limit == 5000, "/team/info never reflected tpm_limit=5000"
)
assert after.model_copy(update={"tpm_limit": _TEAM_TPM_LIMIT}) == before, (
f"the update changed more than tpm_limit: before {before}, after {after}"
)
@pytest.mark.covers("mgmt.team.update.team_admin_limited_to_enabled_fields")
@pytest.mark.parametrize(
"change",
[
pytest.param(TeamSettingsChange(rpm_limit=10), id="rpm_limit"),
pytest.param(TeamSettingsChange(max_budget=0.5), id="max_budget"),
pytest.param(TeamSettingsChange(team_alias="renamed-by-team-admin"), id="team_alias"),
pytest.param(TeamSettingsChange(models=["gemini-2.5-flash"]), id="models"),
pytest.param(TeamSettingsChange(budget_duration="1d"), id="budget_duration"),
pytest.param(TeamSettingsChange(metadata=TeamCustomMetadata(cost_center="team-admin")), id="metadata"),
],
)
def test_team_admin_cannot_change_a_setting_that_is_not_enabled(
self, client: ManagementClient, resources: ResourceManager, change: TeamSettingsChange
) -> None:
(field,) = change.model_fields_set
team_id, admin_key = _team_with_admin(client, resources)
before = _read_team(client, team_id).team_info
outcome = _update_team_as(
client,
admin_key,
TeamSettingsUpdate.model_validate(
{**change.model_dump(exclude_unset=True), "team_id": team_id, "tpm_limit": 5000}
),
)
assert outcome.status_code == 403, (
f"a team admin changing {field} must be 403, got {outcome.status_code}: {outcome.body[:300]}"
)
assert f"'{field}'" in outcome.body, f"403 body should name {field}, got: {outcome.body[:300]}"
after = _read_team(client, team_id).team_info
assert after == before, (
f"the refused update still wrote to the team, the enabled tpm_limit included: before {before}, "
f"after {after}"
)
@pytest.mark.covers("mgmt.team.update.team_admin_resend_keeps_budget_reset")
def test_team_admin_resending_the_budget_settings_keeps_the_next_budget_reset(
self, client: ManagementClient, resources: ResourceManager
) -> None:
"""A 120s budget resets at the start of the minute after next. Resending it once the next minute has
started would push that reset a minute later, while the stored reset is still a minute out, so the
proxy's budget reset job cannot be what moves it."""
team_id, admin_key = _team_with_admin(client, resources)
_ = unwrap(
client.proxy.transport.post(
"/team/update",
headers=client.proxy.transport.master,
json=TeamSettingsUpdate(
team_id=team_id,
budget_duration="120s",
budget_limits=[BudgetWindow(budget_duration="120s", max_budget=5.0)],
),
response_type=NoBody,
)
)
budgeted = _poll_team(
client,
team_id,
lambda info: info.budget_reset_at is not None and bool(info.budget_limits),
"/team/info never reflected the 120s budget the proxy admin set",
)
assert budgeted.budget_reset_at is not None
next_minute = budgeted.budget_reset_at - timedelta(seconds=58)
time.sleep(max(0.0, (next_minute - datetime.now(UTC)).total_seconds()))
outcome = _update_team_as(
client,
admin_key,
TeamSettingsUpdate(
team_id=team_id,
tpm_limit=5000,
budget_duration=budgeted.budget_duration,
budget_limits=budgeted.budget_limits,
),
)
assert outcome.status_code == 200, (
f"resending unchanged budget settings with a new tpm_limit must succeed, got {outcome.status_code}: "
f"{outcome.body[:300]}"
)
after = _poll_team(
client, team_id, lambda info: info.tpm_limit == 5000, "/team/info never reflected tpm_limit=5000"
)
assert after.budget_reset_at == budgeted.budget_reset_at, (
f"the team admin pushed the budget reset from {budgeted.budget_reset_at} to {after.budget_reset_at}"
)
assert after.budget_limits == budgeted.budget_limits, (
f"the team admin pushed the budget window resets from {budgeted.budget_limits} to {after.budget_limits}"
)

View file

@ -10,12 +10,11 @@ Pins the five helpers
Driven through /team/new + /team/update.
Structural finding, updated: /team/new loads the org via `get_org_object`
WITH `include_budget_table=True`, so the org max_budget / org tpm / org rpm
guards inside `_check_org_team_limits` are live there and are pinned as
enforced below. /team/update still loads the org without the budget
relation, so its budget guards remain no-ops. The `models` subset guard IS
reachable on both because it reads `org_table.models` directly. The
Structural finding, updated: /team/new and /team/update both load the org
via `get_org_object` WITH `include_budget_table=True`, so the org max_budget /
org tpm / org rpm guards inside `_check_org_team_limits` are live on both and
are pinned as enforced below. The `models` subset guard reads
`org_table.models` directly. The
`_check_user_team_limits` guards reach all branches through
`user_api_key_dict`, no relation include needed.
"""
@ -139,9 +138,8 @@ async def test_check_org_team_limits_models_subset(
# ---------------------------------------------------------------------------
# _check_org_team_limits — budget / tpm / rpm live on /team/new since its
# get_org_object call passes include_budget_table=True. (/team/update still
# loads the org without the budget relation, so its guards remain no-ops.)
# _check_org_team_limits — budget / tpm / rpm live on /team/new and
# /team/update since both get_org_object calls pass include_budget_table=True.
# ---------------------------------------------------------------------------
_ORG_BUDGET_ENFORCED_SCENARIOS = [
@ -216,6 +214,35 @@ async def test_check_org_team_limits_budget_enforced(
assert len(rows) == (1 if expected_status == 200 else 0)
@pytest.mark.parametrize(
"org_budget,body_extras,expected_status",
[(b, c, d) for (_id, b, c, d) in _ORG_BUDGET_ENFORCED_SCENARIOS],
ids=[s[0] for s in _ORG_BUDGET_ENFORCED_SCENARIOS],
)
async def test_check_org_team_limits_budget_enforced_on_update(
org_budget,
body_extras: Dict[str, Any],
expected_status: int,
proxy_client,
prisma,
scratch,
world,
):
org_id = await create_scratch_org(prisma, scratch.prefix, **org_budget)
team_id = await create_scratch_team(prisma, scratch.tag("team"), organization_id=org_id)
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
resp = await proxy_client.post(
"/team/update",
headers={"Authorization": f"Bearer {seeder}"},
json={"team_id": team_id, **body_extras},
)
assert resp.status_code == expected_status, f"{body_extras!r}{resp.status_code}: {resp.text}"
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
persisted = {field: getattr(row, field) for field in body_extras}
assert (persisted == body_extras) == (expected_status == 200)
# ---------------------------------------------------------------------------
# _check_user_team_limits — fires for standalone (no-org) teams created by
# a non-admin caller. Each guard reads from user_api_key_dict / user_obj.
@ -310,65 +337,40 @@ async def test_check_user_team_limits(
# /team/update path — budget authority.
#
# The caller's PERSONAL limits are never applied on update (that compared the
# wrong thing). But raising a team's spend ceiling is reserved for proxy admins:
# a team admin may keep or LOWER the budget, only a proxy admin may RAISE it.
# _check_user_team_limits() only runs on /team/new.
# wrong thing). Raising a team's spend ceiling is reserved for proxy admins.
# max_budget is not on the team-admin allow-list yet (LIT-5722), so a team
# admin is refused in either direction; the raise-only guard underneath the
# allow-list is pinned in the unit tests. _check_user_team_limits() only runs
# on /team/new.
# ---------------------------------------------------------------------------
async def test_team_admin_raise_budget_blocked(proxy_client, prisma, scratch):
"""A team admin cannot raise the team's budget; the block is NOT based on
their personal budget (which here is higher than the requested value)."""
caller_cleartext = await _seed_scratch_actor_with_caps(
prisma,
scratch.prefix,
max_budget=100000.0, # generous personal budget; must not matter
)
creator_user_id = f"{scratch.prefix}-team-creator"
@pytest.mark.parametrize(
"personal_budget,requested_budget",
[(100000.0, 999.0), (10.0, 300.0)],
ids=["raise_with_generous_personal_budget", "lower_with_tiny_personal_budget"],
)
async def test_team_admin_cannot_change_budget_while_max_budget_is_not_editable(
proxy_client, prisma, scratch, personal_budget: float, requested_budget: float
):
caller_cleartext = await _seed_scratch_actor_with_caps(prisma, scratch.prefix, max_budget=personal_budget)
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[creator_user_id],
max_budget=50.0,
)
# Raise the team budget 50 -> 999 as a team admin.
resp = await proxy_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller_cleartext}"},
json={"team_id": team_id, "max_budget": 999.0},
)
assert resp.status_code == 403, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget == 50.0, "team budget must not change on a blocked raise"
async def test_team_admin_lower_budget_allowed(proxy_client, prisma, scratch):
"""A team admin may freely lower (or keep) the team's budget."""
caller_cleartext = await _seed_scratch_actor_with_caps(
prisma,
scratch.prefix,
max_budget=10.0, # below both the old and new team budget; must not matter
)
creator_user_id = f"{scratch.prefix}-team-creator"
team_id = await create_scratch_team(
prisma,
team_id=scratch.tag("team"),
admin_user_ids=[creator_user_id],
admin_user_ids=[f"{scratch.prefix}-team-creator"],
max_budget=500.0,
)
# Lower the team budget 500 -> 300 as a team admin.
resp = await proxy_client.post(
"/team/update",
headers={"Authorization": f"Bearer {caller_cleartext}"},
json={"team_id": team_id, "max_budget": 300.0},
json={"team_id": team_id, "max_budget": requested_budget},
)
assert resp.status_code == 200, resp.text
assert resp.status_code == 403, resp.text
assert "Team admin editable fields" in resp.text, resp.text
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
assert row is not None
assert row.max_budget == 300.0, "team admin should be able to lower the budget"
assert row.max_budget == 500.0, "a refused update must leave the team budget unchanged"
async def test_proxy_admin_raise_budget_allowed(proxy_client, prisma, scratch):

View file

@ -9,31 +9,31 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
# POST /team/update — actor x team-shape matrix (shapes built by _seed_target).
# Each request carries the team's own organization_id so a non-proxy-admin can
# reach the org-scoped branch of the route-permission gate (401 on denial),
# which fronts the handler's _verify_team_access. Only PROXY_ADMIN and an
# ORG_ADMIN of the team's org pass: an internal_user team admin is filtered by
# the route gate before _verify_team_access's team-admin branch is reached.
# The route is self-managed (LIT-5722), so every authenticated caller reaches
# update_team and denials are the handler's 403, never the route gate's 401.
# Only PROXY_ADMIN and an ORG_ADMIN of the team's org pass: a team admin is
# admitted by _resolve_team_access but then refused because no team field is
# enabled for team admins (team_admin_editable_team_fields defaults to empty).
MARKER_ALIAS = "behavior-pin-update-marker-alias"
_MATRIX = [
("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200),
("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200),
("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 401),
("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 401),
("alpha/owner", Actor.OWNER, "alpha", 401),
("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 401),
("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 401),
("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 401),
("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 401),
("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 403),
("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403),
("alpha/owner", Actor.OWNER, "alpha", 403),
("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403),
("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403),
("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403),
("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403),
("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200),
("beta/org_admin", Actor.ORG_ADMIN, "beta", 401),
("beta/team_admin", Actor.TEAM_ADMIN, "beta", 401),
("beta/internal_user", Actor.INTERNAL_USER, "beta", 401),
("beta/owner", Actor.OWNER, "beta", 401),
("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 401),
("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 401),
("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 401),
("beta/org_admin", Actor.ORG_ADMIN, "beta", 403),
("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403),
("beta/internal_user", Actor.INTERNAL_USER, "beta", 403),
("beta/owner", Actor.OWNER, "beta", 403),
("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 403),
("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 403),
("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 403),
("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200),
]
@ -110,8 +110,9 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context(
):
"""With no organization_id in the body the route gate resolves the target
team's org from team_id, so an org admin of the team's own org is allowed
(200), same as PROXY_ADMIN. A team admin of that same team stays denied
(401): the resolution grants org admins access, not team admins."""
(200), same as PROXY_ADMIN. A team admin of that same team reaches the
handler but is refused (403) until a proxy admin enables fields for team
admins, and the response says so."""
await _seed_target(prisma, world, "alpha", scratch.prefix)
allowed_org_admin = await proxy_client.post(
@ -133,21 +134,25 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context(
headers={"Authorization": f"Bearer {world.keys[Actor.TEAM_ADMIN].cleartext}"},
json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS},
)
assert denied_team_admin.status_code == 401, denied_team_admin.text
assert denied_team_admin.status_code == 403, denied_team_admin.text
assert "cannot edit team settings" in denied_team_admin.text, denied_team_admin.text
assert "Team admin editable fields" in denied_team_admin.text, denied_team_admin.text
# Relocation gate — moving a team to a different org. The scratch team starts
# in ORG_A; each scenario relocates it to ORG_B. PROXY_ADMIN bypasses;
# ORG_B_ADMIN clears the route gate (dest-org admin) but fails
# _verify_team_access on the source team (403); the rest fail the route gate
# (401). The relocation-*allowed* branch (caller is org admin of both orgs) is
# covered by test_team_update_org_relocation_allowed_for_dual_org_admin below.
# ORG_B_ADMIN reaches the handler but holds no role on the source team (403);
# ORG_ADMIN holds the source team but not the destination org (403 from the
# relocation gate); the team admin is refused by the empty field allow-list and
# the internal user holds no role at all (403). The relocation-*allowed* branch
# (caller is org admin of both orgs) is covered by
# test_team_update_org_relocation_allowed_for_dual_org_admin below.
_RELOCATION = [
("proxy_admin", Actor.PROXY_ADMIN, 200),
("org_b_admin", Actor.ORG_B_ADMIN, 403),
("org_admin", Actor.ORG_ADMIN, 401),
("team_admin", Actor.TEAM_ADMIN, 401),
("internal_user", Actor.INTERNAL_USER, 401),
("org_admin", Actor.ORG_ADMIN, 403),
("team_admin", Actor.TEAM_ADMIN, 403),
("internal_user", Actor.INTERNAL_USER, 403),
]

View file

@ -809,6 +809,7 @@ def test_img_gen(mock_aimage_generation, client_no_auth):
n=1,
size="1024x1024",
imageConfig={"aspectRatio": "9:16", "imageSize": "1K"},
litellm_call_id=mock.ANY,
metadata=mock.ANY,
proxy_server_request=mock.ANY,
secret_fields=mock.ANY,

View file

@ -87,7 +87,7 @@ class TestSkipPreCallLogic:
await processor.base_process_llm_request(
request=MagicMock(spec=Request),
fastapi_response=MagicMock(spec=Response),
user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
user_api_key_dict=UserAPIKeyAuth(),
route_type="aresponses",
proxy_logging_obj=mock_proxy_logging,
llm_router=MagicMock(),

View file

@ -1934,3 +1934,18 @@ async def test_discovery_auth_fingerprint_tracks_effective_credentials(resolved:
assert original != replaced
assert len(original) == 64
assert "private-original-credential" not in original
@pytest.mark.asyncio
async def test_request_auth_preview_uses_the_same_effective_headers_as_egress() -> None:
from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth
client: Final = MCPClient(
server_url="https://upstream.example/mcp", auth_type=MCPAuth.bearer_token,
resolved_auth=StaticHeaderAuth("Bearer resolved"), extra_headers={"X-Trace": "trace"},
)
request: Final = await client.prepare_request_auth()
assert request.method == "POST"
assert str(request.url) == "https://upstream.example/mcp"
assert request.headers["Authorization"] == "Bearer resolved"
assert request.headers["X-Trace"] == "trace"

View file

@ -19,12 +19,12 @@ to 0 when the only update we saw was the cursor, allowing the
text-based fallback to estimate from the real completion text.
"""
import pytest
import litellm
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
Delta,
ModelResponseStream,
StreamingChoices,
@ -35,6 +35,7 @@ from litellm.types.utils import (
def _make_chunk(
*,
content: str = "",
reasoning_content: str | None = None,
usage: Usage = None,
finish_reason: str = None,
custom_llm_provider: str = "anthropic",
@ -48,7 +49,7 @@ def _make_chunk(
StreamingChoices(
finish_reason=finish_reason,
index=0,
delta=Delta(content=content, role="assistant"),
delta=Delta(content=content, role="assistant", reasoning_content=reasoning_content),
)
],
usage=usage,
@ -69,9 +70,7 @@ class TestAnthropicCursorBug:
token_counter fallback can estimate from completion text.
"""
# Anthropic message_start: input_tokens accurate, output_tokens=1 cursor
message_start = _make_chunk(
usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
)
message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025))
# Several content_block_delta chunks (no usage attached)
text_chunks = [
_make_chunk(content="Hello"),
@ -97,9 +96,7 @@ class TestAnthropicCursorBug:
Normal complete stream: message_start cursor=1, then message_delta=3847.
Last-wins must give 3847 (the real value).
"""
message_start = _make_chunk(
usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
)
message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025))
text_chunks = [_make_chunk(content=t) for t in ["Hello", " world", "!"]]
# message_delta with the real cumulative output_tokens
message_delta = _make_chunk(
@ -119,19 +116,14 @@ class TestAnthropicCursorBug:
End-to-end via calculate_usage(): cursor-only stream + real completion
text should produce a token-counter estimate, NOT 1.
"""
message_start = _make_chunk(
usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
)
message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025))
# ~50 visible chars ≈ ~12 tokens (anthropic-style tokenizer ballpark)
text_chunks = [
_make_chunk(content="Based on your question, I think the answer is "),
_make_chunk(content="forty-two. Here is my reasoning: "),
]
chunks = [message_start, *text_chunks]
completion_output = (
"Based on your question, I think the answer is forty-two. "
"Here is my reasoning: "
)
completion_output = "Based on your question, I think the answer is forty-two. Here is my reasoning: "
processor = ChunkProcessor(chunks=chunks, messages=[])
usage = processor.calculate_usage(
@ -149,9 +141,7 @@ class TestAnthropicCursorBug:
def test_cache_fields_preserved_from_message_start(self):
"""cache_read / cache_creation come from message_start and must survive."""
message_start_usage = Usage(
prompt_tokens=1024, completion_tokens=1, total_tokens=1025
)
message_start_usage = Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
# Anthropic puts these in message_start
message_start_usage.cache_read_input_tokens = 512
message_start_usage.cache_creation_input_tokens = 128
@ -193,9 +183,7 @@ class TestAnthropicCursorBug:
on a 1-token string also gives ~1, so billing is still approximately
correct. This test pins that the result is sane (1 or 0).
"""
message_start = _make_chunk(
usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21)
)
message_start = _make_chunk(usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21))
text_chunk = _make_chunk(content="Yes.")
# Anthropic's message_delta also gives output_tokens=1 in this case
message_delta = _make_chunk(
@ -231,9 +219,7 @@ class TestAnthropicCursorBug:
must fire so token_counter estimates from completion text instead of
billing the placeholder.
"""
message_start_usage = Usage(
prompt_tokens=1024, completion_tokens=1, total_tokens=1025
)
message_start_usage = Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
message_start_usage.cache_read_input_tokens = 4096
message_start = _make_chunk(usage=message_start_usage)
# Subsequent chunks with cache fields but no completion_tokens
@ -253,6 +239,114 @@ class TestAnthropicCursorBug:
"Reset to 0 forces token_counter fallback."
)
@pytest.mark.parametrize("placeholder", [1, 3, 8])
def test_interrupted_reasoning_only_stream_estimates_from_reasoning(self, placeholder: int):
message_start = _make_chunk(
usage=Usage(
prompt_tokens=100,
completion_tokens=placeholder,
total_tokens=100 + placeholder,
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=placeholder),
)
)
reasoning_text = "Let me work through the scheduling constraints step by step. " * 40
reasoning_chunks = [
_make_chunk(reasoning_content=reasoning_text[i : i + 50]) for i in range(0, len(reasoning_text), 50)
]
response = litellm.stream_chunk_builder(
chunks=[message_start, *reasoning_chunks],
messages=[{"role": "user", "content": "Plan the schedule."}],
)
assert response.choices[0].message.reasoning_content == reasoning_text
reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens
assert reasoning_tokens > placeholder
assert response.usage.completion_tokens == reasoning_tokens, (
f"Expected completion_tokens to be the reasoning estimate, got "
f"completion_tokens={response.usage.completion_tokens} reasoning_tokens={reasoning_tokens}"
)
assert response.usage.total_tokens == response.usage.prompt_tokens + reasoning_tokens
details = response.usage.completion_tokens_details
assert details.text_tokens + details.reasoning_tokens == response.usage.completion_tokens
def test_fallback_counts_reasoning_and_text_together(self):
reasoning = "First I should check whether the input is sorted. " * 10
text = "The list is already sorted, so no work is needed."
chunks = [_make_chunk(reasoning_content=reasoning), _make_chunk(content=text)]
response = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "Sort it."}])
text_only = litellm.token_counter(model="claude-sonnet-4-6", text=text, count_response_tokens=True)
details = response.usage.completion_tokens_details
assert details.reasoning_tokens > 0
assert response.usage.completion_tokens == text_only + details.reasoning_tokens
assert details.text_tokens == text_only
def test_lone_usage_event_with_finish_reason_is_trusted(self):
chunks = [
_make_chunk(content="Yes, "),
_make_chunk(content="that works."),
_make_chunk(
usage=Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25),
finish_reason="stop",
),
]
processor = ChunkProcessor(chunks=chunks, messages=[])
result = processor._calculate_usage_per_chunk(chunks=chunks)
assert result["completion_tokens"] == 5
def test_dict_chunks_with_finish_reason_are_trusted(self):
chunks = [
{
"_hidden_params": {"custom_llm_provider": "anthropic"},
"choices": [{"delta": {"content": "Yes, "}, "finish_reason": None}],
},
{
"_hidden_params": {"custom_llm_provider": "anthropic"},
"choices": [{"delta": {"content": "that works."}, "finish_reason": "stop"}],
"usage": Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25),
},
]
processor = ChunkProcessor(chunks=chunks, messages=[])
result = processor._calculate_usage_per_chunk(chunks=chunks)
assert result["completion_tokens"] == 5
def test_dict_chunks_without_finish_reason_reset_placeholder(self):
chunks = [
{
"_hidden_params": {"custom_llm_provider": "anthropic"},
"choices": [],
"usage": Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21),
},
{
"_hidden_params": {"custom_llm_provider": "anthropic"},
"choices": [{"delta": {"content": "partial"}, "finish_reason": None}],
},
]
processor = ChunkProcessor(chunks=chunks, messages=[])
result = processor._calculate_usage_per_chunk(chunks=chunks)
assert result["completion_tokens"] == 0
assert result["completion_tokens_details"] is None
def test_estimated_reasoning_is_capped_to_trusted_completion_total(self):
chunks = [
_make_chunk(reasoning_content="Let me reason about this carefully and at length. " * 20),
_make_chunk(
finish_reason="stop",
usage=Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25),
),
]
response = litellm.stream_chunk_builder(
chunks=chunks,
messages=[{"role": "user", "content": "Go."}],
)
details = response.usage.completion_tokens_details
assert response.usage.completion_tokens == 5
assert details.reasoning_tokens <= response.usage.completion_tokens
assert details.reasoning_tokens + details.text_tokens == response.usage.completion_tokens
assert details.text_tokens >= 0
class TestProviderGuard:
"""Class A: the cursor-reset heuristic must NOT silently affect non-Anthropic
@ -297,11 +391,12 @@ class TestNonAnthropicStreamingIntact:
"""Make sure providers without cursor pattern still work."""
def test_completion_tokens_above_one_never_resets(self):
"""Any chunk reporting completion_tokens > 1 sets saw_non_cursor
and prevents the reset."""
"""A non-Anthropic provider reporting completion_tokens > 1 from a
single usage event keeps that value."""
chunks = [
_make_chunk(
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
custom_llm_provider="openai",
),
]
processor = ChunkProcessor(chunks=chunks, messages=[])

View file

@ -677,3 +677,14 @@ def test_azure_responses_gpt6_astra_rejects_temperature_while_reasoning(local_mo
model="gpt-6-astra",
drop_params=False,
)
def test_azure_responses_sends_the_deployment_name_when_azure_ai_prefix_survives_provider_remap():
request = AzureOpenAIResponsesAPIConfig().transform_responses_api_request(
model="azure_ai/gpt-5.4-nano",
input="hi",
response_api_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert request["model"] == "gpt-5.4-nano"

View file

@ -0,0 +1,311 @@
import json
import httpx
import pytest
import respx
import litellm
from litellm.llms.azure_ai.responses.transformation import AzureAIResponsesAPIConfig
from litellm.responses.main import _will_bridge_to_chat_completions
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager
FOUNDRY_PROJECT_BASE = "https://res.services.ai.azure.com/api/projects/proj"
FOUNDRY_RESPONSES_URL = f"{FOUNDRY_PROJECT_BASE}/openai/v1/responses"
SERVERLESS_BASE = "https://endpoint.eastus.models.ai.azure.com"
WEATHER_TOOL = {
"type": "function",
"name": "get_weather",
"description": "Get weather",
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
}
@pytest.fixture(autouse=True)
def clear_azure_ai_env(monkeypatch):
monkeypatch.setattr(litellm, "api_base", None)
monkeypatch.setattr(litellm, "api_key", None)
monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
for env_var in (
"AZURE_AI_API_BASE",
"AZURE_AI_API_KEY",
"AZURE_AD_TOKEN",
"AZURE_TENANT_ID",
"AZURE_CLIENT_ID",
"AZURE_CLIENT_SECRET",
):
monkeypatch.delenv(env_var, raising=False)
def _responses_payload(model: str) -> dict:
return {
"id": "resp_123",
"object": "response",
"created_at": 1741369938,
"status": "completed",
"model": model,
"output": [],
"parallel_tool_calls": False,
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
"error": None,
"tool_choice": "auto",
"tools": [],
"metadata": None,
"temperature": None,
"top_p": None,
"max_output_tokens": None,
"previous_response_id": None,
"reasoning": None,
"truncation": None,
"instructions": None,
"incomplete_details": None,
"user": None,
}
def _chat_completion_payload(model: str) -> dict:
return {
"id": "chatcmpl-123",
"object": "chat.completion",
"created": 1741369938,
"model": model,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
@pytest.mark.parametrize("model", ["gpt-5.6-luna-20260710154139", "gpt-5.5-20260504143601", "DeepSeek-R1-0528", None])
@pytest.mark.parametrize(
"api_base", [FOUNDRY_PROJECT_BASE, "https://res.services.ai.azure.com", "https://res.openai.azure.com"]
)
def test_azure_openai_v1_hosts_resolve_native_config(model, api_base):
config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model=model, api_base=api_base)
assert isinstance(config, AzureAIResponsesAPIConfig)
def test_api_base_from_env_resolves_native_config(monkeypatch):
monkeypatch.setenv("AZURE_AI_API_BASE", FOUNDRY_PROJECT_BASE)
config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model="gpt-5.6-luna", api_base=None)
assert isinstance(config, AzureAIResponsesAPIConfig)
@pytest.mark.parametrize("model", ["gpt-5.6-luna", None])
@pytest.mark.parametrize(
"api_base",
[SERVERLESS_BASE, "https://endpoint.eastus.inference.ml.azure.com/score", "https://res.cognitiveservices.azure.com"],
)
def test_other_hosts_keep_chat_bridge(model, api_base):
config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model=model, api_base=api_base)
assert config is None
@pytest.mark.parametrize("model", ["claude-3-5-sonnet", "model_router/gpt-5", "agents/my-agent"])
def test_non_openai_surfaces_keep_chat_bridge(model):
config = ProviderConfigManager.get_provider_responses_api_config(
provider="azure_ai", model=model, api_base=FOUNDRY_PROJECT_BASE
)
assert config is None
@pytest.mark.parametrize("api_base,bridged", [(FOUNDRY_PROJECT_BASE, False), (SERVERLESS_BASE, True)])
def test_will_bridge_to_chat_completions_follows_host(api_base, bridged):
assert _will_bridge_to_chat_completions("gpt-5.6-luna", "azure_ai", False, None, api_base) is bridged
@pytest.mark.parametrize(
"api_base,expected",
[
(FOUNDRY_PROJECT_BASE, FOUNDRY_RESPONSES_URL),
(f"{FOUNDRY_PROJECT_BASE}/", FOUNDRY_RESPONSES_URL),
(f"{FOUNDRY_PROJECT_BASE}/openai/v1", FOUNDRY_RESPONSES_URL),
(FOUNDRY_RESPONSES_URL, FOUNDRY_RESPONSES_URL),
("https://res.services.ai.azure.com", "https://res.services.ai.azure.com/openai/v1/responses"),
("https://res.services.ai.azure.com/models", "https://res.services.ai.azure.com/openai/v1/responses"),
(
"https://res.services.ai.azure.com/models/chat/completions?api-version=2024-05-01-preview",
"https://res.services.ai.azure.com/openai/v1/responses",
),
("https://res.openai.azure.com", "https://res.openai.azure.com/openai/v1/responses"),
(
"https://res.openai.azure.com/openai/deployments/gpt-5?api-version=2025-04-01-preview",
"https://res.openai.azure.com/openai/v1/responses",
),
],
)
def test_get_complete_url(api_base, expected):
assert AzureAIResponsesAPIConfig().get_complete_url(api_base=api_base, litellm_params={}) == expected
def test_get_complete_url_ignores_api_version():
url = AzureAIResponsesAPIConfig().get_complete_url(
api_base=FOUNDRY_PROJECT_BASE, litellm_params={"api_version": "2025-04-01-preview"}
)
assert url == FOUNDRY_RESPONSES_URL
def test_get_complete_url_uses_env_api_base(monkeypatch):
monkeypatch.setenv("AZURE_AI_API_BASE", FOUNDRY_PROJECT_BASE)
assert AzureAIResponsesAPIConfig().get_complete_url(api_base=None, litellm_params={}) == FOUNDRY_RESPONSES_URL
def test_get_complete_url_raises_without_api_base():
with pytest.raises(ValueError, match="AZURE_AI_API_BASE"):
AzureAIResponsesAPIConfig().get_complete_url(api_base=None, litellm_params={})
def test_native_websocket_stays_off():
assert AzureAIResponsesAPIConfig().supports_native_websocket() is False
def test_validate_environment_sends_api_key_header():
headers = AzureAIResponsesAPIConfig().validate_environment(
headers={"x-custom": "1"},
model="gpt-5.6-luna",
litellm_params=GenericLiteLLMParams(api_key="secret", api_base=FOUNDRY_PROJECT_BASE),
)
assert headers == {"x-custom": "1", "api-key": "secret", "Content-Type": "application/json"}
def test_validate_environment_reads_api_key_from_env(monkeypatch):
monkeypatch.setenv("AZURE_AI_API_KEY", "env-secret")
headers = AzureAIResponsesAPIConfig().validate_environment(
headers={}, model="gpt-5.6-luna", litellm_params=GenericLiteLLMParams(api_base=FOUNDRY_PROJECT_BASE)
)
assert headers["api-key"] == "env-secret"
def test_validate_environment_uses_entra_token_without_api_key():
headers = AzureAIResponsesAPIConfig().validate_environment(
headers={},
model="gpt-5.6-luna",
litellm_params=GenericLiteLLMParams(azure_ad_token="entra-token", api_base=FOUNDRY_PROJECT_BASE),
)
assert headers["Authorization"] == "Bearer entra-token"
assert "api-key" not in headers
def test_validate_environment_raises_without_credentials():
with pytest.raises(ValueError, match="AZURE_AI_API_KEY"):
AzureAIResponsesAPIConfig().validate_environment(
headers={}, model="gpt-5.6-luna", litellm_params=GenericLiteLLMParams(api_base=FOUNDRY_PROJECT_BASE)
)
NATIVE_RESPONSES_CASES = [
("azure_ai/gpt-5.6-luna-20260710154139", FOUNDRY_PROJECT_BASE, FOUNDRY_RESPONSES_URL, "gpt-5.6-luna-20260710154139"),
(
"azure_ai/gpt-5.6-luna",
"https://res.services.ai.azure.com/models",
"https://res.services.ai.azure.com/openai/v1/responses",
"gpt-5.6-luna",
),
(
"azure_ai/gpt-5.6-sol",
"https://res.services.ai.azure.com",
"https://res.services.ai.azure.com/openai/v1/responses",
"gpt-5.6-sol",
),
(
"azure_ai/gpt-5.6-luna-20260710154139",
"https://res.openai.azure.com",
"https://res.openai.azure.com/openai/v1/responses",
"gpt-5.6-luna-20260710154139",
),
(
"azure_ai/gpt-5.6-sol",
"https://res.openai.azure.com",
"https://res.openai.azure.com/openai/v1/responses",
"gpt-5.6-sol",
),
]
def _assert_native_responses_request(route, expected_url, expected_model):
request = route.calls.last.request
body = json.loads(request.content)
assert f"{request.url.scheme}://{request.url.host}{request.url.path}" == expected_url
assert request.headers["api-key"] == "fake-key"
assert body["model"] == expected_model
assert body["input"] == "What is the weather in SF?"
assert "messages" not in body
assert body["reasoning"] == {"effort": "high"}
assert body["tools"] == [WEATHER_TOOL]
@pytest.mark.asyncio
@respx.mock
@pytest.mark.parametrize("model,api_base,expected_url,expected_model", NATIVE_RESPONSES_CASES)
async def test_aresponses_sends_reasoning_and_tools_to_native_endpoint(model, api_base, expected_url, expected_model):
route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock(
return_value=httpx.Response(200, json=_responses_payload(expected_model))
)
await litellm.aresponses(
model=model,
input="What is the weather in SF?",
reasoning_effort="high",
tools=[WEATHER_TOOL],
api_base=api_base,
api_key="fake-key",
)
_assert_native_responses_request(route, expected_url, expected_model)
@pytest.mark.asyncio
@respx.mock
async def test_aresponses_catalog_name_remapped_to_azure_sends_bare_deployment_name(monkeypatch):
monkeypatch.setenv("AZURE_AI_API_BASE", "https://res.openai.azure.com")
route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock(
return_value=httpx.Response(200, json=_responses_payload("gpt-5.4-nano"))
)
await litellm.aresponses(
model="azure_ai/gpt-5.4-nano",
input="What is the weather in SF?",
api_base="https://res.openai.azure.com",
api_key="fake-key",
)
assert json.loads(route.calls.last.request.content)["model"] == "gpt-5.4-nano"
@pytest.mark.asyncio
@respx.mock
@pytest.mark.parametrize("model,api_base,expected_url,expected_model", NATIVE_RESPONSES_CASES)
async def test_router_aresponses_sends_bare_deployment_name(model, api_base, expected_url, expected_model):
route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock(
return_value=httpx.Response(200, json=_responses_payload(expected_model))
)
router = litellm.Router(
model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": model, "api_base": api_base, "api_key": "fake-key"}}],
num_retries=0,
)
await router.aresponses(
model="gpt-5.6", input="What is the weather in SF?", reasoning={"effort": "high"}, tools=[WEATHER_TOOL]
)
_assert_native_responses_request(route, expected_url, expected_model)
@pytest.mark.asyncio
@respx.mock
async def test_aresponses_serverless_host_stays_on_chat_bridge():
chat_route = respx.post(url__regex=r".*/chat/completions$").mock(
return_value=httpx.Response(200, json=_chat_completion_payload("gpt-5.6-luna"))
)
responses_route = respx.post(url__regex=r".*/responses$")
await litellm.aresponses(
model="azure_ai/gpt-5.6-luna-20260710154139",
input="What is the weather in SF?",
tools=[WEATHER_TOOL],
api_base=SERVERLESS_BASE,
api_key="fake-key",
)
assert chat_route.called
assert not responses_route.called
assert chat_route.calls.last.request.headers["Authorization"] == "Bearer fake-key"

View file

@ -164,6 +164,21 @@ class TestDashScopeConfig:
assert transformed_messages[0].get("cache_control") == {"type": "ephemeral"}
@pytest.mark.parametrize("reasoning_effort", ["none", "minimal", "low", "high"])
def test_dashscope_forwards_reasoning_effort(self, reasoning_effort: str):
"""DashScope supports reasoning_effort, so it must reach the provider instead of being dropped."""
assert "reasoning_effort" in DashScopeChatConfig().get_supported_openai_params(
model="qwen3.7-plus"
)
optional_params = litellm.get_optional_params(
model="qwen3.7-plus",
custom_llm_provider="dashscope",
reasoning_effort=reasoning_effort,
)
assert optional_params["reasoning_effort"] == reasoning_effort
def test_dashscope_preserves_cache_control_in_tools(self):
"""DashScope should NOT strip cache_control from tools."""
config = DashScopeChatConfig()

View file

@ -8,7 +8,7 @@ from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config
from litellm.llms.openai.openai import OpenAIConfig
from litellm.utils import (
_is_explicitly_disabled_factory,
is_explicitly_disabled_factory,
peek_reasoning_summary_aliases,
strip_reasoning_summary_aliases_from_optional_params,
)
@ -524,19 +524,19 @@ def test_gpt5_minimal_explicitly_disabled_check(gpt5_config: OpenAIGPT5Config):
def test_is_explicitly_disabled_factory_minimal():
"""_is_explicitly_disabled_factory returns True only for explicit False entries.
"""is_explicitly_disabled_factory returns True only for explicit False entries.
Verifies the shared helper used by _is_reasoning_effort_level_explicitly_disabled
directly so future changes to the helper are caught without going through the
method wrapper.
"""
key = "supports_minimal_reasoning_effort"
assert _is_explicitly_disabled_factory("gpt-5.4-mini", None, key)
assert _is_explicitly_disabled_factory("gpt-5.4-nano", None, key)
assert _is_explicitly_disabled_factory("openai/gpt-5.4-mini", None, key)
assert _is_explicitly_disabled_factory("gpt-5.4", None, key)
assert _is_explicitly_disabled_factory("gpt-5.4-pro", None, key)
assert not _is_explicitly_disabled_factory("gpt-5.4-turbo-preview", None, key)
assert is_explicitly_disabled_factory("gpt-5.4-mini", None, key)
assert is_explicitly_disabled_factory("gpt-5.4-nano", None, key)
assert is_explicitly_disabled_factory("openai/gpt-5.4-mini", None, key)
assert is_explicitly_disabled_factory("gpt-5.4", None, key)
assert is_explicitly_disabled_factory("gpt-5.4-pro", None, key)
assert not is_explicitly_disabled_factory("gpt-5.4-turbo-preview", None, key)
def test_gpt5_unknown_model_passes_through_minimal(config: OpenAIConfig):

View file

@ -5,11 +5,14 @@ from copy import deepcopy
from typing import Final, List, cast
from unittest.mock import MagicMock, patch
import httpx
import pytest
from pydantic import BaseModel
import litellm
from litellm import ModelResponse, completion
from litellm.llms.anthropic.experimental_pass_through.messages import handler as anthropic_messages_handler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
@ -2678,6 +2681,118 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3():
assert result["thinkingConfig"]["includeThoughts"] is False
@pytest.mark.parametrize(
"model",
[
"gemini-3.7-flash",
"vertex_ai/gemini-3.8-flash",
"gemini/gemini-3.8-flash",
],
)
@pytest.mark.parametrize(
("reasoning_effort", "include_thoughts"),
[("minimal", True), ("none", False), ("disable", False)],
)
def test_gemini_37_38_flash_floor_minimal_thinking_level(
local_model_cost_map, model, reasoning_effort, include_thoughts
):
result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(
reasoning_effort, model
)
assert result["thinkingLevel"] == "low"
assert result["includeThoughts"] is include_thoughts
@pytest.mark.parametrize(
("model", "reasoning_effort", "expected_level", "include_thoughts"),
[
("gemini-3-flash-preview", "minimal", "minimal", True),
("gemini-3-flash-preview", "none", "minimal", False),
("gemini-3-flash-preview", "disable", "minimal", False),
("gemini-3.6-flash", "minimal", "minimal", True),
("gemini-3.6-flash", "none", "minimal", False),
("gemini-3.6-flash", "disable", "minimal", False),
("gemini-3.5-flash", "minimal", "minimal", True),
("gemini-3.5-flash", "none", "minimal", False),
("gemini-3.5-flash", "disable", "minimal", False),
("gemini-3.8-flash", "medium", "medium", True),
],
)
def test_gemini_flash_minimal_thinking_support(
local_model_cost_map, model, reasoning_effort, expected_level, include_thoughts
):
result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(
reasoning_effort, model
)
assert result["thinkingLevel"] == expected_level
assert result["includeThoughts"] is include_thoughts
def test_gemini_38_flash_feature_flag_uses_low_thinking_level(local_model_cost_map, monkeypatch):
monkeypatch.setattr(litellm, "enable_gemini_default_thinking_level_low", True)
thinking_param = {"type": "enabled", "budget_tokens": 1024}
result_38 = VertexGeminiConfig._map_thinking_param(
thinking_param, model="gemini-3.8-flash"
)
result_36 = VertexGeminiConfig._map_thinking_param(
thinking_param, model="gemini-3.6-flash"
)
assert result_38["thinkingLevel"] == "low"
assert result_36["thinkingLevel"] == "minimal"
def test_gemini_38_flash_public_reasoning_effort_none_uses_low(local_model_cost_map):
result = VertexGeminiConfig().map_openai_params(
non_default_params={"reasoning_effort": "none"},
optional_params={},
model="gemini-3.8-flash",
drop_params=False,
)
assert result["thinkingConfig"] == {
"thinkingLevel": "low",
"includeThoughts": False,
}
@pytest.mark.asyncio
async def test_gemini_38_flash_messages_bridge_thinking_disabled_sends_low_thinking_level(local_model_cost_map):
captured: dict[str, dict] = {}
def upstream(request: httpx.Request) -> httpx.Response:
captured["body"] = json.loads(request.content)
return httpx.Response(
200,
json={
"candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP"}],
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
},
request=request,
)
client = AsyncHTTPHandler()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream))
await anthropic_messages_handler.anthropic_messages(
max_tokens=16,
messages=[{"role": "user", "content": "hi"}],
model="gemini/gemini-3.8-flash",
custom_llm_provider="gemini",
thinking={"type": "disabled"},
api_key="fake-gemini-key",
client=client,
)
assert captured["body"]["generationConfig"]["thinkingConfig"] == {
"thinkingLevel": "low",
"includeThoughts": False,
}
def test_reasoning_effort_dict_format_gemini_3():
"""
Test that reasoning_effort works when passed as dict format from OpenAI Agents SDK.

View file

@ -7,6 +7,7 @@ maps each CredError onto its HTTP status. These pin the parity-critical mapping
import base64
from types import SimpleNamespace
from typing import Final
import pytest
from fastapi import HTTPException
@ -20,7 +21,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import
raise_user_oauth_challenge,
to_server_spec,
to_subject,
validate_static_credential,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
ApiKeyConfig,
AuthorizationCodeConfig,
@ -34,10 +37,28 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
SharedKey,
TokenExchangeConfig,
)
from litellm.types.mcp import MCPAuth, MCPTransport
from litellm.types.mcp import MCPAuth, MCPAuthType, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@pytest.mark.parametrize("auth_type,header,value", [
(MCPAuth.api_key, "Authorization", "Bearer fixture-key"),
(MCPAuth.api_key, "Authorization", "ApiKey fixture-key"),
(MCPAuth.api_key, "Authorization", "token fixture-key"),
(MCPAuth.api_key, "Authorization", "Bearer token"),
(MCPAuth.api_key, "Authorization", "opaque-key"),
(MCPAuth.api_key, "Authorization", "Custom Custom"),
(MCPAuth.api_key, "X-API-Key", "Bearer Bearer"),
(MCPAuth.api_key, "X-Custom", "ApiKey ApiKey"),
(MCPAuth.authorization, "Authorization", "opaque-secret-value"),
])
def test_static_credential_preserves_supported_api_key_and_raw_headers(
auth_type: MCPAuthType, header: str, value: str,
) -> None:
result: Final = validate_static_credential(auth_type, {header: value}, upstream_token_header=header)
assert isinstance(result, Ok)
def _server(**kwargs) -> MCPServer:
return MCPServer(server_id="s", name="n", transport=MCPTransport.http, **kwargs)
@ -155,12 +176,6 @@ def test_oauth2_user_token_maps_to_authorization_code(oauth2_flow):
_server(auth_type=MCPAuth.api_key), # no token configured
_server(auth_type=MCPAuth.bearer_token), # no token configured
_server(auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True), # delegated upstream OAuth -> v1
_server(auth_type=MCPAuth.oauth2_token_exchange), # no endpoint/client creds -> incomplete -> v1
_server(
auth_type=MCPAuth.oauth2_token_exchange,
token_exchange_endpoint="https://idp/token",
client_id="cid",
), # missing client_secret -> incomplete -> v1
_server(auth_type=MCPAuth.aws_sigv4),
_server(auth_type=None, oauth_passthrough=True, extra_headers=["Authorization"]),
],
@ -802,3 +817,14 @@ def test_a_blank_header_name_means_unset_rather_than_an_error(blank):
spec = to_server_spec(server)
assert spec is not None
assert spec.config.header_name == "Authorization"
@pytest.mark.parametrize("client_secret", [None, ""])
@pytest.mark.parametrize("is_byok", [False, True])
def test_incomplete_obo_keeps_exchange_ownership(client_secret: str | None, is_byok: bool) -> None:
spec = to_server_spec(_server(auth_type=MCPAuth.oauth2_token_exchange, client_id="client",
client_secret=client_secret, is_byok=is_byok))
assert spec is not None
assert isinstance(spec.config, TokenExchangeConfig)
assert spec.config.client_id == "client"
assert spec.config.client_secret is None

View file

@ -5,11 +5,13 @@ import logging
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Final, Literal, Optional
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from respx import MockRouter
from litellm.proxy._experimental.mcp_server.exceptions import (
MCPServerListError,
@ -5127,7 +5129,8 @@ class TestMCPServerManager:
captured: dict = {}
def fake_create_tool_function(
path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False
path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False,
auth_type=None, upstream_token_header=None,
):
captured["headers"] = headers
captured["server_label"] = server_label
@ -5212,7 +5215,8 @@ class TestMCPServerManager:
captured: dict = {}
def fake_create_tool_function(
path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False
path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False,
auth_type=None, upstream_token_header=None,
):
captured["headers"] = headers
@ -9401,12 +9405,13 @@ class TestCreateMcpClientV2Graft:
assert "misconfigured" in str(exc_info.value.detail)
assert "token_url" in str(exc_info.value.detail)
async def test_static_token_missing_defers_to_v1(self):
client = await MCPServerManager()._create_mcp_client(
self._http_server(auth_type=MCPAuth.api_key, authentication_token=None)
)
assert client._resolved_auth is None
async def test_static_token_missing_rejects_before_connecting(self):
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(
self._http_server(auth_type=MCPAuth.api_key, authentication_token=None)
)
assert exc.value.status_code == 500
assert "credential" in str(exc.value.detail)
async def test_stdio_migrated_auth_type_still_defers_to_v1(self):
client = await MCPServerManager()._create_mcp_client(
@ -13467,3 +13472,400 @@ async def test_discovery_cache_returns_oversized_results_without_retaining_them(
result: Final = await cache.get(("server", None), fetch)
assert result[0].description == description
assert fetch.await_count == 2
class TestProtectedCredentialPreparation:
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type,credential", [
(MCPAuth.bearer_token, None),
(MCPAuth.bearer_token, "Bearer"),
(MCPAuth.api_key, None),
(MCPAuth.basic, "Basic"),
])
@pytest.mark.parametrize("dispatch", ["managed", "local"])
async def test_openapi_dispatch_rejects_unusable_effective_credentials(
self, tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch,
auth_type: MCPAuthType, credential: str | None, dispatch: str,
) -> None:
from litellm.proxy._experimental.mcp_server.server import _handle_local_mcp_tool
from litellm.proxy._experimental.mcp_server.utils import add_server_prefix_to_name, get_server_prefix
spec_path: Final = tmp_path / "openapi.json"
spec_path.write_text(json.dumps({"openapi": "3.0.0", "info": {"title": "Auth", "version": "1"},
"paths": {"/echo": {"get": {"operationId": "echo"}}}}))
server: Final = MCPServer(
server_id="dispatch-auth", name="dispatch-auth", url="https://upstream.example",
transport=MCPTransport.http, auth_type=auth_type, authentication_token=credential,
)
manager: Final = MCPServerManager()
await manager._register_openapi_tools(str(spec_path), server, server.url)
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="unexpected success")
result: Final = (
await manager._call_openapi_tool_handler(server, "echo", {})
if dispatch == "managed"
else await _handle_local_mcp_tool(add_server_prefix_to_name("echo", get_server_prefix(server)), {})
)
assert result.isError is True
assert "requires a usable upstream credential" in result.content[0].text
assert destination.call_count == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("transport", [MCPTransport.http, MCPTransport.sse])
@pytest.mark.parametrize("client_secret", [None, ""])
@pytest.mark.parametrize("subject", [None, "caller-subject"])
async def test_incomplete_obo_rejects_caller_and_static_fallback(
self, transport: MCPTransport, client_secret: str | None, subject: str | None
) -> None:
server = MCPServer(
server_id="incomplete-obo", name="incomplete-obo", url="https://upstream.example/mcp",
transport=transport, auth_type=MCPAuth.oauth2_token_exchange,
client_id="gateway", client_secret=client_secret,
token_exchange_endpoint="https://idp.example/token", authentication_token="static-fallback",
)
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(
server, mcp_auth_header="Bearer override", subject_token=subject,
)
assert exc.value.status_code == (401 if subject is None else 500)
assert "static-fallback" not in str(exc.value.detail)
assert "override" not in str(exc.value.detail)
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type", [MCPAuth.api_key, MCPAuth.bearer_token])
@pytest.mark.parametrize("credential", [None, "", " ", {"X-Trace": "trace"}])
async def test_static_auth_without_usable_credential_rejects(
self, auth_type: MCPAuthType, credential: str | dict[str, str] | None
) -> None:
server = MCPServer(
server_id="empty-static", name="empty-static", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=auth_type,
)
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(server, mcp_auth_header=credential)
assert exc.value.status_code == 500
assert "credential" in str(exc.value.detail).lower()
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type,headers", [
(MCPAuth.api_key, {"X-API-Key": "key"}),
(MCPAuth.bearer_token, {"Authorization": "Bearer token"}),
])
async def test_static_auth_accepts_actual_forwarded_credential(
self, auth_type: MCPAuthType, headers: dict[str, str]
) -> None:
server = MCPServer(
server_id="header-static", name="header-static", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=auth_type,
)
client = await MCPServerManager()._create_mcp_client(server, extra_headers=headers)
assert client._get_auth_headers() == headers
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange])
async def test_openapi_protected_auth_rejects_missing_credentials(self, auth_type: MCPAuthType) -> None:
server = MCPServer(
server_id="openapi-empty", name="openapi-empty", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=auth_type,
token_exchange_endpoint="https://idp.example/token",
)
with pytest.raises(HTTPException) as exc:
await MCPServerManager().resolve_openapi_upstream_auth(
mcp_server=server, oauth2_headers=None, raw_headers=None, mcp_auth_header=None,
user_api_key_auth=None, forwarded_headers=None,
)
assert exc.value.status_code in (401, 500)
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type,slot,value", [
(MCPAuth.api_key, "X-API-Key", "token"),
(MCPAuth.authorization, "Authorization", "opaque-secret-value"),
(MCPAuth.authorization, "Authorization", "Bearer abc"),
(MCPAuth.authorization, "Authorization", "Custom abc"),
])
async def test_raw_static_credentials_are_forwarded_unchanged(
self, auth_type: MCPAuthType, slot: str, value: str,
) -> None:
server = MCPServer(server_id="raw-key", name="raw-key", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=auth_type, authentication_token=value)
client = await MCPServerManager()._create_mcp_client(server)
assert client._resolved_auth is not None
request = httpx.Request("GET", server.url)
flow = client._resolved_auth.auth_flow(request)
try:
assert next(flow).headers[slot] == value
finally:
flow.close()
@pytest.mark.asyncio
@pytest.mark.parametrize("value", ["Bearer", "basic", "token", "ApiKey", " bEaReR ", "\tTOKEN\t"])
@pytest.mark.parametrize("source", ["configured", "caller", "forwarded"])
async def test_raw_authorization_rejects_bare_schemes_before_dispatch(
self, respx_mock: MockRouter, value: str, source: str,
) -> None:
server: Final = MCPServer(
server_id="raw-empty", name="raw-empty", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.authorization,
authentication_token=value if source == "configured" else None,
)
destination: Final = respx_mock.route().respond(200)
with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc:
await MCPServerManager()._create_mcp_client(
server, mcp_auth_header=value if source == "caller" else None,
extra_headers={"Authorization": value} if source == "forwarded" else None,
)
assert exc.value.status_code == 500
assert destination.call_count == 0
@pytest.mark.asyncio
async def test_byok_flag_cannot_bypass_incomplete_obo(self) -> None:
server = MCPServer(server_id="obo-byok", name="obo-byok", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.oauth2_token_exchange, is_byok=True,
token_exchange_endpoint="https://idp.example/token")
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(server, mcp_auth_header="Bearer override")
assert exc.value.status_code == 401
@pytest.mark.asyncio
@pytest.mark.parametrize("configured,override", [(None, "Bearer usable"), ("shared", "Bearer usable")])
async def test_bearer_override_remains_usable(self, configured: str | None, override: str) -> None:
server = MCPServer(server_id="override", name="override", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=configured)
client = await MCPServerManager()._create_mcp_client(server, mcp_auth_header=override)
assert client._get_auth_headers()["Authorization"] == override
@pytest.mark.asyncio
@pytest.mark.parametrize("token", [None, "shared"])
async def test_empty_injected_header_cannot_satisfy_protected_auth(self, token: str | None) -> None:
server = MCPServer(server_id="empty-header", name="empty-header", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=token)
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(server, extra_headers={"authorization": " "})
assert exc.value.status_code == 500
@pytest.mark.asyncio
async def test_custom_slot_uses_its_actual_credential(self) -> None:
server = MCPServer(server_id="custom", name="custom", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.api_key,
upstream_token_header="X-Custom", authentication_token="key")
client = await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Trace": "trace"})
assert client._credential_slot == "X-Custom"
assert await client.discovery_auth_fingerprint()
@pytest.mark.asyncio
@pytest.mark.parametrize("static,forwarded,caller", [
({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None),
({}, {"X-API-Key": "forwarded"}, None),
({}, None, "ApiKey caller"),
({"X-API-Key": "static"}, {"Authorization": ""}, None),
])
async def test_openapi_static_credentials_remain_supported(
self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch,
static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None
) -> None:
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
_request_auth_header, _request_extra_headers, create_tool_function,
)
tool: Final = create_tool_function(
"/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.api_key,
)
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated")
caller_token: Final = _request_auth_header.set(caller)
extra_token: Final = _request_extra_headers.set(forwarded)
try:
assert await tool() == "authenticated"
sent: Final = destination.calls.last.request.headers
assert sent.get("x-api-key") == static.get("X-API-Key", (forwarded or {}).get("X-API-Key"))
if caller:
assert sent["authorization"] == caller
assert destination.call_count == 1
finally:
_request_auth_header.reset(caller_token)
_request_extra_headers.reset(extra_token)
@pytest.mark.asyncio
async def test_static_resolution_cancellation_closes_flow(self) -> None:
from collections.abc import AsyncGenerator
from litellm.experimental_mcp_client.client import MCPClient
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import prepare_mcp_client
class CancelledAuth(httpx.Auth):
closed = False
async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
try:
raise asyncio.CancelledError()
yield request
finally:
self.closed = True
auth = CancelledAuth()
server = MCPServer(server_id="cancel", name="cancel", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.api_key)
client = MCPClient(server_url=server.url, auth_type=MCPAuth.api_key, resolved_auth=auth)
with pytest.raises(asyncio.CancelledError):
await prepare_mcp_client(server, client)
assert auth.closed
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type", [MCPAuth.basic, MCPAuth.token, MCPAuth.authorization])
async def test_other_static_schemes_reject_whitespace_credentials(self, auth_type: MCPAuthType) -> None:
server = MCPServer(server_id="blank-static", name="blank-static", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=auth_type, authentication_token=" ")
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(server)
assert exc.value.status_code == 500
@pytest.mark.asyncio
@pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc", "Basic QmFzaWM=", "Basic bm8tY29sb24="])
async def test_basic_headers_without_usable_credentials_reject(self, header: str) -> None:
server = MCPServer(server_id="bad-basic", name="bad-basic", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.basic)
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(server, extra_headers={"Authorization": header})
assert exc.value.status_code == 500
@pytest.mark.asyncio
@pytest.mark.parametrize("value", ["Basic", "Basic ", "basic"])
@pytest.mark.parametrize("source", ["configured", "caller"])
async def test_basic_scheme_alone_is_not_a_credential(self, value: str, source: str) -> None:
server = MCPServer(server_id="basic-scheme", name="basic-scheme", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.basic,
authentication_token=value if source == "configured" else None)
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None)
assert exc.value.status_code == 500
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type,value,default_slot", [
(MCPAuth.api_key, "fixture-key", "X-API-Key"),
(MCPAuth.bearer_token, "fixture-key", "Authorization"),
(MCPAuth.basic, "user:pass", "Authorization"),
(MCPAuth.token, "fixture-key", "Authorization"),
(MCPAuth.authorization, "fixture-key", "Authorization"),
])
@pytest.mark.parametrize("source", ["configured", "caller"])
async def test_usable_credential_survives_an_empty_alternate_header(
self, auth_type: MCPAuthType, value: str, default_slot: str, source: str
) -> None:
server: Final = MCPServer(
server_id="alternate", name="alternate", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=auth_type, upstream_token_header="X-Custom",
authentication_token=value if source == "configured" else None,
)
empty_slot: Final = default_slot if source == "configured" else "X-Custom"
selected_slot: Final = "X-Custom" if source == "configured" else default_slot
client: Final = await MCPServerManager()._create_mcp_client(
server, mcp_auth_header=value if source == "caller" else None, extra_headers={empty_slot: ""},
)
request: Final = await client.prepare_request_auth()
assert request.headers[selected_slot]
assert request.headers[empty_slot] == ""
@pytest.mark.asyncio
async def test_empty_custom_and_default_headers_do_not_satisfy_auth(self) -> None:
server: Final = MCPServer(
server_id="both-empty", name="both-empty", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.api_key, upstream_token_header="X-Custom",
)
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Custom": "", "X-API-Key": ""})
assert exc.value.status_code == 500
@pytest.mark.asyncio
@pytest.mark.parametrize("custom_slot", [None, "X-Custom"])
@pytest.mark.parametrize("source", ["caller", "forwarded"])
async def test_api_key_preserves_explicit_authorization_credential(
self, custom_slot: str | None, source: str
) -> None:
server: Final = MCPServer(
server_id="caller-auth", name="caller-auth", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.api_key, upstream_token_header=custom_slot,
)
headers: Final = {"Authorization": "Bearer caller-credential", "X-API-Key": ""}
client: Final = await MCPServerManager()._create_mcp_client(
server, mcp_auth_header=headers if source == "caller" else None,
extra_headers=headers if source == "forwarded" else None,
)
request: Final = await client.prepare_request_auth()
assert request.headers["Authorization"] == "Bearer caller-credential"
assert request.headers["X-API-Key"] == ""
assert custom_slot is None or custom_slot not in request.headers
@pytest.mark.asyncio
@pytest.mark.parametrize("value", [
"", " ", "Bearer", "Basic", "token", "ApiKey",
"Bearer Bearer", "ApiKey ApiKey", "token token", "bEaReR BEARER", "aPiKeY\tAPIKEY",
])
async def test_api_key_rejects_authorization_without_a_credential(self, value: str) -> None:
server: Final = MCPServer(
server_id="caller-empty", name="caller-empty", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.api_key,
)
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(server, mcp_auth_header={"Authorization": value})
assert exc.value.status_code == 500
@pytest.mark.asyncio
@pytest.mark.parametrize("value", ["no-colon", "Basic bm8tY29sb24="])
@pytest.mark.parametrize("source", ["configured", "caller"])
async def test_basic_requires_a_username_password_separator(self, value: str, source: str) -> None:
server: Final = MCPServer(
server_id="basic-pair", name="basic-pair", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.basic,
authentication_token=value if source == "configured" else None,
)
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None)
assert exc.value.status_code == 500
@pytest.mark.asyncio
@pytest.mark.parametrize("value", ["user:pass", "user:", ":pass", ":"])
async def test_basic_preserves_username_password_pairs(self, value: str) -> None:
import base64
server: Final = MCPServer(
server_id="basic-valid", name="basic-valid", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.basic, authentication_token=value,
)
client: Final = await MCPServerManager()._create_mcp_client(server)
request: Final = await client.prepare_request_auth()
scheme, encoded = request.headers["Authorization"].split(" ", 1)
assert scheme == "Basic"
assert base64.b64decode(encoded) == value.encode()
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type,value", [
(MCPAuth.bearer_token, "Bearer"), (MCPAuth.bearer_token, "Bearer "), (MCPAuth.bearer_token, "bearer"),
(MCPAuth.token, "token"), (MCPAuth.token, "token "), (MCPAuth.token, "TOKEN"),
])
@pytest.mark.parametrize("source", ["configured", "caller"])
async def test_static_scheme_only_input_cannot_hide_behind_rendered_prefix(
self, auth_type: MCPAuthType, value: str, source: str
) -> None:
server: Final = MCPServer(
server_id="empty-scheme", name="empty-scheme", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=auth_type,
authentication_token=value if source == "configured" else None,
)
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None)
assert exc.value.status_code == 500
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type,value,expected", [
(MCPAuth.bearer_token, "token", "Bearer token"),
(MCPAuth.bearer_token, "Bearertoken", "Bearer Bearertoken"),
(MCPAuth.token, "tokenish", "token tokenish"),
])
async def test_static_credentials_that_resemble_schemes_remain_usable(
self, auth_type: MCPAuthType, value: str, expected: str
) -> None:
server: Final = MCPServer(
server_id="real-token", name="real-token", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=auth_type, authentication_token=value,
)
client: Final = await MCPServerManager()._create_mcp_client(server)
request: Final = await client.prepare_request_auth()
assert request.headers["Authorization"] == expected

View file

@ -10,9 +10,14 @@ This test suite ensures that:
"""
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, patch
import pytest
from fastapi import HTTPException
from respx import MockRouter
from litellm.types.mcp import MCPAuth, MCPAuthType
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
_request_auth_header,
@ -35,6 +40,120 @@ from litellm.proxy._experimental.mcp_server.exceptions import (
GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.get_async_httpx_client"
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type,value,accepted", [
(MCPAuth.api_key, "Bearer Bearer", False), (MCPAuth.api_key, "ApiKey ApiKey", False),
(MCPAuth.api_key, "token token", False), (MCPAuth.api_key, "bEaReR BEARER", False),
(MCPAuth.api_key, "aPiKeY\tAPIKEY", False), (MCPAuth.api_key, "Bearer fixture-key", True),
(MCPAuth.api_key, "ApiKey fixture-key", True), (MCPAuth.api_key, "token fixture-key", True),
(MCPAuth.authorization, "Bearer", False), (MCPAuth.authorization, "basic", False),
(MCPAuth.authorization, "token", False), (MCPAuth.authorization, "ApiKey", False),
(MCPAuth.authorization, " bEaReR ", False), (MCPAuth.authorization, "\tTOKEN\t", False),
(MCPAuth.authorization, "opaque-secret-value", True), (MCPAuth.authorization, "Bearer abc", True),
(MCPAuth.authorization, "Custom abc", True),
])
async def test_authorization_validates_credentials_before_http(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, auth_type: MCPAuthType, value: str, accepted: bool,
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
tool: Final = create_tool_function(
"/echo", "get", {}, "https://upstream.example", auth_type=auth_type,
)
destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated")
caller_token: Final = _request_auth_header.set(value)
try:
if accepted:
assert await tool() == "authenticated"
assert destination.call_count == 1
assert destination.calls.last.request.headers["authorization"] == value
else:
with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc:
await tool()
assert exc.value.status_code == 500
assert destination.call_count == 0
finally:
_request_auth_header.reset(caller_token)
@pytest.mark.asyncio
@pytest.mark.parametrize("static,forwarded,caller,resolved,expected", [
({"Authorization": "Bearer configured"}, {"authorization": "Bearer forwarded"}, None, None, "Bearer configured"),
({"Authorization": "Bearer configured"}, None, "Bearer caller", None, "Bearer caller"),
({"Authorization": "Bearer configured"}, None, "Bearer", None, None),
({"Authorization": "Bearer configured"}, None, "Bearer caller", {"authorization": " "}, None),
({"Authorization": "Bearer configured"}, None, "Bearer", {"authorization": "Bearer resolved"}, "Bearer resolved"),
])
async def test_static_auth_validates_headers_after_existing_precedence(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch,
static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None,
resolved: dict[str, str] | None, expected: str | None,
) -> None:
tool: Final = create_tool_function(
"/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.bearer_token,
)
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated")
caller_token: Final = _request_auth_header.set(caller)
extra_token: Final = _request_extra_headers.set(forwarded)
resolved_token: Final = _request_resolved_auth_headers.set(resolved)
try:
if expected is None:
with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc:
await tool()
assert exc.value.status_code == 500
assert destination.call_count == 0
else:
assert await tool() == "authenticated"
assert destination.call_count == 1
assert destination.calls.last.request.headers["authorization"] == expected
finally:
_request_auth_header.reset(caller_token)
_request_extra_headers.reset(extra_token)
_request_resolved_auth_headers.reset(resolved_token)
@pytest.mark.asyncio
@pytest.mark.parametrize("credential", ["custom-key", ""])
async def test_static_auth_uses_configured_custom_header(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, credential: str,
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
tool: Final = create_tool_function(
"/echo", "get", {}, "https://upstream.example", headers={"x-custom": credential},
auth_type=MCPAuth.api_key, upstream_token_header="X-Custom",
)
destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated")
if credential:
assert await tool() == "authenticated"
assert destination.call_count == 1
assert destination.calls.last.request.headers["x-custom"] == credential
else:
with pytest.raises(HTTPException, match="requires a usable upstream credential"):
await tool()
assert destination.call_count == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type,resolved", [
(MCPAuth.none, None),
(MCPAuth.oauth2, {"Authorization": "Bearer user-oauth"}),
])
async def test_static_validation_preserves_no_auth_and_resolved_oauth(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch,
auth_type: MCPAuthType, resolved: dict[str, str] | None,
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
tool: Final = create_tool_function("/echo", "get", {}, "https://upstream.example", auth_type=auth_type)
destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="echo")
token: Final = _request_resolved_auth_headers.set(resolved)
try:
assert await tool() == "echo"
assert destination.call_count == 1
assert destination.calls.last.request.headers.get("authorization") == (resolved or {}).get("Authorization")
finally:
_request_resolved_auth_headers.reset(token)
def _create_mock_client(method: str, response_text: str, status_code: int = 200) -> AsyncMock:
"""Utility to create a mocked async httpx client for the given method.
@ -1458,3 +1577,21 @@ class TestBoundedOpenAPISpecLoading:
else:
assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) == {"paths": {}}
assert destination.call_count == 1
def test_openapi_generator_import_does_not_require_mcp_sdk() -> None:
import subprocess
import sys
script = """
import builtins
original_import = builtins.__import__
def without_mcp(name, *args, **kwargs):
if name == 'mcp' or name.startswith('mcp.'):
raise ModuleNotFoundError('MCP SDK unavailable')
return original_import(name, *args, **kwargs)
builtins.__import__ = without_mcp
import litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator
"""
result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True)
assert result.returncode == 0, result.stderr

View file

@ -3,12 +3,14 @@ Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objec
"""
import json
import logging
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi.testclient import TestClient
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@ -285,6 +287,115 @@ class TestFailureHookRequestData:
assert hook_request_data["litellm_logging_obj"] == "logging-obj-sentinel"
class TestErrorLogCarriesCallId:
"""LIT-7836: the /v1/messages and /v1/messages/count_tokens error lines must carry
the request's litellm_call_id, rendered in the message and as a structured field."""
@pytest.fixture(autouse=True)
def propagating_proxy_logger(self):
verbose_proxy_logger.propagate = True
try:
yield
finally:
verbose_proxy_logger.propagate = False
@staticmethod
def _error_record(caplog: pytest.LogCaptureFixture) -> logging.LogRecord:
return next(r for r in caplog.records if "Exception occured" in r.getMessage())
@pytest.mark.asyncio
async def test_messages_failure_log_carries_call_id(self, caplog: pytest.LogCaptureFixture):
import litellm.proxy.anthropic_endpoints.endpoints as ep
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import UserAPIKeyAuth
call_id = "messages-call-7836"
async def fake_process(self, **kwargs):
self.data = {**self.data, "litellm_call_id": call_id}
raise RuntimeError("provider timeout")
request = MagicMock()
request.headers = {}
with (
patch.object(ep, "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam
patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the provider failure happens inside this call; the test targets the endpoint's except block
patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
):
mock_logging.post_call_failure_hook = AsyncMock()
response = await ep.anthropic_response(
fastapi_response=MagicMock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(),
)
assert response.status_code == 500
record = self._error_record(caplog)
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
@pytest.mark.asyncio
async def test_messages_already_shaped_failure_answers_with_the_call_id(self):
import litellm.proxy.anthropic_endpoints.endpoints as ep
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
call_id = "messages-call-7836-shaped"
async def fake_process(self, **kwargs):
self.data = {**self.data, "litellm_call_id": call_id}
raise ProxyException(message="budget exceeded", type=ProxyErrorTypes.budget_exceeded, param="key", code=402)
request = MagicMock()
request.headers = {}
with (
patch.object(ep, "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam
patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the proxy shaped failure happens inside this call; the test targets the endpoint's except block
patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
):
mock_logging.post_call_failure_hook = AsyncMock()
response = await ep.anthropic_response(
fastapi_response=MagicMock(),
request=request,
user_api_key_dict=UserAPIKeyAuth(),
)
assert response.status_code == 402
assert response.headers["x-litellm-call-id"] == call_id
@pytest.mark.asyncio
async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture):
from fastapi import HTTPException
import litellm.proxy.anthropic_endpoints.endpoints as ep
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import UserAPIKeyAuth
call_id = "count-tokens-call-7836"
request = MagicMock()
request.headers = {"x-litellm-call-id": call_id}
with (
patch.object( # test-quality-ok: endpoint reads the body via a module function; no injection seam
ep,
"_read_request_body",
new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}),
),
patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
pytest.raises(HTTPException) as raised,
):
await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
assert raised.value.status_code == 500
record = self._error_record(caplog)
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
class TestEventLoggingBatchEndpoint:
"""Test the stubbed event logging batch endpoint"""

View file

@ -2892,45 +2892,49 @@ def test_team_update_gate_allows_org_admin_with_resolved_org():
)
def test_team_update_gate_rejects_without_org_context():
"""Without organization_id (i.e. resolution found no org, or a non-org-admin),
the gate still rejects /team/update the fix adds no blanket allow. Guards
against re-widening the route (e.g. dropping it into self_managed_routes)."""
def test_team_update_gate_admits_internal_user_without_org_context(): # test-quality-ok: the gate's only success signal is not raising; the handler's team-admin 403s are pinned in test_team_endpoints
"""/team/update is self-managed (LIT-5722): the coarse gate admits any authenticated
caller and update_team resolves proxy, org or team admin itself, then filters team admins
through the team_admin_editable_team_fields setting. Before that the gate 401'd every
team admin, which left the handler's team-admin branch unreachable."""
user_obj = LiteLLM_UserTable(
user_id="team-admin-user",
user_role=LitellmUserRoles.INTERNAL_USER.value,
organization_memberships=None,
)
valid_token = UserAPIKeyAuth(user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value)
request = MagicMock(spec=Request)
request.method = "POST"
request.query_params = {}
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
route="/team/update",
request=request,
valid_token=valid_token,
request_data={"team_id": "team-1", "max_budget": 42},
)
def test_team_update_gate_defers_cross_org_admin_to_the_handler(): # test-quality-ok: the gate's only success signal is not raising; the handler's 403 it defers to is pinned in test_team_endpoints
"""An org admin of a DIFFERENT org clears the coarse gate like any internal user;
update_team's _resolve_team_access finds no role on the team and 403s (pinned in
test_team_endpoints), so there is still no cross-org escalation."""
user_obj = _make_org_admin_user("org-1")
valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value)
request = MagicMock(spec=Request)
request.method = "POST"
request.query_params = {}
with pytest.raises(Exception, match="Only proxy admin can be used to generate"):
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
route="/team/update",
request=request,
valid_token=valid_token,
request_data={"team_id": "team-1", "max_budget": 42},
)
def test_team_update_gate_rejects_cross_org_admin_with_resolved_org():
"""Even after the target team's org is resolved, an org admin of a DIFFERENT
org is rejected at the gate (no cross-org escalation)."""
user_obj = _make_org_admin_user("org-1")
valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value)
request = MagicMock(spec=Request)
request.method = "POST"
request.query_params = {}
with pytest.raises(Exception, match="Only proxy admin can be used to generate"):
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
route="/team/update",
request=request,
valid_token=valid_token,
request_data={"team_id": "team-1", "organization_id": "org-2"},
)
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
route="/team/update",
request=request,
valid_token=valid_token,
request_data={"team_id": "team-1", "organization_id": "org-2"},
)
# ── PATCH /team/{team_id}: same org-context + role reach as POST /team/update ──
@ -2993,23 +2997,6 @@ async def test_add_team_org_context_noop_for_static_team_route():
assert out == body
def test_patch_team_route_has_same_reach_as_team_update():
"""/team/{team_id} is reachable by org admins (in org_admin_allowed_routes) but
NOT by regular internal users or the role-agnostic self_managed_routes the
latter would open /team/new (the collision footgun) to any authenticated user."""
from litellm.proxy._types import LiteLLMRoutes
assert RouteChecks.check_route_access(
route="/team/abc-123", allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value
)
assert not RouteChecks.check_route_access(
route="/team/abc-123", allowed_routes=LiteLLMRoutes.internal_user_routes.value
)
assert not RouteChecks.check_route_access(
route="/team/abc-123", allowed_routes=LiteLLMRoutes.self_managed_routes.value
)
def _patch_team_request() -> MagicMock:
request = MagicMock(spec=Request)
request.method = "PATCH"
@ -3897,7 +3884,6 @@ def test_team_disable_logging_stays_proxy_admin_only():
"route",
[
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112",
"/team/update",
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/model/add",
],
)

View file

@ -31,6 +31,7 @@ cannot drift without a test failure.
import base64
import json
import logging
from contextlib import ExitStack
from dataclasses import dataclass
from typing import Any, Dict, Optional
@ -1088,6 +1089,28 @@ async def test_create__exception_calls_failure_hook(harness, openai_env_creds):
assert harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom"
async def test_create__exception_carries_the_litellm_call_id(harness, openai_env_creds, caplog):
call_id = "lit7836-batch-call-id"
set_body(
harness,
{
"input_file_id": "file-plain",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"litellm_call_id": call_id,
},
)
harness.litellm_acreate.side_effect = ValueError("provider boom")
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised:
await call_create(harness)
assert raised.value.headers["x-litellm-call-id"] == call_id
record = next(r for r in caplog.records if "Exception occured" in r.getMessage())
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
# =========================================================================== #
# #
# GET /v1/batches/{batch_id} - retrieve_batch routing-contract tests #
@ -1953,6 +1976,24 @@ async def test_list__exception_calls_failure_hook(list_harness):
assert list_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom"
@pytest.mark.asyncio
async def test_list__failure_hook_and_response_share_the_request_litellm_call_id(list_harness):
call_id = "lit7836-list-batches-call-id"
list_harness.pre_call.side_effect = lambda **kw: (
{**list_harness.body["body"], "litellm_call_id": call_id},
MagicMock(),
)
list_harness.litellm_alist.side_effect = ValueError("provider boom")
with pytest.raises(ProxyException) as raised:
await call_list(list_harness, after="batch-0", limit=5)
failure_request_data = list_harness.logging.post_call_failure_hook.call_args.kwargs["request_data"]
assert failure_request_data["litellm_call_id"] == call_id
assert (failure_request_data["after"], failure_request_data["limit"]) == ("batch-0", 5)
assert raised.value.headers["x-litellm-call-id"] == call_id
# =========================================================================== #
# #
# POST /v1/batches/{batch_id}/cancel - cancel_batch routing-contract tests #

View file

@ -6,8 +6,10 @@ from fastapi import HTTPException
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
litellm_call_id_headers,
openai_error_param,
openai_error_type,
with_litellm_call_id,
)
@ -158,3 +160,32 @@ def test_a_stringified_none_type_or_param_is_treated_as_absent():
assert carried.type == "None"
assert openai_error_type(carried, 400) == "invalid_request_error"
assert openai_error_param(carried) is None
def test_a_failed_request_answers_with_the_call_id_it_was_logged_under():
assert litellm_call_id_headers("call-7836") == {"x-litellm-call-id": "call-7836"}
assert litellm_call_id_headers(None) is None
def test_an_already_shaped_proxy_error_answers_with_the_call_id_it_was_logged_under():
raised_without_id = ProxyException(message="budget exceeded", type="budget_exceeded", param="key", code=402)
carried = with_litellm_call_id(raised_without_id, "call-7836")
assert carried is raised_without_id
assert carried.headers == {"x-litellm-call-id": "call-7836"}
assert (carried.message, carried.type, carried.param, carried.code) == (
"budget exceeded",
"budget_exceeded",
"key",
"402",
)
def test_a_proxy_error_keeps_the_call_id_it_was_raised_with():
raised_with_id = ProxyException(
message="nope", type="None", param=None, code=400, headers={"x-litellm-call-id": "first"}
)
assert with_litellm_call_id(raised_with_id, "second").headers == {"x-litellm-call-id": "first"}
assert with_litellm_call_id(ProxyException(message="nope", type="None", param=None, code=400), None).headers == {}

View file

@ -1,5 +1,7 @@
import asyncio
import copy
import logging
from collections.abc import Iterator, Mapping
from types import SimpleNamespace
from typing import Any, Dict
@ -10,6 +12,7 @@ from fastapi.testclient import TestClient
from starlette.requests import Request
from starlette.responses import Response
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.image_endpoints import endpoints
@ -211,3 +214,117 @@ async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(mon
await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth())
assert (raised.value.type, raised.value.param, raised.value.code) == ("invalid_request_error", None, "404")
@pytest.fixture
def propagating_proxy_logger() -> Iterator[None]:
verbose_proxy_logger.propagate = True
try:
yield
finally:
verbose_proxy_logger.propagate = False
@pytest.mark.asyncio
async def test_failure_log_carries_the_callers_litellm_call_id(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, propagating_proxy_logger: None
) -> None:
"""LIT-7836: the /v1/images/generations error line must carry the litellm_call_id
the client sent, both rendered in the message and as a structured record field."""
call_id = "images-call-7836"
async def fake_add_litellm_data_to_request(**kwargs: object) -> object:
return kwargs["data"]
async def fake_pre_call_hook(*, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> dict[str, object]:
return data
async def fake_post_call_failure_hook(**_: object) -> None:
return None
async def failing_route_request(**_: object) -> None:
raise HTTPException(status_code=401, detail={"error": "invalid api key"})
monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {})
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj",
SimpleNamespace(pre_call_hook=fake_pre_call_hook, post_call_failure_hook=fake_post_call_failure_hook),
)
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version")
monkeypatch.setattr("litellm.proxy.image_endpoints.endpoints.route_request", failing_route_request)
body = orjson.dumps({"model": "dall-e-3", "prompt": "a lighthouse at dusk"})
async def receive() -> dict[str, object]:
return {"type": "http.request", "body": body, "more_body": False}
request = Request(
{
"type": "http",
"method": "POST",
"path": "/v1/images/generations",
"headers": [(b"x-litellm-call-id", call_id.encode())],
},
receive,
)
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised:
await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth())
assert raised.value.headers["x-litellm-call-id"] == call_id
record = next(r for r in caplog.records if "Exception occured" in r.getMessage())
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
@pytest.mark.asyncio
async def test_failure_before_the_provider_call_bills_the_callers_litellm_call_id(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""LIT-7836: when the request is rejected while it is still being prepared, the
failure hook must see the same litellm_call_id the response header answers with,
otherwise the spend row is stored under a freshly minted id nobody can look up."""
call_id = "images-early-7836"
hook_request_data: list[Mapping[str, object]] = []
async def rejecting_add_litellm_data_to_request(**_: object) -> object:
raise HTTPException(status_code=400, detail={"error": "tag not allowed"})
async def fake_post_call_failure_hook(*, request_data: Mapping[str, object], **_: object) -> None:
hook_request_data.append(request_data)
monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", rejecting_add_litellm_data_to_request)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {})
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj",
SimpleNamespace(post_call_failure_hook=fake_post_call_failure_hook),
)
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version")
body = orjson.dumps({"model": "dall-e-3", "prompt": "a lighthouse at dusk", "litellm_call_id": "from-the-body"})
async def receive() -> dict[str, object]:
return {"type": "http.request", "body": body, "more_body": False}
request = Request(
{
"type": "http",
"method": "POST",
"path": "/v1/images/generations",
"headers": [(b"x-litellm-call-id", call_id.encode())],
},
receive,
)
with pytest.raises(ProxyException) as raised:
await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth())
assert raised.value.headers["x-litellm-call-id"] == call_id
assert [data["litellm_call_id"] for data in hook_request_data] == [call_id]

View file

@ -0,0 +1,138 @@
import pytest
from fastapi import HTTPException
from litellm.proxy._types import LiteLLM_ModelTable, LiteLLM_TeamTable, UpdateTeamRequest
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
TeamAdminEditAllowed,
TeamAdminEditingDisabled,
TeamAdminFieldNotPermitted,
changed_team_fields,
resolve_team_admin_editable_fields,
team_admin_edit_verdict,
team_admin_request_or_raise,
)
_SUPPORTED = frozenset({"tpm_limit", "rpm_limit", "team_alias"})
def _team(**overrides):
return LiteLLM_TeamTable(team_id="team-1", **overrides)
class TestResolveTeamAdminEditableFields:
def test_missing_setting_means_nothing_editable(self):
assert resolve_team_admin_editable_fields({}, _SUPPORTED) == frozenset()
def test_keeps_only_supported_names(self):
configured = {"team_admin_editable_team_fields": ["tpm_limit", "blocked", "organization_id"]}
assert resolve_team_admin_editable_fields(configured, _SUPPORTED) == frozenset({"tpm_limit"})
@pytest.mark.parametrize("raw", ["tpm_limit", 7, {"tpm_limit": True}, [1, 2]])
def test_malformed_setting_fails_closed(self, raw):
assert resolve_team_admin_editable_fields({"team_admin_editable_team_fields": raw}, _SUPPORTED) == frozenset()
class TestChangedTeamFields:
def test_team_id_alone_changes_nothing(self):
assert changed_team_fields(UpdateTeamRequest(team_id="team-1"), _team()) == frozenset()
def test_column_echoing_stored_value_is_not_a_change(self):
data = UpdateTeamRequest(team_id="team-1", tpm_limit=5, team_alias="alpha", max_budget=None)
assert changed_team_fields(data, _team(tpm_limit=5, team_alias="alpha")) == frozenset()
def test_column_with_different_value_is_a_change(self):
data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha")
assert changed_team_fields(data, _team(tpm_limit=5, team_alias="alpha")) == frozenset({"tpm_limit"})
def test_explicit_null_clearing_a_stored_column_is_a_change(self):
data = UpdateTeamRequest(team_id="team-1", max_budget=None)
assert changed_team_fields(data, _team(max_budget=30.0)) == frozenset({"max_budget"})
def test_folded_field_sent_top_level_is_named_not_metadata(self):
data = UpdateTeamRequest(team_id="team-1", guardrails=["b"])
assert changed_team_fields(data, _team(metadata={"guardrails": ["a"]})) == frozenset({"guardrails"})
def test_folded_field_sent_inside_metadata_is_named_not_metadata(self):
data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["b"]})
assert changed_team_fields(data, _team(metadata={"guardrails": ["a"]})) == frozenset({"guardrails"})
def test_custom_metadata_key_change_is_attributed_to_metadata(self):
data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["a"], "cost_center": "b"})
existing = _team(metadata={"guardrails": ["a"], "cost_center": "a"})
assert changed_team_fields(data, existing) == frozenset({"metadata"})
def test_metadata_echo_with_top_level_override_only_names_the_override(self):
data = UpdateTeamRequest(team_id="team-1", guardrails=["b"], metadata={"guardrails": ["a"], "cost_center": "a"})
existing = _team(metadata={"guardrails": ["a"], "cost_center": "a"})
assert changed_team_fields(data, existing) == frozenset({"guardrails"})
def test_dropping_a_stored_key_from_submitted_metadata_is_a_change(self):
data = UpdateTeamRequest(team_id="team-1", metadata={"cost_center": "a"})
existing = _team(metadata={"cost_center": "a", "tags": ["x"], "logging": [{"callback": "langfuse"}]})
assert changed_team_fields(data, existing) == frozenset({"tags", "logging"})
def test_server_managed_metadata_key_is_ignored(self):
data = UpdateTeamRequest(team_id="team-1", metadata={"cost_center": "a"})
existing = _team(metadata={"cost_center": "a", "team_member_budget_id": "budget-1"})
assert changed_team_fields(data, existing) == frozenset()
def test_model_aliases_compare_against_the_model_table(self):
table = LiteLLM_ModelTable(model_aliases='{"fast": "gpt-4o-mini"}', created_by="a", updated_by="a")
same = UpdateTeamRequest(team_id="team-1", model_aliases={"fast": "gpt-4o-mini"})
different = UpdateTeamRequest(team_id="team-1", model_aliases={"fast": "gpt-4o"})
assert changed_team_fields(same, _team(litellm_model_table=table)) == frozenset()
assert changed_team_fields(different, _team(litellm_model_table=table)) == frozenset({"model_aliases"})
def test_empty_model_aliases_against_no_model_table_is_not_a_change(self):
assert changed_team_fields(UpdateTeamRequest(team_id="team-1", model_aliases={}), _team()) == frozenset()
def test_field_without_a_stored_counterpart_counts_as_changed_when_sent(self):
data = UpdateTeamRequest(team_id="team-1", team_member_budget=10.0)
assert changed_team_fields(data, _team()) == frozenset({"team_member_budget"})
class TestTeamAdminEditVerdict:
def test_no_permitted_fields_disables_editing_even_for_a_no_op(self):
verdict = team_admin_edit_verdict(UpdateTeamRequest(team_id="team-1"), _team(), frozenset())
assert verdict == TeamAdminEditingDisabled()
def test_allowed_request_keeps_only_the_changed_fields(self):
data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha", budget_duration="30d")
existing = _team(team_alias="alpha", budget_duration="30d")
verdict = team_admin_edit_verdict(data, existing, frozenset({"tpm_limit"}))
assert isinstance(verdict, TeamAdminEditAllowed)
assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "tpm_limit": 6}
def test_permitted_field_changed_inside_metadata_keeps_the_metadata(self):
data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["b"]}, team_alias="alpha")
existing = _team(team_alias="alpha", metadata={"guardrails": ["a"]})
verdict = team_admin_edit_verdict(data, existing, frozenset({"guardrails"}))
assert isinstance(verdict, TeamAdminEditAllowed)
assert verdict.request.model_dump(exclude_unset=True) == {
"team_id": "team-1",
"metadata": {"guardrails": ["b"]},
}
def test_first_blocked_field_in_sorted_order_is_reported(self):
data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, rpm_limit=6, blocked=True)
verdict = team_admin_edit_verdict(data, _team(), frozenset({"tpm_limit"}))
assert verdict == TeamAdminFieldNotPermitted(field="blocked")
class TestTeamAdminRequestOrRaise:
def test_allowed_hands_back_its_request(self):
request = UpdateTeamRequest(team_id="team-1", tpm_limit=6)
assert team_admin_request_or_raise(TeamAdminEditAllowed(request=request)) is request
def test_disabled_is_a_403_pointing_at_the_proxy_admin(self):
with pytest.raises(HTTPException) as exc:
team_admin_request_or_raise(TeamAdminEditingDisabled())
assert exc.value.status_code == 403
assert "cannot edit team settings" in exc.value.detail
assert "Settings > UI > Team admin editable fields" in exc.value.detail
def test_field_not_permitted_is_a_403_naming_the_field(self):
with pytest.raises(HTTPException) as exc:
team_admin_request_or_raise(TeamAdminFieldNotPermitted(field="blocked"))
assert exc.value.status_code == 403
assert "'blocked'" in exc.value.detail

View file

@ -1,6 +1,6 @@
import asyncio
import json
from contextlib import asynccontextmanager
from contextlib import asynccontextmanager, contextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Final, Optional, cast
@ -76,6 +76,31 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
client = TestClient(app)
@contextmanager
def _team_admin_may_edit(*fields: str):
"""Let team admins change ``fields`` on /team/update for the duration of the block.
The registry only lists the fields shipped so far (LIT-5722 adds them one PR at a time), so tests that
exercise the gates layered underneath the allow-list widen it here instead of asserting the early 403."""
with (
patch( # test-quality-ok: the registry is a module constant update_team reads directly; no seam to inject
"litellm.proxy.management_endpoints.team_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS",
frozenset(fields),
),
patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": list(fields)}), # test-quality-ok: update_team reads general_settings as a proxy_server module global
):
yield
def _not_org_admin():
"""update_team asks whether the caller administers the team's org before it settles for team admin;
a MagicMock prisma cannot answer that lookup, so pin it to False."""
return patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide
"litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team",
AsyncMock(return_value=False),
)
def _wire_team_create_tx(prisma_client):
"""`/team/new` inserts the team and mirrors it onto the access groups in one transaction,
so a mocked client has to hand its team table back out of `db.tx()`.
@ -6393,6 +6418,7 @@ async def test_update_team_standalone_budget_raise_blocked_for_team_admin():
dummy_request = MagicMock(spec=Request)
with (
_team_admin_may_edit("max_budget"),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -6549,6 +6575,7 @@ async def test_update_team_standalone_budget_removal_blocked_for_team_admin():
dummy_request = MagicMock(spec=Request)
with (
_team_admin_may_edit("max_budget"),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -6618,6 +6645,7 @@ async def test_update_team_standalone_uncapped_team_admin_sets_finite_allowed(
dummy_request = MagicMock(spec=Request)
with (
_team_admin_may_edit("max_budget"),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -6712,6 +6740,7 @@ async def test_update_team_standalone_unchanged_budget_allowed(
dummy_request = MagicMock(spec=Request)
with (
_team_admin_may_edit("max_budget", "tpm_limit"),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -6810,6 +6839,7 @@ async def test_update_team_standalone_lower_budget_allowed(
dummy_request = MagicMock(spec=Request)
with (
_team_admin_may_edit("max_budget"),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -6912,6 +6942,8 @@ async def test_update_team_org_scoped_budget_exceeds_org_limit():
mock_org.litellm_budget_table = mock_budget_table
with (
_team_admin_may_edit("max_budget"),
_not_org_admin(),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -6992,6 +7024,7 @@ async def test_update_team_standalone_models_not_gated_by_user_limit(
dummy_request = MagicMock(spec=Request)
with (
_team_admin_may_edit("models"),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -7091,6 +7124,8 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit(
mock_org.litellm_budget_table = mock_budget_table
with (
_team_admin_may_edit("max_budget"),
_not_org_admin(),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -7202,6 +7237,8 @@ async def test_update_team_org_scoped_models_bypasses_user_limit(
mock_org.litellm_budget_table = None
with (
_team_admin_may_edit("models"),
_not_org_admin(),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -7304,6 +7341,8 @@ async def test_update_team_org_scoped_models_not_in_org_models():
mock_org.litellm_budget_table = None
with (
_team_admin_may_edit("models"),
_not_org_admin(),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -7393,6 +7432,8 @@ async def test_update_team_org_scoped_models_with_all_proxy_models(
mock_org.litellm_budget_table = None
with (
_team_admin_may_edit("models"),
_not_org_admin(),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -7502,6 +7543,7 @@ async def test_update_team_tpm_limit_not_gated_by_user_limit(
dummy_request = MagicMock(spec=Request)
with (
_team_admin_may_edit("tpm_limit"),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -7584,6 +7626,7 @@ async def test_update_team_rpm_limit_not_gated_by_user_limit(
dummy_request = MagicMock(spec=Request)
with (
_team_admin_may_edit("rpm_limit"),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -7981,6 +8024,8 @@ async def test_update_team_org_scoped_tpm_exceeds_org_limit():
mock_org.litellm_budget_table = mock_budget_table
with (
_team_admin_may_edit("tpm_limit"),
_not_org_admin(),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -8067,6 +8112,8 @@ async def test_update_team_org_scoped_rpm_exceeds_org_limit():
mock_org.litellm_budget_table = mock_budget_table
with (
_team_admin_may_edit("rpm_limit"),
_not_org_admin(),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -8158,6 +8205,8 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(
mock_org.litellm_budget_table = mock_budget_table
with (
_team_admin_may_edit("tpm_limit", "rpm_limit"),
_not_org_admin(),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -8286,6 +8335,7 @@ async def test_update_team_guardrails_with_org_id(
}
with (
_team_admin_may_edit("guardrails", "organization_id"),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
@ -11177,8 +11227,8 @@ async def test_update_team_blocks_non_admin_passthrough_routes(mock_db_client):
mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing)
with patch(
"litellm.proxy.management_endpoints.team_endpoints._verify_team_access",
AsyncMock(return_value=None),
"litellm.proxy.management_endpoints.team_endpoints._resolve_team_access",
AsyncMock(return_value="org_admin"),
):
with pytest.raises(ProxyException) as exc:
await update_team(
@ -13246,6 +13296,7 @@ async def test_update_team_output_token_estimate_lowered_rejected_for_team_admin
with contextlib.ExitStack() as stack:
_wire_update_team(stack, {_TEAM_ESTIMATE: 4000})
stack.enter_context(_team_admin_may_edit("default_estimated_output_tokens"))
with pytest.raises(ProxyException) as exc:
await update_team(
data=UpdateTeamRequest(team_id="test_team_id", default_estimated_output_tokens=1),
@ -13277,6 +13328,7 @@ async def test_update_team_output_token_estimate_unchanged_allows_team_admin_edi
with contextlib.ExitStack() as stack:
prisma = _wire_update_team(stack, {_TEAM_ESTIMATE: 4000})
stack.enter_context(_team_admin_may_edit("team_alias"))
await update_team(
data=UpdateTeamRequest(
team_id="test_team_id",
@ -13336,6 +13388,7 @@ async def test_update_team_batch_enqueued_token_limit_raised_rejected_for_team_a
with contextlib.ExitStack() as stack:
_wire_update_team(stack, {_TEAM_BATCH_LIMIT: 100000})
stack.enter_context(_team_admin_may_edit("metadata"))
with pytest.raises(ProxyException) as exc:
await update_team(
data=UpdateTeamRequest(team_id="test_team_id", metadata={_TEAM_BATCH_LIMIT: 10**12}),
@ -14894,3 +14947,369 @@ async def test_update_team_model_max_budget_raise_blocked_for_team_admin():
assert exc.value.code == "403"
assert "proxy admin" in str(exc.value.message).lower()
mock_prisma.db.litellm_teamtable.update.assert_not_awaited()
# ---------------------------------------------------------------------------
# LIT-5722: team admins reach update_team through self_managed_routes and are
# filtered by the team_admin_editable_team_fields setting.
# ---------------------------------------------------------------------------
_TEAM_ADMIN_CALLER = UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-team-admin", user_id="team-admin"
)
_PROXY_ADMIN_CALLER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin")
def _update_request_stub():
from unittest.mock import Mock
from fastapi import Request
return Mock(spec=Request)
@pytest.mark.asyncio
async def test_update_team_team_admin_is_refused_before_any_write_when_no_fields_are_enabled():
import contextlib
with contextlib.ExitStack() as stack:
prisma = _wire_update_team(stack, {})
stack.enter_context(_team_admin_may_edit())
with pytest.raises(ProxyException) as exc:
await update_team(
data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"),
http_request=_update_request_stub(),
user_api_key_dict=_TEAM_ADMIN_CALLER,
)
assert str(exc.value.code) == "403"
assert "cannot edit team settings" in str(exc.value.message)
assert not prisma.db.litellm_teamtable.update.called
@pytest.mark.asyncio
async def test_update_team_configured_but_unsupported_field_does_not_open_editing():
"""Only fields in SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS count, whatever general_settings says."""
import contextlib
with contextlib.ExitStack() as stack:
prisma = _wire_update_team(stack, {})
stack.enter_context(
patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": ["team_alias"]}) # test-quality-ok: update_team reads general_settings as a proxy_server module global
)
with pytest.raises(ProxyException) as exc:
await update_team(
data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"),
http_request=_update_request_stub(),
user_api_key_dict=_TEAM_ADMIN_CALLER,
)
assert str(exc.value.code) == "403"
assert "cannot edit team settings" in str(exc.value.message)
assert not prisma.db.litellm_teamtable.update.called
@pytest.mark.asyncio
async def test_update_team_team_admin_changing_an_unpermitted_field_is_refused_by_name():
import contextlib
with contextlib.ExitStack() as stack:
prisma = _wire_update_team(stack, {})
stack.enter_context(_team_admin_may_edit("team_alias"))
with pytest.raises(ProxyException) as exc:
await update_team(
data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed", tpm_limit=10),
http_request=_update_request_stub(),
user_api_key_dict=_TEAM_ADMIN_CALLER,
)
assert str(exc.value.code) == "403"
assert "'tpm_limit'" in str(exc.value.message)
assert not prisma.db.litellm_teamtable.update.called
@pytest.mark.asyncio
async def test_update_team_team_admin_echoing_unpermitted_fields_unchanged_is_allowed(
disable_audit_logging_for_mocked_team,
):
"""The dashboard resends the whole form, so only a value that differs from what is stored counts."""
import contextlib
with contextlib.ExitStack() as stack:
prisma = _wire_update_team(stack, {})
stack.enter_context(_team_admin_may_edit("team_alias"))
result = await update_team(
data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed", tpm_limit=None, models=[]),
http_request=_update_request_stub(),
user_api_key_dict=_TEAM_ADMIN_CALLER,
)
assert result["data"].team_id == "test_team_id"
assert prisma.db.litellm_teamtable.update.called
@pytest.mark.asyncio
async def test_update_team_team_admin_changes_tpm_limit_once_a_proxy_admin_enables_it(
disable_audit_logging_for_mocked_team,
):
"""tpm_limit is the first field a proxy admin can open to team admins; every other field stays admin-only."""
import contextlib
with contextlib.ExitStack() as stack:
prisma = _wire_update_team(stack, {})
stack.enter_context(
patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": ["tpm_limit"]}) # test-quality-ok: update_team reads general_settings as a proxy_server module global
)
await update_team(
data=UpdateTeamRequest(team_id="test_team_id", tpm_limit=5000),
http_request=_update_request_stub(),
user_api_key_dict=_TEAM_ADMIN_CALLER,
)
with pytest.raises(ProxyException) as refused:
await update_team(
data=UpdateTeamRequest(team_id="test_team_id", tpm_limit=6000, rpm_limit=10),
http_request=_update_request_stub(),
user_api_key_dict=_TEAM_ADMIN_CALLER,
)
assert prisma.db.litellm_teamtable.update.await_count == 1
assert prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["tpm_limit"] == 5000
assert str(refused.value.code) == "403"
assert "'rpm_limit'" in str(refused.value.message)
@pytest.mark.asyncio
async def test_update_team_team_admin_resending_budget_settings_does_not_push_back_budget_resets(
disable_audit_logging_for_mocked_team,
):
"""A resent budget_duration or budget_limits would otherwise recompute the reset timestamps from now."""
import contextlib
stored_windows = [{"budget_duration": "7d", "max_budget": 5.0, "reset_at": "2026-09-20T00:00:00Z"}]
budgeted_team = MagicMock()
budgeted_team.metadata = {}
budgeted_team.model_dump.return_value = {
"team_id": "test_team_id",
"team_alias": "test_team",
"metadata": {},
"budget_duration": "30d",
"budget_limits": stored_windows,
"members_with_roles": [{"user_id": "team-admin", "role": "admin"}],
}
with contextlib.ExitStack() as stack:
prisma = _wire_update_team(stack, {})
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=budgeted_team)
stack.enter_context(_team_admin_may_edit("tpm_limit"))
await update_team(
data=UpdateTeamRequest(
team_id="test_team_id", tpm_limit=5000, budget_duration="30d", budget_limits=stored_windows
),
http_request=_update_request_stub(),
user_api_key_dict=_TEAM_ADMIN_CALLER,
)
written = prisma.db.litellm_teamtable.update.call_args.kwargs["data"]
assert written["tpm_limit"] == 5000
assert not {"budget_duration", "budget_reset_at", "budget_limits"} & written.keys()
@pytest.mark.asyncio
async def test_update_team_holds_a_team_admin_to_the_org_tpm_limit(disable_audit_logging_for_mocked_team):
"""The org ceiling lives on the org's budget row, so /team/update must load it to enforce the cap."""
import contextlib
capped_org = LiteLLM_OrganizationTable(
organization_id="capped-org",
budget_id="capped-budget",
created_by="admin",
updated_by="admin",
litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=10000),
)
async def org_lookup(**kwargs):
return capped_org if kwargs.get("include_budget_table") else capped_org.model_copy(
update={"litellm_budget_table": None}
)
org_team = MagicMock()
org_team.metadata = {}
org_team.organization_id = "capped-org"
org_team.model_dump.return_value = {
"team_id": "test_team_id",
"team_alias": "test_team",
"organization_id": "capped-org",
"metadata": {},
"members_with_roles": [{"user_id": "team-admin", "role": "admin"}],
}
with contextlib.ExitStack() as stack:
prisma = _wire_update_team(stack, {})
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=org_team)
stack.enter_context(_team_admin_may_edit("tpm_limit"))
stack.enter_context(
patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide
"litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team",
AsyncMock(return_value=False),
)
)
stack.enter_context(
patch( # test-quality-ok: update_team reads orgs through this module-level import; no seam to inject
"litellm.proxy.management_endpoints.team_endpoints.get_org_object",
AsyncMock(side_effect=org_lookup),
)
)
with pytest.raises(ProxyException) as over_cap:
await update_team(
data=UpdateTeamRequest(team_id="test_team_id", tpm_limit=20000),
http_request=_update_request_stub(),
user_api_key_dict=_TEAM_ADMIN_CALLER,
)
await update_team(
data=UpdateTeamRequest(team_id="test_team_id", tpm_limit=8000),
http_request=_update_request_stub(),
user_api_key_dict=_TEAM_ADMIN_CALLER,
)
assert str(over_cap.value.code) == "400"
assert "exceeds organization's tpm_limit (10000)" in str(over_cap.value.message)
assert prisma.db.litellm_teamtable.update.await_count == 1
assert prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["tpm_limit"] == 8000
@pytest.mark.asyncio
async def test_update_team_org_admin_is_not_filtered_by_the_team_admin_field_list(
disable_audit_logging_for_mocked_team,
):
"""A caller who is both org admin and roster admin keeps unrestricted edits."""
import contextlib
with contextlib.ExitStack() as stack:
prisma = _wire_update_team(stack, {})
stack.enter_context(_team_admin_may_edit())
stack.enter_context(
patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide
"litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team",
AsyncMock(return_value=True),
)
)
result = await update_team(
data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"),
http_request=_update_request_stub(),
user_api_key_dict=_TEAM_ADMIN_CALLER,
)
assert result["data"].team_id == "test_team_id"
assert prisma.db.litellm_teamtable.update.called
@pytest.mark.asyncio
async def test_update_team_unknown_team_is_403_for_non_proxy_admins_and_404_for_proxy_admins():
"""Now that any authenticated caller reaches the handler, 'team not found' must not leak team ids."""
import contextlib
with contextlib.ExitStack() as stack:
prisma = _wire_update_team(stack, {})
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
with pytest.raises(ProxyException) as denied:
await update_team(
data=UpdateTeamRequest(team_id="no-such-team", team_alias="renamed"),
http_request=_update_request_stub(),
user_api_key_dict=_TEAM_ADMIN_CALLER,
)
with pytest.raises(ProxyException) as missing:
await update_team(
data=UpdateTeamRequest(team_id="no-such-team", team_alias="renamed"),
http_request=_update_request_stub(),
user_api_key_dict=_PROXY_ADMIN_CALLER,
)
assert str(denied.value.code) == "403"
assert "do not have access to this team" in str(denied.value.message)
assert "no-such-team" not in str(denied.value.message)
assert str(missing.value.code) == "404"
@pytest.mark.asyncio
async def test_resolve_team_access_ranks_proxy_admin_then_org_admin_then_team_admin():
from litellm.proxy.management_endpoints.team_endpoints import _resolve_team_access
team = LiteLLM_TeamTable(
team_id="team-1",
organization_id="org-1",
members_with_roles=[Member(user_id="team-admin", role="admin")],
)
roster_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin")
outsider = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="someone-else")
org_lookup = AsyncMock(return_value=False)
with patch("litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", org_lookup): # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide
assert await _resolve_team_access(team_obj=team, user_api_key_dict=_PROXY_ADMIN_CALLER) == "proxy_admin"
assert org_lookup.await_count == 0
assert await _resolve_team_access(team_obj=team, user_api_key_dict=roster_admin) == "team_admin"
assert await _resolve_team_access(team_obj=team, user_api_key_dict=outsider) is None
org_lookup.return_value = True
assert await _resolve_team_access(team_obj=team, user_api_key_dict=roster_admin) == "org_admin"
_ROSTER_ADMIN_CALLER = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="admin-1")
_MEMBER_CALLER = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="member-1")
@pytest.mark.parametrize(
"caller, org_admin, enabled_fields, expected",
[
pytest.param(_PROXY_ADMIN_CALLER, False, (), {"kind": "unrestricted"}, id="proxy-admin"),
pytest.param(
UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, user_id="viewer"),
False,
("tpm_limit",),
{"kind": "none"},
id="proxy-admin-viewer",
),
pytest.param(_ROSTER_ADMIN_CALLER, True, (), {"kind": "unrestricted"}, id="org-admin-who-is-also-team-admin"),
pytest.param(_ROSTER_ADMIN_CALLER, False, (), {"kind": "team_admin_disabled"}, id="team-admin-nothing-enabled"),
pytest.param(
_ROSTER_ADMIN_CALLER,
False,
("tpm_limit",),
{"kind": "team_admin", "editable_fields": ["tpm_limit"]},
id="team-admin-field-enabled",
),
pytest.param(_MEMBER_CALLER, False, ("tpm_limit",), {"kind": "none"}, id="plain-member"),
],
)
@pytest.mark.asyncio
async def test_team_info_reports_what_the_caller_may_edit(caller, org_admin, enabled_fields, expected):
"""The dashboard gates its edit form on this field instead of guessing the caller's role from the org list,
which is premium-gated and can be empty for a dual-role org admin."""
from fastapi import Request
from litellm.proxy.management_endpoints import team_endpoints
team_row = LiteLLM_TeamTable(
team_id="team-1",
organization_id="org-1",
members_with_roles=[Member(user_id="admin-1", role="admin"), Member(user_id="member-1", role="user")],
)
mock_prisma = MagicMock()
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[])
mock_prisma.get_data = AsyncMock(return_value=[])
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: no seam on team_info
patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[])), # test-quality-ok: no seam on team_info
patch.object( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide
team_endpoints, "_is_user_org_admin_for_team", AsyncMock(return_value=org_admin)
),
_team_admin_may_edit(*enabled_fields),
):
response = await team_endpoints.team_info(
http_request=MagicMock(spec=Request),
team_id="team-1",
user_api_key_dict=caller,
)
assert response["team_info"].caller_edit_access.model_dump(mode="json") == expected

View file

@ -1983,6 +1983,144 @@ class TestBedrockAgentRuntimePassthroughToggle:
create_route.assert_called_once()
class TestBedrockAgentRuntimePassthroughVirtualKeyLeak:
VKEY: Final = "sk-litellm-victim-key"
MASTER_KEY: Final = "sk-master-1234"
ENDPOINT: Final = "knowledgebases/KB1234567/retrieve"
AMBIENT_AWS_ENV: Final = (
"AWS_BEARER_TOKEN_BEDROCK",
"AWS_SESSION_TOKEN",
"AWS_SESSION_NAME",
"AWS_PROFILE_NAME",
"AWS_ROLE_NAME",
"AWS_WEB_IDENTITY_TOKEN",
"AWS_STS_ENDPOINT",
"AWS_EXTERNAL_ID",
)
async def _upstream_headers(self, monkeypatch, headers: list[tuple[bytes, bytes]]) -> dict:
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import HttpPassThroughEndpointHelpers
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", self.MASTER_KEY)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
for ambient in self.AMBIENT_AWS_ENV:
monkeypatch.delenv(ambient, raising=False)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "ak")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "sk")
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
caller: Final = UserAPIKeyAuth(api_key=self.VKEY)
async def receive():
return {"type": "http.request", "body": b'{"retrievalQuery": {"text": "hi"}}', "more_body": False}
request: Final = Request(
{
"type": "http",
"method": "POST",
"path": f"/bedrock/{self.ENDPOINT}",
"headers": headers,
"query_string": b"",
},
receive=receive,
)
captured: dict = {}
def fake_create_pass_through_route(**kwargs):
captured.update(kwargs)
return AsyncMock(return_value={"status": "success"})
module: Final = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
with (
patch(f"{module}.create_request_copy", Mock()),
patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route),
):
await bedrock_proxy_route(
endpoint=self.ENDPOINT,
request=request,
fastapi_response=Response(),
user_api_key_dict=caller,
)
return HttpPassThroughEndpointHelpers.forward_headers_from_request(
request_headers=dict(request.headers),
headers=dict(captured["custom_headers"] or {}),
forward_headers=captured.get("_forward_headers", False),
)
@staticmethod
def _blob(upstream: dict) -> str:
return " ".join(f"{name}:{value}" for name, value in upstream.items())
@staticmethod
def _names_matching(upstream: dict, lowercase_name: str) -> list[str]:
return [name for name in upstream if name.lower() == lowercase_name]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"header_name", ["x-api-key", "x-litellm-api-key", "api-key", "x-goog-api-key", "ocp-apim-subscription-key"]
)
async def test_virtual_key_in_a_credential_header_never_reaches_aws(self, monkeypatch, header_name: str):
upstream: Final = await self._upstream_headers(
monkeypatch,
[
(header_name.encode(), self.VKEY.encode()),
(b"content-type", b"application/json"),
(b"x-request-id", b"trace-1"),
],
)
assert self.VKEY not in self._blob(upstream)
assert self._names_matching(upstream, header_name) == []
assert upstream["x-request-id"] == "trace-1", "a benign caller header still reaches AWS"
assert upstream["Authorization"].startswith("AWS4-HMAC-SHA256")
assert self._names_matching(upstream, "content-type") == ["Content-Type"], "the signed header is the only one"
@pytest.mark.asyncio
async def test_credential_headers_are_dropped_by_name_even_when_they_carry_someone_elses_key(self, monkeypatch):
other_key: Final = "sk-other-tenant-key"
upstream: Final = await self._upstream_headers(
monkeypatch,
[
(b"x-api-key", other_key.encode()),
(b"x-litellm-api-key", other_key.encode()),
(b"x-request-id", b"trace-3"),
],
)
assert other_key not in self._blob(upstream)
assert self._names_matching(upstream, "x-api-key") == []
assert self._names_matching(upstream, "x-litellm-api-key") == []
assert upstream["x-request-id"] == "trace-3"
@pytest.mark.asyncio
async def test_virtual_key_in_authorization_bearer_is_replaced_by_the_sigv4_signature(self, monkeypatch):
upstream: Final = await self._upstream_headers(
monkeypatch,
[(b"authorization", f"Bearer {self.VKEY}".encode()), (b"content-type", b"application/json")],
)
assert self.VKEY not in self._blob(upstream)
assert self._names_matching(upstream, "authorization") == ["Authorization"]
assert upstream["Authorization"].startswith("AWS4-HMAC-SHA256")
@pytest.mark.asyncio
async def test_authenticated_secrets_in_any_other_header_never_reach_aws(self, monkeypatch):
upstream: Final = await self._upstream_headers(
monkeypatch,
[
(b"x-api-key", self.VKEY.encode()),
(b"x-forwarded-key", self.VKEY.encode()),
(b"x-operator-token", self.MASTER_KEY.encode()),
(b"x-request-id", b"trace-2"),
],
)
assert self.VKEY not in self._blob(upstream) and self.MASTER_KEY not in self._blob(upstream)
assert self._names_matching(upstream, "x-forwarded-key") == []
assert self._names_matching(upstream, "x-operator-token") == []
assert upstream["x-request-id"] == "trace-2"
class TestLLMPassthroughFactoryProxyRoute:
@pytest.mark.asyncio
async def test_llm_passthrough_factory_proxy_route_success(self):

View file

@ -6143,3 +6143,42 @@ async def test_chat_completion_pass_through_endpoint_answers_an_openai_typed_err
)
assert (raised.value.type, raised.value.param, raised.value.code) == ("invalid_request_error", None, "400")
@pytest.mark.asyncio
async def test_chat_completion_pass_through_endpoint_failure_carries_the_callers_litellm_call_id(
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
):
call_id = "lit7836-pass-through-call-id"
proxy_logging = MagicMock()
proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"])
proxy_logging.post_call_failure_hook = AsyncMock()
async def fake_add_litellm_data_to_request(**kwargs: object) -> object:
return kwargs["data"]
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging)
monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
request = MagicMock(spec=Request)
request.headers = Headers({"x-litellm-call-id": call_id})
request.body = AsyncMock(
return_value=json.dumps({"model": "unknown-model", "messages": [{"role": "user", "content": "hi"}]}).encode()
)
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised:
await chat_completion_pass_through_endpoint(
fastapi_response=Response(),
request=request,
adapter_id="anthropic",
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
)
assert raised.value.headers["x-litellm-call-id"] == call_id
record = next(r for r in caplog.records if "Exception occured" in r.getMessage())
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()

View file

@ -3938,3 +3938,58 @@ async def test_ProxyConfig__init_guardrails_in_db_skips_only_the_unloadable_row(
assert sorted(handler.IN_MEMORY_GUARDRAILS) == ["first", "last"]
assert handler.reconciled_with == [{"first", "broken", "last"}]
# ---------------------------------------------------------------------------
# add_deployment: UI settings convergence
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_add_deployment_re_reads_ui_settings_so_other_pods_converge(monkeypatch):
"""The periodic config reload picks up a UI setting written through another pod.
Startup used to be the only read, so a proxy admin flipping a runtime flag reached the pod
that served the PATCH and nowhere else until every other pod restarted.
"""
general_settings: Dict[str, Any] = {"allow_agents_for_team_admins": False}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
prisma_client = MagicMock()
prisma_client.db.litellm_config.find_many = AsyncMock(return_value=[])
prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None)
prisma_client.db.litellm_credentialstable.find_many = AsyncMock(return_value=[])
prisma_client.db.litellm_uisettings.find_unique = AsyncMock(
return_value=SimpleNamespace(
ui_settings=json.dumps({"allow_agents_for_team_admins": True, "enable_chat_ui": False})
)
)
config = ProxyConfig()
config._should_load_db_object = MagicMock(return_value=False)
config._init_non_llm_objects_in_db = AsyncMock()
await config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=MagicMock())
prisma_client.db.litellm_uisettings.find_unique.assert_awaited_once_with(where={"id": "ui_settings"})
assert general_settings["allow_agents_for_team_admins"] is True
assert "enable_chat_ui" not in general_settings
@pytest.mark.asyncio
async def test_add_deployment_syncs_ui_settings_even_when_the_model_reconcile_fails(monkeypatch):
"""A broken model reconcile must not strand every pod on stale settings."""
general_settings: Dict[str, Any] = {"allow_agents_for_team_admins": False}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
prisma_client = MagicMock()
prisma_client.db.litellm_uisettings.find_unique = AsyncMock(
return_value=SimpleNamespace(ui_settings={"allow_agents_for_team_admins": True})
)
config = ProxyConfig()
config._should_load_db_object = MagicMock(side_effect=RuntimeError("db down"))
await config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=MagicMock())
assert general_settings["allow_agents_for_team_admins"] is True

View file

@ -9,7 +9,7 @@ Pins (PR2):
from __future__ import annotations
from unittest.mock import MagicMock
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -128,7 +128,6 @@ def test_v1_model_info_no_model_list_error(client, auth_as, null_router, path):
assert "LLM Model List not loaded" in response.text
def test_get_proxy_model_info_surfaces_supports_parallel_function_calling(local_model_cost_map):
"""``GET /v1/model/info`` enriches each deployment through ``_get_proxy_model_info``; a registry
entry declaring parallel function calling must land in ``model_info`` instead of null."""
@ -161,9 +160,7 @@ def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monkeypatch
router.get_model_list = MagicMock(return_value=[deployment])
monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models)
expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info(
[deployment]
)
expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info([deployment])
allowed_model_names = proxy_server._get_v1_model_info_allowed_model_names(
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-test",
@ -308,6 +305,80 @@ def test_model_group_info_invalid_method(client, auth_as, null_router):
assert len(response.content) > 0
@pytest.fixture
def model_group_info_router(monkeypatch):
from litellm.types.proxy.management_endpoints.model_management_endpoints import ModelGroupInfoProxy
model_names = ["gpt-4", "claude-3"]
router = MagicMock()
router.get_model_names.return_value = model_names
router.get_model_access_groups.return_value = {}
router.get_model_list.return_value = []
def model_group_info(*, llm_router, all_models_str, model_group):
return [ModelGroupInfoProxy(model_group=name, providers=[]) for name in all_models_str]
async def append_agents_to_model_group(*, model_groups, user_api_key_dict):
return model_groups
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "llm_model_list", [{"model_name": name} for name in model_names])
monkeypatch.setattr(proxy_server, "user_model", None)
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", None)
monkeypatch.setattr(proxy_server, "user_api_key_cache", None)
monkeypatch.setattr(proxy_server, "_get_model_group_info", model_group_info)
from litellm.proxy.agent_endpoints import model_list_helpers
monkeypatch.setattr(
model_list_helpers,
"append_agents_to_model_group",
AsyncMock(side_effect=append_agents_to_model_group),
)
return router
@pytest.mark.parametrize("admin_role", ["proxy_admin", "proxy_admin_viewer"])
def test_model_group_info_proxy_admin_ignores_key_model_restriction(
client, auth_as, model_group_info_router, admin_role
):
from litellm.proxy._types import LitellmUserRoles
with auth_as(LitellmUserRoles(admin_role), models=["no-default-models"]):
response = client.get("/model_group/info")
assert response.status_code == 200
assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4", "claude-3"]
@pytest.mark.parametrize("admin_role", ["proxy_admin", "proxy_admin_viewer"])
def test_model_group_info_proxy_admin_expands_wildcard_deployments(client, auth_as, model_group_info_router, admin_role):
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
model_group_info_router.get_model_names.return_value = ["gpt-4", "anthropic/*"]
known_anthropic_models = get_known_models_from_wildcard(wildcard_model="anthropic/*")
assert known_anthropic_models
with auth_as(LitellmUserRoles(admin_role), models=["no-default-models"]):
response = client.get("/model_group/info")
assert response.status_code == 200
assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4", *known_anthropic_models]
def test_model_group_info_internal_user_key_model_restriction_applies(client, auth_as, model_group_info_router):
from litellm.proxy._types import LitellmUserRoles
with auth_as(LitellmUserRoles.INTERNAL_USER, models=["gpt-4"]):
response = client.get("/model_group/info")
assert response.status_code == 200
assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4"]
# ---------------------------------------------------------------------------
# GET /v2/model/info?exclude_auto_routers
# ---------------------------------------------------------------------------
@ -399,14 +470,10 @@ def test_v2_model_info_exclude_auto_routers_shrinks_total_count(client, auth_as,
assert len(payload["data"]) == payload["total_count"]
def test_v2_model_info_exclude_auto_routers_paginates_over_the_filtered_set(
client, auth_as, mixed_auto_router_router
):
def test_v2_model_info_exclude_auto_routers_paginates_over_the_filtered_set(client, auth_as, mixed_auto_router_router):
"""Page size applies to the filtered list, so no page silently comes back short."""
with auth_as():
response = client.get(
"/v2/model/info", params={"exclude_auto_routers": "true", "page": 1, "size": 1}
)
response = client.get("/v2/model/info", params={"exclude_auto_routers": "true", "page": 1, "size": 1})
payload = response.json()
assert payload["total_count"] == 2
assert payload["total_pages"] == 2

View file

@ -3,6 +3,8 @@ Tests for rerank_endpoints/endpoints.py response headers.
"""
import json
import logging
from collections.abc import Iterator
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -10,6 +12,7 @@ from fastapi import HTTPException, Request, Response
import litellm.proxy.common_request_processing as common_request_processing_mod
import litellm.proxy.proxy_server as proxy_server_mod
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.rerank_endpoints.endpoints import rerank
from litellm.types.utils import RerankResponse
@ -28,7 +31,7 @@ HIDDEN_PARAMS = {
}
def _build_request() -> Request:
def _build_request(headers: tuple[tuple[bytes, bytes], ...] = ()) -> Request:
body = json.dumps({"model": "rerank-model", "query": "q", "documents": ["a", "b"]}).encode()
async def receive():
@ -39,7 +42,7 @@ def _build_request() -> Request:
"type": "http",
"method": "POST",
"path": "/rerank",
"headers": [(b"content-type", b"application/json")],
"headers": [(b"content-type", b"application/json"), *headers],
"query_string": b"",
},
receive=receive,
@ -56,7 +59,7 @@ async def _call_rerank(hidden_params: dict = HIDDEN_PARAMS) -> Response:
proxy_logging_obj.update_request_status = AsyncMock()
async def fake_add_litellm_data_to_request(**kwargs):
return {**kwargs["data"], "litellm_call_id": "call-123"}
return dict(kwargs["data"])
async def fake_route_request(**kwargs):
async def _call():
@ -72,7 +75,7 @@ async def _call_rerank(hidden_params: dict = HIDDEN_PARAMS) -> Response:
patch.object(proxy_server_mod, "version", "1.2.3"), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
):
await rerank(
request=_build_request(),
request=_build_request(headers=((b"x-litellm-call-id", b"call-123"),)),
fastapi_response=fastapi_response,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
)
@ -121,7 +124,11 @@ async def test_rerank_omits_detailed_timing_headers_when_disabled():
async def _rerank_failure(
failure: Exception, *, raised_before_routing: bool, monkeypatch: pytest.MonkeyPatch
failure: Exception,
*,
raised_before_routing: bool,
monkeypatch: pytest.MonkeyPatch,
headers: tuple[tuple[bytes, bytes], ...] = (),
) -> ProxyException:
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(
@ -143,13 +150,45 @@ async def _rerank_failure(
with pytest.raises(ProxyException) as raised:
await rerank(
request=_build_request(),
request=_build_request(headers),
fastapi_response=Response(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
)
return raised.value
@pytest.fixture
def propagating_proxy_logger() -> Iterator[None]:
verbose_proxy_logger.propagate = True
try:
yield
finally:
verbose_proxy_logger.propagate = False
@pytest.mark.asyncio
async def test_failure_log_carries_the_callers_litellm_call_id(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, propagating_proxy_logger: None
) -> None:
"""LIT-7836: the /rerank error line must carry the same litellm_call_id the client
sent, both in the rendered message and as a structured log record field."""
call_id = "rerank-call-7836"
failure = HTTPException(status_code=401, detail={"error": "invalid api key"})
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
raised = await _rerank_failure(
failure,
raised_before_routing=False,
monkeypatch=monkeypatch,
headers=((b"x-litellm-call-id", call_id.encode()),),
)
assert raised.headers["x-litellm-call-id"] == call_id
record = next(r for r in caplog.records if "Exception occured" in r.getMessage())
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
@pytest.mark.asyncio
async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(monkeypatch: pytest.MonkeyPatch):
"""A bare HTTPException carries no type or param, so the tail used to ship the

View file

@ -44,6 +44,7 @@ from litellm.proxy.common_request_processing import (
create_response,
)
from litellm.proxy.dd_span_tagger import DDSpanTagger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import ProxyException
from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
@ -6365,6 +6366,206 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
call_type="acompletion",
)
@staticmethod
def _v3_limiter_rig(
monkeypatch: pytest.MonkeyPatch,
user_api_key_dict: ProxyUserAPIKeyAuth,
fallbacks: list[dict[str, list[str]]],
) -> tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]]:
"""Real v3 limiter (the default ``parallel_request_limiter``) wired in through the
``proxy_logging_obj`` seam, so ``common_processing_pre_call_logic`` runs for real:
``add_litellm_data_to_request`` with a live OTel span, ``function_setup``, then the limiter."""
from litellm.caching.caching import DualCache
from litellm.proxy import proxy_server
from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3
from litellm.proxy.utils import InternalUsageCache
monkeypatch.setattr(proxy_server, "prisma_client", None)
limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache()))
limiter_models: list[str] = []
async def run_limiter(
user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str
) -> dict[str, object]:
limiter_models.append(str(data["model"]))
await limiter.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data=data,
call_type=call_type,
)
return data
proxy_logging_obj = MagicMock(spec=ProxyLogging)
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=run_limiter)
router = litellm.Router(
model_list=[
{"model_name": group, "litellm_params": {"model": "openai/gpt-4.1-nano", "api_key": "fake"}}
for chain in fallbacks
for group in (*chain.keys(), *(m for models in chain.values() for m in models))
],
fallbacks=fallbacks,
)
return proxy_logging_obj, router, proxy_server.ProxyConfig(), limiter_models
@staticmethod
def _otel_key(
rpm_limit: int | None = None,
model_rpm_limit: dict[str, int] | None = None,
disable_fallbacks: bool = False,
) -> ProxyUserAPIKeyAuth:
from opentelemetry.sdk.trace import TracerProvider
span = TracerProvider().get_tracer("test").start_span("proxy-request")
return ProxyUserAPIKeyAuth(
api_key="hashed-key",
parent_otel_span=span,
rpm_limit=rpm_limit,
metadata={
**({"model_rpm_limit": model_rpm_limit} if model_rpm_limit else {}),
**({"disable_fallbacks": True} if disable_fallbacks else {}),
},
)
@staticmethod
def _chat_request() -> Request:
return Request({"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []})
async def _pre_call(
self,
data: dict[str, object],
user_api_key_dict: ProxyUserAPIKeyAuth,
rig: tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]],
) -> tuple[ProxyBaseLLMRequestProcessing, tuple[dict[str, object], LiteLLMLoggingObj]]:
proxy_logging_obj, router, proxy_config, _ = rig
processor = ProxyBaseLLMRequestProcessing(data=data)
result = await processor._pre_call_with_fallbacks(
request=self._chat_request(),
general_settings={},
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
version=None,
proxy_config=proxy_config,
user_model=None,
user_temperature=None,
user_request_timeout=None,
user_max_tokens=None,
user_api_base=None,
model=None,
route_type="acompletion",
llm_router=router,
)
return processor, result
@pytest.mark.asyncio
async def test_v3_limiter_with_otel_span_falls_back_from_client_request(self, monkeypatch: pytest.MonkeyPatch):
"""Customer path: OTel on, per-key model RPM cap on the primary, a router fallback configured.
The first pass enriches ``data["metadata"]`` with the live span, then the limiter raises. The
fallback pass must start from the client's request again, so ``add_litellm_data_to_request``
never deep-copies the span (the ``cannot pickle '_thread.RLock'`` 500)."""
primary_model = "gpt-4.1"
fallback_model = "gpt-4.1-mini"
key = self._otel_key(model_rpm_limit={primary_model: 1})
rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}])
def client_request() -> dict[str, object]:
return {
"model": primary_model,
"messages": [{"role": "user", "content": "hi"}],
"metadata": {"tags": ["client-tag"]},
}
_, (first_data, _) = await self._pre_call(client_request(), key, rig)
processor, (data, logging_obj) = await self._pre_call(client_request(), key, rig)
assert first_data["model"] == primary_model
assert data["model"] == fallback_model
assert processor.data is data
assert data["litellm_logging_obj"] is logging_obj
assert logging_obj.model == fallback_model
requester_metadata = data["metadata"]["requester_metadata"]
assert requester_metadata["tags"] == ["client-tag"]
assert "litellm_parent_otel_span" not in requester_metadata
assert "user_api_key_auth" not in requester_metadata
assert data["metadata"]["litellm_parent_otel_span"] is key.parent_otel_span
assert rig[3] == [primary_model, primary_model, fallback_model]
@pytest.mark.asyncio
async def test_v3_limiter_with_otel_span_returns_429_when_fallbacks_exhausted(
self, monkeypatch: pytest.MonkeyPatch
):
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
primary_model = "gpt-4.1"
fallback_model = "gpt-4.1-mini"
key = self._otel_key(rpm_limit=1)
rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}])
request = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]}
await self._pre_call(dict(request), key, rig)
processor = ProxyBaseLLMRequestProcessing(data=dict(request))
with pytest.raises(ProxyRateLimitError) as exc_info:
await processor._pre_call_with_fallbacks(
request=self._chat_request(),
general_settings={},
proxy_logging_obj=rig[0],
user_api_key_dict=key,
version=None,
proxy_config=rig[2],
user_model=None,
user_temperature=None,
user_request_timeout=None,
user_max_tokens=None,
user_api_base=None,
model=None,
route_type="acompletion",
llm_router=rig[1],
)
assert rig[3] == [primary_model, primary_model, fallback_model]
assert exc_info.value.status_code == 429
assert "Rate limit exceeded" in str(exc_info.value.detail)
assert exc_info.value.headers["retry-after"]
assert processor.data["model"] == primary_model
assert processor.data["litellm_logging_obj"].model == primary_model
assert processor.data["litellm_call_id"]
@pytest.mark.asyncio
async def test_fallback_lookup_uses_alias_resolved_model_group(self, monkeypatch: pytest.MonkeyPatch):
primary_model = "gpt-4.1"
fallback_model = "gpt-4.1-mini"
monkeypatch.setattr(litellm, "model_alias_map", {"my-alias": primary_model})
key = self._otel_key(model_rpm_limit={primary_model: 1})
rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}])
request = {"model": "my-alias", "messages": [{"role": "user", "content": "hi"}]}
await self._pre_call(dict(request), key, rig)
_, (data, _) = await self._pre_call(dict(request), key, rig)
assert data["model"] == fallback_model
assert rig[3] == [primary_model, primary_model, fallback_model]
@pytest.mark.asyncio
async def test_key_metadata_disable_fallbacks_returns_429_instead_of_retrying(
self, monkeypatch: pytest.MonkeyPatch
):
"""``disable_fallbacks`` set in key metadata only lands on ``data`` during the first
pre-call pass (``add_key_level_controls``), so it must be honored after that pass."""
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
primary_model = "gpt-4.1"
fallback_model = "gpt-4.1-mini"
key = self._otel_key(model_rpm_limit={primary_model: 1}, disable_fallbacks=True)
rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}])
request = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]}
await self._pre_call(dict(request), key, rig)
with pytest.raises(ProxyRateLimitError) as exc_info:
await self._pre_call(dict(request), key, rig)
assert exc_info.value.status_code == 429
assert rig[3] == [primary_model, primary_model]
class _RecordingSuccessLogger(CustomLogger):
def __init__(self):
@ -8212,7 +8413,7 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_
"""Regression for LIT-6043: expected 4xx errors log without formatting a
traceback; unexpected errors keep logger.exception behavior."""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_request_processing import _log_llm_api_exception
from litellm.proxy.common_request_processing import log_llm_api_exception
verbose_proxy_logger.propagate = True
try:
@ -8220,7 +8421,7 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_
try:
raise exc
except Exception as raised:
_log_llm_api_exception(raised, "call-id-for-traceback-test")
log_llm_api_exception(raised, "call-id-for-traceback-test")
finally:
verbose_proxy_logger.propagate = False
@ -8778,14 +8979,14 @@ class TestErrorLogCarriesCallId:
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_request_processing import (
_CLIENT_DISCONNECT_DETAIL,
_log_llm_api_exception,
log_llm_api_exception,
)
call_id: Final = str(uuid.uuid4())
verbose_proxy_logger.propagate = True
try:
with caplog.at_level("INFO", logger="LiteLLM Proxy"):
_log_llm_api_exception(
log_llm_api_exception(
HTTPException(status_code=499, detail=_CLIENT_DISCONNECT_DETAIL),
call_id,
)

View file

@ -2,6 +2,7 @@ import asyncio
import contextlib
import importlib
import json
import logging
import os
import re
import socket
@ -19,7 +20,7 @@ import fastapi.routing
import httpx
import pytest
import yaml
from fastapi import FastAPI, Request
from fastapi import FastAPI, HTTPException, Request
from fastapi.encoders import jsonable_encoder
from fastapi.staticfiles import StaticFiles
from fastapi.testclient import TestClient
@ -13003,6 +13004,144 @@ async def test_moderations_response_carries_litellm_call_id_header():
assert fastapi_response.headers["x-litellm-model-id"] == "mod-deployment-1"
@pytest.mark.asyncio
async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplog):
"""LIT-7836: the /v1/moderations error line must carry the litellm_call_id the
client sent, rendered in the message and as a structured log record field."""
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ProxyException
call_id = "moderations-call-7836"
async def passthrough_add_litellm_data(data, **kwargs):
return data
request = MagicMock()
request.headers = {"x-litellm-call-id": call_id}
request.body = AsyncMock(return_value=b'{"input": "hi"}')
fake_logging = MagicMock()
fake_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data)
fake_logging.post_call_failure_hook = AsyncMock()
verbose_proxy_logger.propagate = True
try:
with (
patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point
patch.object(proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key"))), # test-quality-ok: fakes the provider failure so the real route's error log is observable
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
pytest.raises(ProxyException) as raised,
):
await proxy_server_module.moderations(
request=request,
fastapi_response=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0),
)
finally:
verbose_proxy_logger.propagate = False
assert raised.value.headers["x-litellm-call-id"] == call_id
record = next(r for r in caplog.records if "Exception occured" in r.getMessage())
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
@pytest.mark.asyncio
async def test_moderations_unparseable_body_bills_the_callers_litellm_call_id():
"""LIT-7836: a body that fails to parse must still hand the failure hook the
litellm_call_id the response header answers with, so the spend row is findable."""
from litellm.proxy._types import ProxyException
call_id = "moderations-early-7836"
request = MagicMock()
request.headers = {"x-litellm-call-id": call_id}
request.body = AsyncMock(return_value=b'{"input": ')
fake_logging = MagicMock()
fake_logging.post_call_failure_hook = AsyncMock()
with (
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
pytest.raises(ProxyException) as raised,
):
await proxy_server_module.moderations(
request=request,
fastapi_response=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0),
)
assert raised.value.headers["x-litellm-call-id"] == call_id
hook_request_data = fake_logging.post_call_failure_hook.await_args.kwargs["request_data"]
assert hook_request_data["litellm_call_id"] == call_id
@pytest.mark.asyncio
async def test_moderations_already_shaped_failure_answers_with_the_callers_litellm_call_id():
"""LIT-7836: a ProxyException raised inside /v1/moderations is re-raised unwrapped but still
answers with the caller's x-litellm-call-id so the client can join it to the error log."""
call_id = "moderations-call-7836-shaped"
exc = ProxyException(message="budget exceeded", type=ProxyErrorTypes.budget_exceeded, param="key", code=402)
request = MagicMock()
request.headers = {"x-litellm-call-id": call_id}
request.body = AsyncMock(return_value=b'{"input": "hi"}')
fake_logging = MagicMock()
fake_logging.post_call_failure_hook = AsyncMock()
with (
patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
pytest.raises(ProxyException) as raised,
):
await proxy_server_module.moderations(
request=request,
fastapi_response=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0),
)
assert raised.value is exc
assert raised.value.code == "402"
assert raised.value.headers["x-litellm-call-id"] == call_id
@pytest.mark.asyncio
@pytest.mark.parametrize(
"exc",
[
HTTPException(status_code=401, detail="bad key"),
ProxyException(message="budget exceeded", type=ProxyErrorTypes.budget_exceeded, param="key", code=402),
],
ids=["http_exception", "proxy_exception"],
)
async def test_audio_speech_already_shaped_failure_answers_with_the_callers_litellm_call_id(exc: Exception):
"""LIT-7836: /v1/audio/speech re-raises HTTP and proxy shaped failures unchanged, and they must
still answer with the caller's x-litellm-call-id."""
call_id = "speech-call-7836-shaped"
request = MagicMock()
request.headers = {"x-litellm-call-id": call_id}
request.body = AsyncMock(return_value=b'{"model": "tts-1", "input": "hi", "voice": "alloy"}')
fake_logging = MagicMock()
fake_logging.post_call_failure_hook = AsyncMock()
with (
patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point
patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point
pytest.raises(type(exc)) as raised,
):
await proxy_server_module.audio_speech(
request=request,
fastapi_response=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0),
)
if isinstance(exc, HTTPException):
assert (raised.value.status_code, raised.value.detail) == (401, "bad key")
else:
assert raised.value is exc
assert raised.value.headers["x-litellm-call-id"] == call_id
@pytest.mark.asyncio
async def test_init_agents_in_db_rebuilds_registry_under_agent_reconcile_lock(monkeypatch):
from litellm.proxy.agent_endpoints.agent_registry import (

View file

@ -160,6 +160,37 @@ async def test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params():
assert "litellm_metadata" not in captured["optional_params"]
@pytest.mark.asyncio
async def test_proxy_only_error_log_keeps_the_request_litellm_call_id(monkeypatch: pytest.MonkeyPatch):
"""LIT-7836: a route that already stamped the caller's litellm_call_id must
keep it when the failure is a proxy-only error, so the spend-log row and the
error line share one id instead of a fresh uuid minted here."""
from litellm.litellm_core_utils.litellm_logging import Logging
call_id: Final = "caller-supplied-7836"
captured: dict[str, object] = {}
def fake_pre_call(self, *args, **kwargs):
captured["litellm_call_id"] = self.litellm_call_id
async def _noop_async_failure(self, *args, **kwargs):
return None
monkeypatch.setattr(Logging, "pre_call", fake_pre_call)
monkeypatch.setattr(Logging, "async_failure_handler", _noop_async_failure)
request_data: Final[dict[str, object]] = {"model": "gpt-4o", "input": "hi", "litellm_call_id": call_id}
await ProxyLogging(user_api_key_cache=DualCache())._handle_logging_proxy_only_error(
request_data=request_data,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-bad", request_route="/v1/moderations"),
route="/v1/moderations",
original_exception=Exception("bad key"),
)
assert request_data["litellm_call_id"] == call_id
assert captured["litellm_call_id"] == call_id
def test_get_model_group_info_order():
from litellm import Router
from litellm.proxy.proxy_server import _get_model_group_info

View file

@ -3266,3 +3266,174 @@ class TestPtuCostAttributionUISetting:
assert response.status_code == 400
assert "enable_ptu_cost_attribution" in str(response.json()["detail"])
assert not mock_prisma.db.litellm_uisettings.upsert.called
class TestTeamAdminEditableTeamFieldsSetting:
"""team_admin_editable_team_fields: the proxy-wide allow-list update_team applies to team admins."""
def _as_proxy_admin(self, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="test-user-123",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
mock_prisma = MagicMock()
mock_prisma.db.litellm_uisettings.upsert = AsyncMock()
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
return mock_prisma
def test_patch_rejects_field_names_the_proxy_does_not_support(self, monkeypatch):
mock_prisma = self._as_proxy_admin(monkeypatch)
monkeypatch.setattr(
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS",
frozenset({"tpm_limit"}),
)
try:
response = client.patch(
"/update/ui_settings",
json={"team_admin_editable_team_fields": ["tpm_limit", "blocked", "organization_id"]},
)
finally:
app.dependency_overrides.clear()
assert response.status_code == 400
detail = response.json()["detail"]["error"]
assert "['blocked', 'organization_id']" in detail
assert "['tpm_limit']" in detail
assert not mock_prisma.db.litellm_uisettings.upsert.called
def test_patch_rejects_a_non_list_value(self, monkeypatch):
self._as_proxy_admin(monkeypatch)
try:
response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": "tpm_limit"})
finally:
app.dependency_overrides.clear()
assert response.status_code == 422
def test_patch_persists_and_syncs_the_list_to_general_settings(self, monkeypatch):
mock_prisma = self._as_proxy_admin(monkeypatch)
general_settings: dict = {"team_admin_editable_team_fields": []}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
try:
response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": ["tpm_limit"]})
finally:
app.dependency_overrides.clear()
assert response.status_code == 200
stored = json.loads(mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["create"]["ui_settings"])
assert stored["team_admin_editable_team_fields"] == ["tpm_limit"]
assert general_settings["team_admin_editable_team_fields"] == ["tpm_limit"]
def test_patch_with_an_empty_list_turns_team_admin_editing_off_again(self, monkeypatch):
mock_prisma = self._as_proxy_admin(monkeypatch)
general_settings: dict = {"team_admin_editable_team_fields": ["tpm_limit"]}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
try:
response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": []})
finally:
app.dependency_overrides.clear()
assert response.status_code == 200
stored = json.loads(mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["create"]["ui_settings"])
assert stored["team_admin_editable_team_fields"] == []
assert general_settings["team_admin_editable_team_fields"] == []
def test_get_reports_the_stored_list_and_advertises_supported_fields(self, mock_auth, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
mock_prisma = MagicMock()
mock_db_record = MagicMock()
mock_db_record.ui_settings = {"team_admin_editable_team_fields": ["tpm_limit"]}
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=mock_db_record)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
general_settings: dict = {}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
response = client.get("/get/ui_settings")
assert response.status_code == 200
data = response.json()
assert data["values"]["team_admin_editable_team_fields"] == ["tpm_limit"]
assert general_settings["team_admin_editable_team_fields"] == ["tpm_limit"]
field_schema = data["field_schema"]["properties"]["team_admin_editable_team_fields"]
assert field_schema["type"] == "array"
assert field_schema["items"]["type"] == "string"
assert "tpm_limit" in field_schema["items"]["enum"]
class TestSyncUiSettingsToGeneralSettings:
"""The DB re-read each pod runs on startup and on every config reload."""
def _sync(self):
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
sync_ui_settings_to_general_settings,
)
return sync_ui_settings_to_general_settings
@pytest.mark.asyncio
async def test_applies_runtime_flags_and_leaves_other_ui_settings_alone(self, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
general_settings: dict = {"allow_agents_for_team_admins": False}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
mock_prisma = MagicMock()
record = MagicMock()
record.ui_settings = json.dumps(
{
"allow_agents_for_team_admins": True,
"team_admin_editable_team_fields": ["tpm_limit"],
"enable_chat_ui": False,
}
)
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=record)
applied = await self._sync()(mock_prisma)
assert dict(applied) == {
"allow_agents_for_team_admins": True,
"team_admin_editable_team_fields": ["tpm_limit"],
}
assert general_settings["allow_agents_for_team_admins"] is True
assert general_settings["team_admin_editable_team_fields"] == ["tpm_limit"]
assert "enable_chat_ui" not in general_settings
@pytest.mark.asyncio
async def test_reads_a_row_the_prisma_client_already_deserialized(self, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
general_settings: dict = {}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
mock_prisma = MagicMock()
record = MagicMock()
record.ui_settings = {"team_admin_editable_team_fields": ["rpm_limit"]}
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=record)
await self._sync()(mock_prisma)
assert general_settings["team_admin_editable_team_fields"] == ["rpm_limit"]
@pytest.mark.asyncio
async def test_without_a_stored_row_general_settings_is_left_untouched(self, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
general_settings: dict = {"allow_agents_for_team_admins": True}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
mock_prisma = MagicMock()
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None)
applied = await self._sync()(mock_prisma)
assert dict(applied) == {}
assert general_settings == {"allow_agents_for_team_admins": True}

View file

@ -176,6 +176,25 @@ def test_handle_exception_on_proxy_error_path_none_input_wraps_as_500():
}
@pytest.mark.parametrize(
"exc",
[
HTTPException(status_code=401, detail="bad key"),
ValueError("provider boom"),
ProxyException(message="already wrapped", type=ProxyErrorTypes.budget_exceeded.value, param="key", code=402),
],
ids=["http_exception", "generic_exception", "already_proxy_exception"],
)
def test_handle_exception_on_proxy_returns_the_litellm_call_id_header(exc: Exception):
result = handle_exception_on_proxy(exc, "call-7836")
assert result.headers == {"x-litellm-call-id": "call-7836"}
def test_handle_exception_on_proxy_sends_no_call_id_header_when_the_request_has_none():
assert handle_exception_on_proxy(ValueError("provider boom")).headers == {}
@pytest.mark.asyncio
async def test_handle_exception_on_proxy_read_only_transaction_forces_writer_recreate(
monkeypatch: pytest.MonkeyPatch,

View file

@ -115,13 +115,13 @@ def _respx_interceptable_httpx_client(monkeypatch):
],
)
def test_resolver_opt_in_gates_openai_like_config(model_info, expected_type):
config = _resolve_responses_api_provider_config("my-model", "custom_openai", model_info)
config = _resolve_responses_api_provider_config("my-model", "custom_openai", model_info, None)
assert type(config) is expected_type
def test_resolver_keeps_native_provider_config():
"""`openai/` already routes /v1/responses natively; the opt-in must not swap its config."""
config = _resolve_responses_api_provider_config("gpt-4.1", "openai", OPT_IN)
config = _resolve_responses_api_provider_config("gpt-4.1", "openai", OPT_IN, None)
assert type(config) is OpenAIResponsesAPIConfig

View file

@ -17,6 +17,7 @@ import SCIMConfig from "@/components/SCIM";
import LoggingSettings from "@/components/Settings/AdminSettings/LoggingSettings/LoggingSettings";
import SSOSettings from "@/components/Settings/AdminSettings/SSOSettings/SSOSettings";
import UISettings from "@/components/Settings/AdminSettings/UISettings/UISettings";
import TeamAdminEditableFieldsSettings from "@/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings";
import UserBannerSettings from "@/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings";
import CyberArk from "@/components/Settings/AdminSettings/CyberArk/CyberArk";
import HashicorpVault from "@/components/Settings/AdminSettings/HashicorpVault/HashicorpVault";
@ -382,6 +383,7 @@ const AdminPanel: React.FC<AdminPanelProps> = ({ proxySettings }) => {
children: (
<div className="flex flex-col gap-4">
<UISettings />
<TeamAdminEditableFieldsSettings />
<UserBannerSettings />
</div>
),

View file

@ -0,0 +1,175 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
import { fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils";
import { toast } from "@/lib/toast";
import TeamAdminEditableFieldsSettings from "./TeamAdminEditableFieldsSettings";
const mockUseUISettings = vi.hoisted(() => vi.fn());
const mockUseUpdateUISettings = vi.hoisted(() => vi.fn());
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
default: () => ({ accessToken: "test-token" }),
}));
vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({
useUISettings: mockUseUISettings,
}));
vi.mock("@/app/(dashboard)/hooks/uiSettings/useUpdateUISettings", () => ({
useUpdateUISettings: mockUseUpdateUISettings,
}));
const TPM_LABEL = "Tokens per minute Limit (TPM)";
const mockSettings = (supported: readonly string[], enabled: readonly string[]) =>
mockUseUISettings.mockReturnValue({
isLoading: false,
data: {
field_schema: {
properties: {
team_admin_editable_team_fields: {
description: "Fields a team admin may change",
items: { type: "string", enum: supported },
},
},
},
values: { team_admin_editable_team_fields: enabled },
},
});
const mockSave = ({
isPending = false,
outcome = "success",
}: {
isPending?: boolean;
outcome?: "success" | "error";
}) => {
const mutate = vi.fn((_settings: unknown, options: { onSuccess: () => void; onError: (error: Error) => void }) =>
outcome === "success" ? options.onSuccess() : options.onError(new Error("save failed")),
);
mockUseUpdateUISettings.mockReturnValue({ mutate, isPending });
return mutate;
};
const saveButton = () => screen.getByRole("button", { name: "Save" });
describe("TeamAdminEditableFieldsSettings", () => {
beforeEach(() => {
vi.clearAllMocks();
});
it("explains that nothing can be enabled when the proxy supports no fields", () => {
mockSettings([], []);
mockSave({});
renderWithProviders(<TeamAdminEditableFieldsSettings />);
expect(screen.getByText("Team admins cannot edit team settings")).toBeInTheDocument();
expect(screen.getByText(/does not support enabling any team settings fields/)).toBeInTheDocument();
expect(screen.queryByRole("checkbox")).not.toBeInTheDocument();
expect(screen.queryByRole("button", { name: "Save" })).not.toBeInTheDocument();
});
it("renders one checkbox per supported field, checked for the saved ones, with Save disabled until something changes", () => {
mockSettings(["max_budget", "tpm_limit"], ["tpm_limit"]);
mockSave({});
renderWithProviders(<TeamAdminEditableFieldsSettings />);
expect(screen.getByText("Team admin editable fields")).toBeInTheDocument();
expect(screen.getByText("1 field enabled")).toBeInTheDocument();
expect(screen.getByText("Fields a team admin may change")).toBeInTheDocument();
expect(screen.getByRole("checkbox", { name: "max_budget" })).not.toBeChecked();
expect(screen.getByRole("checkbox", { name: TPM_LABEL })).toBeChecked();
expect(saveButton()).toBeDisabled();
});
it("only saves a ticked field once Save is clicked", async () => {
mockSettings(["max_budget", "tpm_limit"], ["tpm_limit"]);
const mutate = mockSave({});
renderWithProviders(<TeamAdminEditableFieldsSettings />);
fireEvent.click(screen.getByRole("checkbox", { name: "max_budget" }));
expect(screen.getByRole("checkbox", { name: "max_budget" })).toBeChecked();
expect(mutate).not.toHaveBeenCalled();
fireEvent.click(saveButton());
await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Team admin editable fields updated successfully"));
expect(mutate).toHaveBeenCalledWith(
{ team_admin_editable_team_fields: ["max_budget", "tpm_limit"] },
expect.anything(),
);
expect(saveButton()).toBeDisabled();
});
it("saves the list without an unticked field", async () => {
mockSettings(["max_budget", "tpm_limit"], ["max_budget", "tpm_limit"]);
const mutate = mockSave({});
renderWithProviders(<TeamAdminEditableFieldsSettings />);
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
fireEvent.click(saveButton());
await waitFor(() => expect(mutate).toHaveBeenCalledTimes(1));
expect(mutate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["max_budget"] }, expect.anything());
});
it("disables Save again when the draft is ticked back to the saved list", () => {
mockSettings(["tpm_limit"], []);
mockSave({});
renderWithProviders(<TeamAdminEditableFieldsSettings />);
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
expect(saveButton()).toBeEnabled();
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
expect(screen.getByRole("checkbox", { name: TPM_LABEL })).not.toBeChecked();
expect(saveButton()).toBeDisabled();
});
it("treats a saved list in another order, or with fields this proxy dropped, as the same selection", () => {
mockSettings(["max_budget", "tpm_limit"], ["tpm_limit", "retired_field", "max_budget"]);
mockSave({});
renderWithProviders(<TeamAdminEditableFieldsSettings />);
expect(screen.getByText("2 fields enabled")).toBeInTheDocument();
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
expect(saveButton()).toBeDisabled();
});
it("keeps the draft and shows the error when the save fails", async () => {
mockSettings(["tpm_limit"], []);
const mutate = mockSave({ outcome: "error" });
renderWithProviders(<TeamAdminEditableFieldsSettings />);
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
fireEvent.click(saveButton());
await waitFor(() => expect(toast.fromError).toHaveBeenCalledTimes(1));
expect(mutate).toHaveBeenCalledTimes(1);
expect(toast.success).not.toHaveBeenCalled();
expect(screen.getByRole("checkbox", { name: TPM_LABEL })).toBeChecked();
expect(saveButton()).toBeEnabled();
});
it("blocks ticking and saving while a save is in flight", () => {
mockSettings(["tpm_limit"], []);
const mutate = mockSave({ isPending: true });
renderWithProviders(<TeamAdminEditableFieldsSettings />);
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
expect(screen.getByRole("checkbox", { name: TPM_LABEL })).not.toBeChecked();
expect(screen.getByRole("button", { name: "Saving..." })).toBeDisabled();
expect(mutate).not.toHaveBeenCalled();
});
});

View file

@ -0,0 +1,138 @@
"use client";
import { Controller } from "react-hook-form";
import { z } from "zod/v4";
import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
import { useUpdateUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUpdateUISettings";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import {
parseSupportedTeamAdminEditableFields,
parseTeamAdminEditableFields,
teamAdminFieldLabel,
} from "@/components/team/teamAdminEditAccess";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
import { Checkbox } from "@/components/ui/checkbox";
import { Skeleton } from "@/components/ui/skeleton";
import { useZodForm } from "@/lib/forms/useZodForm";
import { toast } from "@/lib/toast";
const editableFieldsSchema = z.object({ team_admin_editable_team_fields: z.array(z.string()) });
type SaveEditableFields = ReturnType<typeof useUpdateUISettings>["mutate"];
export default function TeamAdminEditableFieldsSettings() {
const { accessToken } = useAuthorized();
const { data, isLoading } = useUISettings();
const { mutate: saveSettings, isPending } = useUpdateUISettings(accessToken);
const supportedFields = parseSupportedTeamAdminEditableFields(data?.field_schema);
const savedFields = parseTeamAdminEditableFields(data?.values);
const enabledFields = supportedFields.filter((field) => savedFields.includes(field));
return (
<Card>
<CardHeader>
<div className="flex items-center gap-2">
<CardTitle>Team admin editable fields</CardTitle>
<Badge variant={enabledFields.length > 0 ? "secondary" : "outline"}>
{enabledFields.length > 0
? `${enabledFields.length} field${enabledFields.length !== 1 ? "s" : ""} enabled`
: "Team admins cannot edit team settings"}
</Badge>
</div>
<CardDescription>
{data?.field_schema?.properties?.team_admin_editable_team_fields?.description ??
"Team settings fields a team admin may change on the teams they administer."}
</CardDescription>
</CardHeader>
<CardContent>
{isLoading ? (
<Skeleton className="h-16 w-full" />
) : (
<TeamAdminEditableFieldsForm
key={enabledFields.join(",")}
enabledFields={enabledFields}
supportedFields={supportedFields}
isPending={isPending}
saveSettings={saveSettings}
/>
)}
</CardContent>
</Card>
);
}
interface TeamAdminEditableFieldsFormProps {
enabledFields: readonly string[];
supportedFields: readonly string[];
isPending: boolean;
saveSettings: SaveEditableFields;
}
function TeamAdminEditableFieldsForm({
enabledFields,
supportedFields,
isPending,
saveSettings,
}: TeamAdminEditableFieldsFormProps) {
const form = useZodForm(editableFieldsSchema, {
defaultValues: { team_admin_editable_team_fields: [...enabledFields] },
});
const submit = form.handleSubmit((values) =>
saveSettings(values, {
onSuccess: () => {
form.reset(values);
toast.success("Team admin editable fields updated successfully");
},
onError: (error) => {
toast.fromError(error);
},
}),
);
if (supportedFields.length === 0) {
return (
<p className="text-sm italic text-muted-foreground">
This proxy version does not support enabling any team settings fields for team admins yet.
</p>
);
}
return (
<form onSubmit={(event) => void submit(event)} className="space-y-4">
<Controller
control={form.control}
name="team_admin_editable_team_fields"
render={({ field }) => (
<div className="space-y-2">
{supportedFields.map((name) => {
const checkboxId = `team-admin-editable-${name}`;
return (
<label key={name} htmlFor={checkboxId} className="flex cursor-pointer items-center gap-2">
<Checkbox
id={checkboxId}
checked={field.value.includes(name)}
disabled={isPending}
onCheckedChange={(checked) =>
field.onChange(
supportedFields.filter((item) => (item === name ? checked : field.value.includes(item))),
)
}
/>
<span className="text-sm text-foreground">{teamAdminFieldLabel(name)}</span>
</label>
);
})}
</div>
)}
/>
<div className="flex justify-end">
<Button type="submit" disabled={isPending || !form.formState.isDirty}>
{isPending ? "Saving..." : "Save"}
</Button>
</div>
</form>
);
}

View file

@ -145,7 +145,7 @@ function UserBannerSettingsForm({ persisted, isLoading, isPending, saveBanner }:
</div>
)}
<div>
<div className="flex justify-end">
<Button onClick={handleSave} disabled={isPending || messageMissing}>
{isPending ? "Saving..." : "Save banner"}
</Button>

View file

@ -0,0 +1,85 @@
import { describe, expect, it, vi } from "vitest";
import userEvent from "@testing-library/user-event";
import { fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils";
import TeamAdminSettingsForm from "./TeamAdminSettingsForm";
const renderForm = (editableFields: ReadonlySet<string>, overrides: { isSaving?: boolean } = {}) => {
const onSave = vi.fn().mockResolvedValue(undefined);
const onCancel = vi.fn();
renderWithProviders(
<TeamAdminSettingsForm
initialValues={{ tpm_limit: 1000 }}
editableFields={editableFields}
isSaving={overrides.isSaving ?? false}
onCancel={onCancel}
onSave={onSave}
/>,
);
return { onSave, onCancel };
};
describe("TeamAdminSettingsForm", () => {
it("shows the team's current TPM limit when the proxy lets team admins edit it", () => {
renderForm(new Set(["tpm_limit"]));
expect(screen.getByLabelText("Tokens per minute Limit (TPM)")).toHaveValue(1000);
});
it("hides the TPM limit when the proxy has not enabled it for team admins", () => {
renderForm(new Set(["max_budget"]));
expect(screen.queryByLabelText("Tokens per minute Limit (TPM)")).not.toBeInTheDocument();
});
it("saves the new TPM limit and nothing else", async () => {
const user = userEvent.setup();
const { onSave } = renderForm(new Set(["tpm_limit"]));
fireEvent.change(screen.getByLabelText("Tokens per minute Limit (TPM)"), { target: { value: "5000" } });
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(onSave).toHaveBeenCalledWith({ tpm_limit: 5000 }));
});
it("saves a cleared TPM limit as no limit", async () => {
const user = userEvent.setup();
const { onSave } = renderForm(new Set(["tpm_limit"]));
fireEvent.change(screen.getByLabelText("Tokens per minute Limit (TPM)"), { target: { value: "" } });
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(onSave).toHaveBeenCalledWith({ tpm_limit: null }));
});
it("keeps Save disabled until the TPM limit differs from the team's", () => {
renderForm(new Set(["tpm_limit"]));
const tpmInput = screen.getByLabelText("Tokens per minute Limit (TPM)");
const save = screen.getByRole("button", { name: /save changes/i });
expect(save).toBeDisabled();
fireEvent.change(tpmInput, { target: { value: "5000" } });
expect(save).toBeEnabled();
fireEvent.change(tpmInput, { target: { value: "1000" } });
expect(save).toBeDisabled();
});
it("closes without saving on cancel", async () => {
const user = userEvent.setup();
const { onSave, onCancel } = renderForm(new Set(["tpm_limit"]));
await user.click(screen.getByRole("button", { name: "Cancel" }));
expect(onCancel).toHaveBeenCalledTimes(1);
expect(onSave).not.toHaveBeenCalled();
});
it("locks both buttons while a save is in flight", () => {
renderForm(new Set(["tpm_limit"]), { isSaving: true });
fireEvent.change(screen.getByLabelText("Tokens per minute Limit (TPM)"), { target: { value: "5000" } });
expect(screen.getByRole("button", { name: "Cancel" })).toBeDisabled();
expect(screen.getByRole("button", { name: /save changes/i })).toBeDisabled();
});
});

View file

@ -0,0 +1,69 @@
"use client";
import { Save } from "lucide-react";
import { useWatch } from "react-hook-form";
import { z } from "zod/v4";
import { FormField } from "@/components/shared/form/FormField";
import { Button } from "@/components/ui/button";
import { FieldGroup } from "@/components/ui/field";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { useZodForm } from "@/lib/forms/useZodForm";
import NumericalInput from "../shared/numerical_input";
import {
teamAdminFieldLabel,
teamAdminSettingsChanges,
type TeamAdminSettingsChanges,
type TeamAdminSettingsValues,
} from "./teamAdminEditAccess";
const teamAdminSettingsSchema = z.object({
tpm_limit: z.union([z.string(), z.number()]).nullish(),
});
interface TeamAdminSettingsFormProps {
initialValues: TeamAdminSettingsValues;
editableFields: ReadonlySet<string>;
isSaving: boolean;
onCancel: () => void;
onSave: (changes: TeamAdminSettingsChanges) => Promise<void>;
}
export default function TeamAdminSettingsForm({
initialValues,
editableFields,
isSaving,
onCancel,
onSave,
}: TeamAdminSettingsFormProps) {
const form = useZodForm(teamAdminSettingsSchema, { defaultValues: initialValues });
const draft = useWatch({ control: form.control });
const hasChanges = Object.keys(teamAdminSettingsChanges(draft, initialValues, editableFields)).length > 0;
const submit = form.handleSubmit((values) => onSave(teamAdminSettingsChanges(values, initialValues, editableFields)));
return (
<form onSubmit={(event) => void submit(event)}>
<FieldGroup>
<p className="text-sm text-muted-foreground">
A proxy admin chose which settings team admins can change. Ask a proxy admin to change anything else.
</p>
{editableFields.has("tpm_limit") && (
<FormField control={form.control} name="tpm_limit" label={teamAdminFieldLabel("tpm_limit")}>
{({ ref, value, ...field }) => <NumericalInput {...field} ref={ref} value={value ?? ""} step={1} />}
</FormField>
)}
</FieldGroup>
<div className="mt-6 flex items-center justify-end gap-2">
<Button type="button" variant="outline" onClick={onCancel} disabled={isSaving}>
Cancel
</Button>
<Button type="submit" disabled={isSaving || !hasChanges}>
{isSaving ? <UiLoadingSpinner className="size-4" /> : <Save className="size-4" />}
Save Changes
</Button>
</div>
</form>
);
}

View file

@ -69,6 +69,10 @@ vi.mock("@/app/(dashboard)/hooks/teams/useTeamMetadataSchema", () => ({
useTeamMetadataSchema: vi.fn(() => ({ data: [], isLoading: false })),
}));
vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({
useUISettings: vi.fn(),
}));
vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({
useAllProxyModels: vi.fn(),
}));
@ -228,6 +232,7 @@ import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser";
import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers";
import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets";
import { useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups";
import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
const mockUseAllProxyModels = vi.mocked(useAllProxyModels);
const mockUseKeys = vi.mocked(useKeys);
@ -237,6 +242,7 @@ const mockUseCurrentUser = vi.mocked(useCurrentUser);
const mockUseMCPServers = vi.mocked(useMCPServers);
const mockUseMCPToolsets = vi.mocked(useMCPToolsets);
const mockUseAccessGroups = vi.mocked(useAccessGroups);
const mockUseUISettings = vi.mocked(useUISettings);
const createMockTeamData = (overrides = {}) => ({
team_id: "123",
@ -305,6 +311,10 @@ const seedDefaultMocks = () => {
isLoading: false,
isError: false,
} as any);
mockUseUISettings.mockReturnValue({
data: { values: {} },
isLoading: false,
} as any);
mockUseKeys.mockReturnValue({
data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 },
isPending: false,
@ -656,19 +666,9 @@ describe("TeamInfoView", () => {
});
});
it("shows edit tabs when the fetched team data marks the session user as team admin, even without the is_team_admin prop", async () => {
it("shows edit tabs when the proxy reports the session user may edit, even without the is_team_admin prop", async () => {
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({
members_with_roles: [
{
user_id: "user-1",
user_email: "admin@test.com",
role: "admin",
spend: 0,
budget_id: "budget1",
},
],
}),
createMockTeamData({ caller_edit_access: { kind: "team_admin_disabled" } }),
);
renderWithProviders(<TeamInfoView {...defaultProps} is_team_admin={false} is_proxy_admin={false} />);
@ -1863,6 +1863,92 @@ describe("TeamInfoView", () => {
});
});
});
describe("team admin edit access", () => {
const teamAdminProps = { ...defaultProps, is_proxy_admin: false, is_team_admin: true };
beforeEach(() => {
authState.userRole = "Internal User";
});
it("tells a team admin to ask a proxy admin when the proxy reports no team field is enabled for them", async () => {
const user = userEvent.setup({ delay: null });
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({ caller_edit_access: { kind: "team_admin_disabled" } }),
);
renderWithProviders(<TeamInfoView {...teamAdminProps} />);
await user.click(await screen.findByRole("tab", { name: "Settings" }));
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
expect(toast.error).toHaveBeenCalledWith("Team admins cannot edit team settings on this proxy", {
description: "Ask a proxy admin to enable fields under Settings > UI > Team admin editable fields.",
});
expect(screen.queryByLabelText("Team Name")).not.toBeInTheDocument();
expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument();
});
it("gives a team admin only the fields the proxy enabled and sends only those on save", async () => {
const user = userEvent.setup({ delay: null });
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({
tpm_limit: 1000,
caller_edit_access: { kind: "team_admin", editable_fields: ["tpm_limit"] },
}),
);
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any);
renderWithProviders(<TeamInfoView {...teamAdminProps} />);
await user.click(await screen.findByRole("tab", { name: "Settings" }));
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
const tpmInput = await screen.findByLabelText("Tokens per minute Limit (TPM)");
expect(tpmInput).toHaveValue(1000);
expect(screen.queryByLabelText("Team Name")).not.toBeInTheDocument();
expect(screen.queryByLabelText("Requests per minute Limit (RPM)")).not.toBeInTheDocument();
fireEvent.change(tpmInput, { target: { value: "5000" } });
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalledTimes(1));
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1]).toStrictEqual({ team_id: "123", tpm_limit: 5000 });
expect(toast.success).toHaveBeenCalledWith("Team settings updated successfully");
expect(toast.error).not.toHaveBeenCalled();
});
it("opens the form when the proxy reports unrestricted access although the props only mark a team admin", async () => {
const user = userEvent.setup({ delay: null });
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({ caller_edit_access: { kind: "unrestricted" } }),
);
renderWithProviders(<TeamInfoView {...teamAdminProps} />);
await user.click(await screen.findByRole("tab", { name: "Settings" }));
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
expect(await screen.findByLabelText("Team Name")).toBeInTheDocument();
expect(toast.error).not.toHaveBeenCalled();
});
it("never gates a proxy admin on the team admin field list", async () => {
authState.userRole = "Admin";
const user = userEvent.setup({ delay: null });
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({ caller_edit_access: { kind: "unrestricted" } }),
);
renderWithProviders(<TeamInfoView {...defaultProps} />);
await user.click(await screen.findByRole("tab", { name: "Settings" }));
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
expect(await screen.findByLabelText("Team Name")).toBeInTheDocument();
expect(toast.error).not.toHaveBeenCalled();
});
});
});
describe("TeamInfoView - which team member fields reach the update payload depends on the open sections", () => {

View file

@ -48,6 +48,14 @@ import React, { useEffect, useMemo, useState } from "react";
import { useFieldArray } from "react-hook-form";
import { z } from "zod/v4";
import GuardrailsSelect from "./GuardrailsSelect";
import {
type CallerEditAccess,
parseTeamEditAccess,
TEAM_ADMIN_EDITING_DISABLED_DESCRIPTION,
TEAM_ADMIN_EDITING_DISABLED_TITLE,
type TeamAdminSettingsChanges,
} from "./teamAdminEditAccess";
import TeamAdminSettingsForm from "./TeamAdminSettingsForm";
import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
import AccessGroupSelector from "../common_components/AccessGroupSelector";
import BudgetDurationDropdown, { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown";
@ -297,6 +305,7 @@ export interface TeamData {
guardrails?: string[];
policies?: string[];
object_permission?: ObjectPermission | null;
caller_edit_access?: CallerEditAccess;
team_member_budget_table: {
max_budget: number;
budget_duration: string | null;
@ -315,7 +324,6 @@ export interface TeamInfoProps {
accessToken: string | null;
is_team_admin: boolean;
is_proxy_admin: boolean;
is_org_admin?: boolean;
userModels: string[];
editTeam: boolean;
premiumUser?: boolean;
@ -531,7 +539,6 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
accessToken,
is_team_admin,
is_proxy_admin,
is_org_admin = false,
userModels,
editTeam,
premiumUser = false,
@ -575,7 +582,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
const [teamModelMaxBudget, setTeamModelMaxBudget] = useState<ModelMaxBudget>({});
const routerSettingsRef = React.useRef<RouterSettingsAccordionRef>(null);
const [organization, setOrganization] = useState<Organization | null>(null);
const { userRole, userId } = useAuthorized();
const { userRole } = useAuthorized();
const { data: allMcpServers = [], isError: mcpServersFailed, isLoading: mcpServersLoading } = useMCPServers();
const { data: allMcpToolsets = [], isError: mcpToolsetsFailed, isLoading: mcpToolsetsLoading } = useMCPToolsets();
const { data: allAccessGroups = [], isError: accessGroupsFailed, isLoading: accessGroupsLoading } = useAccessGroups();
@ -585,14 +592,6 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
const { data: teamMetadataSchemaFields = [], isLoading: isTeamMetadataSchemaLoading } = useTeamMetadataSchema();
const queryClient = useQueryClient();
// Check if user is org admin for this team's organization
const isOrgAdminForTeam = useMemo(() => {
const teamOrgId = teamData?.team_info?.organization_id;
if (!teamOrgId || !userId) return false;
const org = userOrganizations.find((o) => o.organization_id === teamOrgId);
return org?.members?.some((m: any) => m.user_id === userId && m.user_role === "org_admin") ?? false;
}, [teamData, userOrganizations, userId]);
// Models currently selected in the team edit form, used to scope the per-model
// rate limit dropdown to models this team actually has access to.
const watchedModels = form.watch("models");
@ -616,15 +615,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
return unfurlWildcardModelsInList(selected, userModels);
}, [watchedModels, teamData, userModels]);
const isTeamAdminFromTeamData = useMemo(
() =>
teamData?.team_info?.members_with_roles?.some(
(member) => member.user_id != null && member.user_id === userId && member.role === "admin",
) ?? false,
[teamData, userId],
);
const canEditTeam = is_team_admin || is_proxy_admin || is_org_admin || isOrgAdminForTeam || isTeamAdminFromTeamData;
const teamEditAccess = useMemo(() => parseTeamEditAccess(teamData?.team_info?.caller_edit_access), [teamData]);
const canEditTeam = is_team_admin || is_proxy_admin || teamEditAccess.kind !== "none";
const visibleTabs = useMemo(() => getTeamInfoVisibleTabs(canEditTeam), [canEditTeam]);
const defaultTabKey = useMemo(() => getTeamInfoDefaultTab(editTeam, canEditTeam), [editTeam, canEditTeam]);
const { onTabChange, hasVisited } = useVisitedTabs(defaultTabKey);
@ -644,6 +636,15 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
setIsEditing(true);
};
const openSettingsEditor = (modelAliases: Record<string, string>) => {
if (teamEditAccess.kind === "team_admin_disabled") {
toast.error(TEAM_ADMIN_EDITING_DISABLED_TITLE, { description: TEAM_ADMIN_EDITING_DISABLED_DESCRIPTION });
return;
}
setTeamModelAliases(modelAliases);
startEditing();
};
const applyKillSwitchToGuardrails = (checked: boolean) => {
const current = form.getValues("guardrails") ?? [];
const nonGlobals = current.filter((name) => !globalGuardrailNames.has(name));
@ -863,6 +864,27 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
setMemberToDelete(null);
};
const persistTeamUpdate = async (token: string, updateData: Record<string, unknown>) => {
await teamUpdateCall(token, updateData);
queryClient.invalidateQueries({ queryKey: organizationKeys.all });
toast.success("Team settings updated successfully");
setIsEditing(false);
fetchTeamInfo();
};
const saveTeamAdminSettings = async (changes: TeamAdminSettingsChanges) => {
if (!accessToken) return;
setIsTeamSaving(true);
try {
await persistTeamUpdate(accessToken, { team_id: teamId, ...changes });
} catch (error) {
console.error("Error updating team:", error);
} finally {
setIsTeamSaving(false);
}
};
const handleTeamUpdate = async (values: any) => {
try {
if (!accessToken) return;
@ -1113,12 +1135,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
}
}
await teamUpdateCall(accessToken, updateData);
queryClient.invalidateQueries({ queryKey: organizationKeys.all });
toast.success("Team settings updated successfully");
setIsEditing(false);
fetchTeamInfo();
await persistTeamUpdate(accessToken, updateData);
} catch (error) {
console.error("Error updating team:", error);
} finally {
@ -1136,6 +1153,17 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
const { team_info: info } = teamData;
const teamAdminSettingsEditor =
teamEditAccess.kind === "team_admin" ? (
<TeamAdminSettingsForm
initialValues={{ tpm_limit: info.tpm_limit }}
editableFields={teamEditAccess.editableFields}
isSaving={isTeamSaving}
onCancel={() => setIsEditing(false)}
onSave={saveTeamAdminSettings}
/>
) : null;
const inheritedMcpServers = computeInheritedGrants(
info.access_group_mcp_server_ids,
info.access_group_details,
@ -1340,10 +1368,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
{canEditTeam && !isEditing && (
<Button
variant="outline"
onClick={() => {
setTeamModelAliases(info.litellm_model_table?.model_aliases ?? {});
startEditing();
}}
onClick={() => openSettingsEditor(info.litellm_model_table?.model_aliases ?? {})}
>
<Pencil />
Edit Settings
@ -1351,8 +1376,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
)}
</div>
{isEditing && isGuardrailsLoading ? (
<div className="p-4">Loading...</div>
{isEditing && (teamAdminSettingsEditor !== null || isGuardrailsLoading) ? (
teamAdminSettingsEditor ?? <div className="p-4">Loading...</div>
) : isEditing ? (
<TooltipProvider>
<form onSubmit={(event) => void form.handleSubmit(onTeamUpdateSubmit)(event)}>

View file

@ -0,0 +1,117 @@
import { describe, expect, it } from "vitest";
import {
parseSupportedTeamAdminEditableFields,
parseTeamAdminEditableFields,
parseTeamEditAccess,
teamAdminFieldLabel,
teamAdminSettingsChanges,
} from "./teamAdminEditAccess";
describe("teamAdminFieldLabel", () => {
it("names tpm_limit the way the team settings form does", () => {
expect(teamAdminFieldLabel("tpm_limit")).toBe("Tokens per minute Limit (TPM)");
});
it("falls back to the raw field name for a field the dashboard has no label for", () => {
expect(teamAdminFieldLabel("max_budget")).toBe("max_budget");
});
});
describe("teamAdminSettingsChanges", () => {
const tpmEnabled = new Set(["tpm_limit"]);
const stored = { tpm_limit: 1000 };
it.each([
["a typed number string", "5000", 5000],
["a number", 1200, 1200],
["zero", "0", 0],
["an emptied input", "", null],
["whitespace", " ", null],
["no limit", null, null],
["an unset value", undefined, null],
])("sends tpm_limit changed to %s", (_label, tpm_limit, expected) => {
expect(teamAdminSettingsChanges({ tpm_limit }, stored, tpmEnabled)).toStrictEqual({ tpm_limit: expected });
});
it.each([
["the stored number", 1000, { tpm_limit: 1000 }],
["the stored number typed back in", "1000", { tpm_limit: 1000 }],
["an emptied input over no stored limit", "", { tpm_limit: null }],
["an unset value over no stored limit", undefined, { tpm_limit: null }],
])("sends nothing for %s", (_label, tpm_limit, initialValues) => {
expect(teamAdminSettingsChanges({ tpm_limit }, initialValues, tpmEnabled)).toStrictEqual({});
});
it("leaves tpm_limit out when the proxy did not enable it for team admins", () => {
expect(teamAdminSettingsChanges({ tpm_limit: "5000" }, stored, new Set(["max_budget"]))).toStrictEqual({});
});
});
describe("parseTeamAdminEditableFields", () => {
it("returns the configured list", () => {
expect(parseTeamAdminEditableFields({ team_admin_editable_team_fields: ["tpm_limit", "rpm_limit"] })).toEqual([
"tpm_limit",
"rpm_limit",
]);
});
it.each([
["no values yet", undefined],
["setting missing", {}],
["setting is null", { team_admin_editable_team_fields: null }],
["setting is a string", { team_admin_editable_team_fields: "tpm_limit" }],
["list holds a non-string", { team_admin_editable_team_fields: ["tpm_limit", 7] }],
])("fails closed to an empty list when %s", (_label, values) => {
expect(parseTeamAdminEditableFields(values)).toEqual([]);
});
});
describe("parseSupportedTeamAdminEditableFields", () => {
it("reads the enum the proxy advertises on the setting's items schema", () => {
const schema = {
properties: {
team_admin_editable_team_fields: {
type: "array",
items: { type: "string", enum: ["max_budget", "tpm_limit"] },
},
},
};
expect(parseSupportedTeamAdminEditableFields(schema)).toEqual(["max_budget", "tpm_limit"]);
});
it.each([
["schema not loaded", undefined],
["property absent", { properties: {} }],
["items has no enum", { properties: { team_admin_editable_team_fields: { items: { type: "string" } } } }],
["enum is not a string list", { properties: { team_admin_editable_team_fields: { items: { enum: [1] } } } }],
])("returns no supported fields when %s", (_label, schema) => {
expect(parseSupportedTeamAdminEditableFields(schema)).toEqual([]);
});
});
describe("parseTeamEditAccess", () => {
it.each([
["unrestricted", { kind: "unrestricted" }],
["team_admin_disabled", { kind: "team_admin_disabled" }],
["none", { kind: "none" }],
])("passes the proxy's %s verdict through", (_kind, verdict) => {
expect(parseTeamEditAccess(verdict)).toEqual(verdict);
});
it("hands a team admin the fields the proxy enabled", () => {
expect(parseTeamEditAccess({ kind: "team_admin", editable_fields: ["tpm_limit"] })).toEqual({
kind: "team_admin",
editableFields: new Set(["tpm_limit"]),
});
});
it.each([
["the proxy sent nothing", undefined],
["the kind is unknown", { kind: "owner" }],
["a team admin verdict lacks its field list", { kind: "team_admin" }],
["the field list holds a non-string", { kind: "team_admin", editable_fields: [7] }],
])("fails closed to no access when %s", (_label, value) => {
expect(parseTeamEditAccess(value)).toEqual({ kind: "none" });
});
});

View file

@ -0,0 +1,78 @@
import { z } from "zod/v4";
export const TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING = "team_admin_editable_team_fields";
export const TEAM_ADMIN_EDITING_DISABLED_TITLE = "Team admins cannot edit team settings on this proxy";
export const TEAM_ADMIN_EDITING_DISABLED_DESCRIPTION =
"Ask a proxy admin to enable fields under Settings > UI > Team admin editable fields.";
const callerEditAccessSchema = z.discriminatedUnion("kind", [
z.object({ kind: z.literal("unrestricted") }),
z.object({ kind: z.literal("team_admin"), editable_fields: z.array(z.string()) }),
z.object({ kind: z.literal("team_admin_disabled") }),
z.object({ kind: z.literal("none") }),
]);
export type CallerEditAccess = z.infer<typeof callerEditAccessSchema>;
export type TeamEditAccess =
| { readonly kind: "unrestricted" }
| { readonly kind: "team_admin"; readonly editableFields: ReadonlySet<string> }
| { readonly kind: "team_admin_disabled" }
| { readonly kind: "none" };
const fieldListSchema = z.array(z.string()).catch([]);
export const parseTeamAdminEditableFields = (uiSettingsValues: unknown): readonly string[] => {
const values = z.record(z.string(), z.unknown()).catch({}).parse(uiSettingsValues);
return fieldListSchema.parse(values[TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING]);
};
export const parseSupportedTeamAdminEditableFields = (uiSettingsFieldSchema: unknown): readonly string[] => {
const property = z
.object({ properties: z.object({ [TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING]: z.object({ items: z.unknown() }) }) })
.safeParse(uiSettingsFieldSchema);
if (!property.success) return [];
const items = z
.object({ enum: z.unknown() })
.safeParse(property.data.properties[TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING].items);
return items.success ? fieldListSchema.parse(items.data.enum) : [];
};
const TEAM_ADMIN_FIELD_LABELS: ReadonlyMap<string, string> = new Map([["tpm_limit", "Tokens per minute Limit (TPM)"]]);
export const teamAdminFieldLabel = (field: string): string => TEAM_ADMIN_FIELD_LABELS.get(field) ?? field;
export interface TeamAdminSettingsValues {
readonly tpm_limit?: string | number | null;
}
export interface TeamAdminSettingsChanges {
readonly tpm_limit?: number | null;
}
const numberOrNull = (value: string | number | null | undefined): number | null => {
if (value === null || value === undefined || String(value).trim() === "") return null;
const parsed = Number(value);
return Number.isNaN(parsed) ? null : parsed;
};
export const teamAdminSettingsChanges = (
values: TeamAdminSettingsValues,
initialValues: TeamAdminSettingsValues,
editableFields: ReadonlySet<string>,
): TeamAdminSettingsChanges => {
const tpmLimit = numberOrNull(values.tpm_limit);
return editableFields.has("tpm_limit") && tpmLimit !== numberOrNull(initialValues.tpm_limit)
? { tpm_limit: tpmLimit }
: {};
};
export const parseTeamEditAccess = (callerEditAccess: unknown): TeamEditAccess => {
const parsed = callerEditAccessSchema.safeParse(callerEditAccess);
if (!parsed.success) return { kind: "none" };
if (parsed.data.kind === "team_admin") {
return { kind: "team_admin", editableFields: new Set(parsed.data.editable_fields) };
}
return parsed.data;
};