diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 26e4e06a796..92dc89eb0b8 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -99,7 +99,7 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44358 + "limit": 44802 }, "reportUnknownLambdaType": { "limit": 109 diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index e5e54a3df2c..74cedb9d84d 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.71" +version = "0.1.72" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.71" +version = "0.1.72" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260921190000_agent_identity/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260921190000_agent_identity/migration.sql new file mode 100644 index 00000000000..06cf03b26b5 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260921190000_agent_identity/migration.sql @@ -0,0 +1,97 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "enabled" BOOLEAN NOT NULL DEFAULT true, +ADD COLUMN IF NOT EXISTS "execution_mode" TEXT NOT NULL DEFAULT 'autonomous', +ADD COLUMN IF NOT EXISTS "identity_managed" BOOLEAN NOT NULL DEFAULT false; + +-- AlterTable +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "billing_agent_id" TEXT; + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_AgentIdentity" ( + "agent_id" TEXT NOT NULL, + "active" BOOLEAN NOT NULL DEFAULT true, + "provider" TEXT NOT NULL, + "issuer" TEXT NOT NULL, + "tenant_id" TEXT NOT NULL, + "client_id" TEXT NOT NULL, + "service_principal_id" TEXT, + "required_roles" TEXT[] DEFAULT ARRAY[]::TEXT[], + "required_scopes" TEXT[] DEFAULT ARRAY['user_impersonation']::TEXT[], + "revision" TEXT NOT NULL, + "last_authenticated_at" TIMESTAMP(3), + + CONSTRAINT "LiteLLM_AgentIdentity_pkey" PRIMARY KEY ("agent_id") +); + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_RetiredAgentIdentity" ( + "binding_id" TEXT NOT NULL, + "agent_id" TEXT, + "provider" TEXT NOT NULL, + "issuer" TEXT NOT NULL, + "tenant_id" TEXT NOT NULL, + "client_id" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_RetiredAgentIdentity_pkey" PRIMARY KEY ("binding_id") +); + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_RetiredAgent" ( + "original_agent_id" TEXT NOT NULL, + "retired_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_RetiredAgent_pkey" PRIMARY KEY ("original_agent_id") +); + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_VerifiedSubject" ( + "subject_id" TEXT NOT NULL, + "issuer" TEXT NOT NULL, + "tenant_id" TEXT NOT NULL, + "oid" TEXT NOT NULL, + "kind" TEXT NOT NULL DEFAULT 'human', + "user_id" TEXT, + "verified_via" TEXT NOT NULL DEFAULT 'sso_interactive', + "verified_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_VerifiedSubject_pkey" PRIMARY KEY ("subject_id") +); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentIdentity_provider_tenant_id_client_id_key" ON "LiteLLM_AgentIdentity"("provider", "tenant_id", "client_id"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentIdentity_issuer_service_principal_id_key" ON "LiteLLM_AgentIdentity"("issuer", "service_principal_id"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_RetiredAgentIdentity_provider_tenant_id_client_id_key" ON "LiteLLM_RetiredAgentIdentity"("provider", "tenant_id", "client_id"); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_VerifiedSubject_user_id_idx" ON "LiteLLM_VerifiedSubject"("user_id"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_VerifiedSubject_issuer_tenant_id_oid_key" ON "LiteLLM_VerifiedSubject"("issuer", "tenant_id", "oid"); + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_AgentIdentity_agent_id_fkey') THEN + ALTER TABLE "LiteLLM_AgentIdentity" ADD CONSTRAINT "LiteLLM_AgentIdentity_agent_id_fkey" FOREIGN KEY ("agent_id") REFERENCES "LiteLLM_AgentsTable"("agent_id") ON DELETE CASCADE ON UPDATE CASCADE; + END IF; +END $$; + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_RetiredAgentIdentity_agent_id_fkey') THEN + ALTER TABLE "LiteLLM_RetiredAgentIdentity" ADD CONSTRAINT "LiteLLM_RetiredAgentIdentity_agent_id_fkey" FOREIGN KEY ("agent_id") REFERENCES "LiteLLM_AgentsTable"("agent_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerifiedSubject_user_id_fkey') THEN + ALTER TABLE "LiteLLM_VerifiedSubject" ADD CONSTRAINT "LiteLLM_VerifiedSubject_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE CASCADE ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 03e59257f76..f29caa9ceb7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -78,6 +78,11 @@ model LiteLLM_AgentsTable { object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) + identity_managed Boolean @default(false) + enabled Boolean @default(true) + execution_mode String @default("autonomous") + identity LiteLLM_AgentIdentity? + retired_identities LiteLLM_RetiredAgentIdentity[] tpm_limit Int? rpm_limit Int? session_tpm_limit Int? @@ -88,6 +93,56 @@ model LiteLLM_AgentsTable { updated_by String } +model LiteLLM_AgentIdentity { + agent_id String @id + active Boolean @default(true) + agent LiteLLM_AgentsTable @relation(fields: [agent_id], references: [agent_id], onDelete: Cascade) + provider String + issuer String + tenant_id String + client_id String + service_principal_id String? + required_roles String[] @default([]) + required_scopes String[] @default(["user_impersonation"]) + revision String @default(uuid()) + last_authenticated_at DateTime? + @@unique([provider, tenant_id, client_id]) + @@unique([issuer, service_principal_id]) +} + +model LiteLLM_RetiredAgentIdentity { + binding_id String @id @default(uuid()) + agent_id String? + agent LiteLLM_AgentsTable? @relation(fields: [agent_id], references: [agent_id], onDelete: SetNull) + provider String + issuer String + tenant_id String + client_id String + @@unique([provider, tenant_id, client_id]) +} + +model LiteLLM_RetiredAgent { + original_agent_id String @id + retired_at DateTime @default(now()) +} + +model LiteLLM_VerifiedSubject { + subject_id String @id @default(uuid()) + issuer String + tenant_id String + oid String + kind String @default("human") + user_id String? + user LiteLLM_UserTable? @relation(fields: [user_id], references: [user_id], onDelete: Cascade) + verified_via String @default("sso_interactive") + verified_at DateTime @default(now()) + @@unique([issuer, tenant_id, oid]) + @@index([user_id]) +} + + + + model LiteLLM_OrganizationTable { organization_id String @id @default(uuid()) organization_alias String @@ -241,6 +296,7 @@ model LiteLLM_DeletedTeamTable { // Track spend, rate limit, budget Users model LiteLLM_UserTable { + verified_subjects LiteLLM_VerifiedSubject[] user_id String @id user_alias String? team_id String? @@ -675,6 +731,7 @@ model LiteLLM_SpendLogs { session_id String? status String? mcp_namespaced_tool_name String? + billing_agent_id String? agent_id String? proxy_server_request Json? @default("{}") litellm_call_id String? diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 2835715ef30..e92af4b0861 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.102" +version = "0.4.103" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.102" +version = "0.4.103" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm-rust/crates/cost/tests/calculation.rs b/litellm-rust/crates/cost/tests/calculation.rs index 2acd12b647b..f8152483a7c 100644 --- a/litellm-rust/crates/cost/tests/calculation.rs +++ b/litellm-rust/crates/cost/tests/calculation.rs @@ -203,6 +203,85 @@ fn threshold_tiers_and_boundaries() { assert_eq!(calculate(&specification, &flex).unwrap().input(), 600.0); } +#[rstest] +#[case::ultrafast_above_threshold(ServiceTier::Ultrafast, 300_000, 9_301_000.0, 37_000.0)] +#[case::ultrafast_at_threshold(ServiceTier::Ultrafast, 272_000, 544_500.0, 5_000.0)] +#[case::standard_above_threshold(ServiceTier::Standard, 300_000, 3_300_600.0, 13_000.0)] +#[case::priority_above_threshold(ServiceTier::Priority, 300_000, 5_701_000.0, 23_000.0)] +fn tiered_long_context_rates_are_selected_by_service_tier( + #[case] service_tier: ServiceTier, + #[case] prompt_tokens: u64, + #[case] expected_input: f64, + #[case] expected_output: f64, +) { + let standard = Rates { + cache_read: Rate::Value(3.0), + ..rates(Rate::Value(1.0), Rate::Value(2.0)) + }; + let tiers = [ + TierRates { + tier: ServiceTier::Priority, + rates: Rates { + cache_read: Rate::Value(5.0), + ..rates(Rate::Value(3.0), Rate::Value(4.0)) + }, + }, + TierRates { + tier: ServiceTier::Ultrafast, + rates: Rates { + cache_read: Rate::Value(7.0), + ..rates(Rate::Value(2.0), Rate::Value(5.0)) + }, + }, + ]; + let threshold_tiers = [ + TierRates { + tier: ServiceTier::Priority, + rates: Rates { + cache_read: Rate::Value(29.0), + ..rates(Rate::Value(19.0), Rate::Value(23.0)) + }, + }, + TierRates { + tier: ServiceTier::Ultrafast, + rates: Rates { + cache_read: Rate::Value(41.0), + ..rates(Rate::Value(31.0), Rate::Value(37.0)) + }, + }, + ]; + let thresholds = [ThresholdRates { + above_prompt_tokens: 272_000, + standard: Rates { + cache_read: Rate::Value(17.0), + ..rates(Rate::Value(11.0), Rate::Value(13.0)) + }, + tiers: &threshold_tiers, + }]; + let pricing = Pricing { + standard, + tiers: &tiers, + thresholds: &thresholds, + off_peak: None, + }; + let base = request(); + let long_context_request = Request { + usage: Usage { + prompt_tokens, + completion_tokens: 1_000, + cache_read_tokens: 100, + cache_write_tokens: 0, + ..base.usage + }, + service_tier, + ..base + }; + let cost = calculate(&pricing, &long_context_request).unwrap(); + + assert_eq!(cost.input(), expected_input); + assert_eq!(cost.output(), expected_output); +} + #[test] fn compile_rejects_ambiguous_rates() { let duplicate = ThresholdRates { diff --git a/litellm-rust/crates/python-bridge/src/cache/native/activation.rs b/litellm-rust/crates/python-bridge/src/cache/native/activation.rs index 20f8179ffb0..3ac038d2c39 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/activation.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/activation.rs @@ -1,5 +1,5 @@ use crate::cache::cache_error; -use crate::logger::run_sync_value; +use crate::execution::run_sync_value; use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; use litellm_cache_redis_semantic::RedisSemanticConfig; use litellm_host_python::release_gil; diff --git a/litellm-rust/crates/python-bridge/src/cache/native/backend.rs b/litellm-rust/crates/python-bridge/src/cache/native/backend.rs index 03897d0ddf1..1151ed5cc9d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/backend.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/backend.rs @@ -470,7 +470,7 @@ impl NativeResponseCache { match self { Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); - crate::logger::run_async( + crate::execution::run_async( py, async move { service @@ -495,7 +495,7 @@ impl NativeResponseCache { match self { Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); - crate::logger::run_async( + crate::execution::run_async( py, async move { service.async_lookup(&request, now()).await }, cache_error, @@ -550,7 +550,7 @@ impl NativeResponseCache { match self { Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); - crate::logger::run_async( + crate::execution::run_async( py, async move { service.async_store(&request, response, now()).await }, cache_error, @@ -619,7 +619,7 @@ impl NativeResponseCache { match self { Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); - crate::logger::run_async( + crate::execution::run_async( py, async move { service.async_store_batch(entries, now()).await }, cache_error, diff --git a/litellm-rust/crates/python-bridge/src/cache/native/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/native/semantic.rs index b1c43e1f602..8d9bf270be0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/semantic.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/semantic.rs @@ -1,5 +1,5 @@ use crate::cache::cache_error; -use crate::logger::run_async; +use crate::execution::run_async; use std::{collections::VecDeque, time::Duration}; use litellm_cache::Error; diff --git a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs index 654de75d6bb..0dd70a042e9 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs @@ -144,7 +144,7 @@ impl NativeCacheHandle { self.check_process()?; let request = request(key, None)?; let backend = self.backend.clone(); - crate::logger::run_async( + crate::execution::run_async( py, async move { backend.async_lookup(&request, super::request::now()).await }, cache_error, @@ -163,7 +163,7 @@ impl NativeCacheHandle { let request = request(key, ttl)?; let value: Value = from_py(value)?; let backend = self.backend.clone(); - crate::logger::run_async( + crate::execution::run_async( py, async move { backend @@ -188,7 +188,7 @@ impl NativeCacheHandle { .map(|(key, value)| Ok((request(key, ttl)?, value))) .collect::>>()?; let backend = self.backend.clone(); - crate::logger::run_async( + crate::execution::run_async( py, async move { backend @@ -202,19 +202,19 @@ impl NativeCacheHandle { fn flush(&self, py: Python<'_>) -> PyResult> { self.check_process()?; let backend = self.backend.clone(); - crate::logger::run_sync(py, async move { backend.async_flush().await }, cache_error) + crate::execution::run_sync(py, async move { backend.async_flush().await }, cache_error) } fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { self.check_process()?; let backend = self.backend.clone(); - crate::logger::run_async(py, async move { backend.async_flush().await }, cache_error) + crate::execution::run_async(py, async move { backend.async_flush().await }, cache_error) } fn ping<'py>(&self, py: Python<'py>) -> PyResult> { self.check_process()?; let storage = self.storage.clone(); - crate::logger::run_async( + crate::execution::run_async( py, async move { match storage { @@ -229,7 +229,7 @@ impl NativeCacheHandle { fn disconnect<'py>(&self, py: Python<'py>) -> PyResult> { self.check_process()?; let storage = self.storage.clone(); - crate::logger::run_async( + crate::execution::run_async( py, async move { match storage { @@ -244,7 +244,7 @@ impl NativeCacheHandle { fn delete<'py>(&self, py: Python<'py>, keys: Vec) -> PyResult> { self.check_process()?; let storage = self.storage.clone(); - crate::logger::run_async( + crate::execution::run_async( py, async move { for key in keys { diff --git a/litellm-rust/crates/python-bridge/src/cache/runtime.rs b/litellm-rust/crates/python-bridge/src/cache/runtime.rs index 3bcefad1f1c..eec82f2ba4c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/cache/runtime.rs @@ -1,4 +1,4 @@ -use crate::logger::run_async; +use crate::execution::run_async; use litellm_cache_response::PartialHits; use litellm_host_python::{ExecutionStep, from_py, release_gil, to_py}; use pyo3::{ diff --git a/litellm-rust/crates/python-bridge/src/logger/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs similarity index 71% rename from litellm-rust/crates/python-bridge/src/logger/execution.rs rename to litellm-rust/crates/python-bridge/src/execution.rs index c8d5c0023e3..48f25791b55 100644 --- a/litellm-rust/crates/python-bridge/src/logger/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -13,7 +13,7 @@ where E: Send + 'static, F: Future> + Send + 'static, { - litellm_host_python::run_sync(py, super::capture(py).instrument(future), map_error) + litellm_host_python::run_sync(py, crate::logger::capture(py).instrument(future), map_error) } pub(crate) fn run_async( @@ -26,7 +26,7 @@ where E: Send + 'static, F: Future> + Send + 'static, { - litellm_host_python::run_async(py, super::capture(py).instrument(future), map_error) + litellm_host_python::run_async(py, crate::logger::capture(py).instrument(future), map_error) } pub(crate) fn run_sync_value(py: Python<'_>, future: F) -> PyResult @@ -34,7 +34,7 @@ where T: Send + 'static, F: Future> + Send + 'static, { - litellm_host_python::run_sync_value(py, super::capture(py).instrument(future)) + litellm_host_python::run_sync_value(py, crate::logger::capture(py).instrument(future)) } pub(crate) fn run_async_value(py: Python<'_>, future: F) -> PyResult> @@ -42,5 +42,5 @@ where T: for<'py> IntoPyObject<'py> + Send + 'static, F: Future> + Send + 'static, { - litellm_host_python::run_async_value(py, super::capture(py).instrument(future)) + litellm_host_python::run_async_value(py, crate::logger::capture(py).instrument(future)) } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index c8b0fd2f8bc..37c21cec2de 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -4,6 +4,7 @@ mod coercion; mod credentials; mod diagnostics; mod errors; +mod execution; mod http; mod lifecycle; mod logger; diff --git a/litellm-rust/crates/python-bridge/src/logger/mod.rs b/litellm-rust/crates/python-bridge/src/logger/mod.rs index 6421fe1d554..bf5c735360b 100644 --- a/litellm-rust/crates/python-bridge/src/logger/mod.rs +++ b/litellm-rust/crates/python-bridge/src/logger/mod.rs @@ -1,7 +1,5 @@ -mod execution; mod machine; -pub(crate) use execution::{run_async, run_async_value, run_sync, run_sync_value}; pub(crate) use machine::LoggedMachine; use litellm_host_python::Pythonized; diff --git a/litellm-rust/crates/python-bridge/src/logger/tests.rs b/litellm-rust/crates/python-bridge/src/logger/tests.rs index 21ebf432c99..1fca3be720e 100644 --- a/litellm-rust/crates/python-bridge/src/logger/tests.rs +++ b/litellm-rust/crates/python-bridge/src/logger/tests.rs @@ -76,7 +76,7 @@ async fn traced_operation(_secret: &str) -> PyResult<()> { #[pyfunction] fn span_warning(py: Python<'_>) -> PyResult> { - super::run_async_value(py, traced_operation("private-key-sentinel")) + crate::execution::run_async_value(py, traced_operation("private-key-sentinel")) } #[pyfunction] @@ -93,7 +93,7 @@ fn levels(py: Python<'_>) { #[pyfunction] fn asynchronous_warning(py: Python<'_>) -> PyResult> { - super::run_async_value(py, async { + crate::execution::run_async_value(py, async { tokio::task::yield_now().await; litellm_tracing::warn!("async warning"); Ok(()) @@ -102,7 +102,7 @@ fn asynchronous_warning(py: Python<'_>) -> PyResult> { #[pyfunction] fn synchronous_warning(py: Python<'_>) -> PyResult<()> { - super::run_sync_value(py, async { + crate::execution::run_sync_value(py, async { tokio::task::yield_now().await; litellm_tracing::warn!("sync warning"); Ok(()) @@ -111,7 +111,7 @@ fn synchronous_warning(py: Python<'_>) -> PyResult<()> { #[pyfunction] fn synchronous_failure(py: Python<'_>) -> PyResult<()> { - super::run_sync_value(py, async { + crate::execution::run_sync_value(py, async { litellm_tracing::warn!("failure diagnostic"); Err(pyo3::exceptions::PyValueError::new_err("request failed")) }) diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index 32369890dea..8d434dbbc74 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -1,4 +1,4 @@ -use crate::logger::{run_async, run_sync}; +use crate::execution::{run_async, run_sync}; use litellm_core::audio_transcription::{ AudioTranscriptionRoute, Error, types::AudioTranscriptionRequest, }; diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index bf1d1645c0a..5955729d6e9 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -2,7 +2,7 @@ mod host; use pyo3::types::{PyDict, PyTuple}; -use crate::logger::{run_async, run_sync}; +use crate::execution::{run_async, run_sync}; use litellm_core::chat_completions::{ChatCompletionsRoute, Error, types::ChatCompletionsRequest}; use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse; use pyo3::prelude::*; diff --git a/litellm-rust/crates/python-bridge/src/routes/responses.rs b/litellm-rust/crates/python-bridge/src/routes/responses.rs index bcddaa2bf8d..4e1426cd298 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses.rs @@ -142,7 +142,7 @@ impl ResponsesWebSocketConnection { ) -> PyResult> { let headers = marshal_headers(headers)?; let timeout = optional_timeout(timeout_seconds); - crate::logger::run_async_value(py, async move { + crate::execution::run_async_value(py, async move { let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout) .await .map_err(route_error_to_pyerr)?; @@ -152,21 +152,21 @@ impl ResponsesWebSocketConnection { fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult> { let inner = self.inner.clone(); - crate::logger::run_async_value(py, async move { + crate::execution::run_async_value(py, async move { inner.send_text(text).await.map_err(route_error_to_pyerr) }) } fn recv_text<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); - crate::logger::run_async_value(py, async move { + crate::execution::run_async_value(py, async move { inner.recv_text().await.map_err(route_error_to_pyerr) }) } fn close<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); - crate::logger::run_async_value(py, async move { + crate::execution::run_async_value(py, async move { inner.close().await.map_err(route_error_to_pyerr) }) } diff --git a/litellm-rust/crates/python-bridge/src/routes/token_counter.rs b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs index 21589aa3fe9..2c26311231f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs @@ -1,4 +1,4 @@ -use crate::logger::run_async; +use crate::execution::run_async; use std::sync::Arc; use std::{num::NonZero, thread::available_parallelism}; diff --git a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs index 4d2e88115c8..ab8fea3697d 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs @@ -1,7 +1,7 @@ use std::{collections::BTreeMap, sync::Arc}; use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; -use litellm_host_python::{from_py, json_object_field, run_async_value, run_sync_value, to_py}; +use litellm_host_python::{from_py, json_object_field, to_py}; use litellm_secrets::{ KeyManagementSettings, KeyManagementSystem, Secret, SecretManager, load_native_manager, read_secret_from_python_manager, @@ -13,6 +13,8 @@ use pyo3::{ types::PyDict, }; +use crate::execution::{run_async_value, run_sync_value}; + #[derive(Clone, PartialEq)] struct Configuration { system: KeyManagementSystem, diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index 3a938007633..0e0eeff83fe 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -33,7 +33,8 @@ "thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01", "token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19", "web-fetch-2025-09-10": "web-fetch-2025-09-10", - "web-search-2025-03-05": "web-search-2025-03-05" + "web-search-2025-03-05": "web-search-2025-03-05", + "mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01" }, "azure_ai": { "advisor-tool-2026-03-01": null, @@ -134,7 +135,8 @@ "token-efficient-tools-2025-02-19": null, "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19", "web-fetch-2025-09-10": null, - "web-search-2025-03-05": null + "web-search-2025-03-05": null, + "mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01" }, "bedrock_mantle": { "advanced-tool-use-2025-11-20": "tool-search-tool-2025-10-19", diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 996273d558a..47ce1d35895 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -24,7 +24,7 @@ from litellm.types.caching import RedisPipelineIncrementOperation from .base_cache import BaseCache from .in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache -from .redis_batch import BatchResult, RedisBatch, active_post_call_redis_batch +from .redis_batch import BatchResult, RedisBatch, active_post_call_redis_batch, active_request_redis_batch from .redis_cache import RedisCache, RedisCircuitBreakerOpenError, log_redis_failure if TYPE_CHECKING: @@ -279,6 +279,9 @@ class DualCache(BaseCache): result = in_memory_result if result is None and self.redis_cache is not None and local_only is False: + request_batch: Final = active_request_redis_batch(self.redis_cache) + if request_batch is not None and request_batch.read_as_missing(key): + return None # If not found in in-memory cache, try fetching from Redis redis_result: Final = await self.redis_cache.async_get_cache(key, parent_otel_span=parent_otel_span) @@ -323,6 +326,20 @@ class DualCache(BaseCache): return sublist_keys, previous_access_times + def reserve_redis_batch_reads(self, keys: Sequence[str]) -> tuple[list[str], dict[str, float | None]]: + """Reserve the memory-missed keys whose throttled Redis reads are due, as a batch read would.""" + if self.redis_cache is None: + return [], {} # mutable-ok: API contract returns an empty list and dictionary + key_list: Final = list(keys) # mutable-ok: batch_get_cache takes a list + memory: Final = self.in_memory_cache + in_memory_result: Final = ( + None + if memory is None # pyright: ignore[reportUnnecessaryComparison] # handle an absent in-memory tier + else memory.batch_get_cache(key_list) + ) + result: Final = in_memory_result if in_memory_result is not None else tuple(None for _ in key_list) + return self._reserve_redis_batch_keys(time.time(), key_list, result) + def _rollback_redis_batch_key_reservations(self, previous_access_times: dict[str, float | None]) -> None: with self._last_redis_batch_access_time_lock: for key, previous_time in previous_access_times.items(): @@ -502,12 +519,29 @@ class DualCache(BaseCache): verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async add_cache", e, with_traceback=True ) + async def async_set_cache_pre_call(self, key: str, value: object, ttl: float | None) -> BatchResult[None] | None: + """Memory now, the Redis SET on the request's pipeline, sent with the next read any caller awaits; None + when no pipeline is open, so the caller takes its direct path.""" + batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache) + return None if batch is None else await self._set_on_batch(batch, key, value, ttl) + async def async_set_cache_post_call(self, key: str, value: object, ttl: float | None) -> BatchResult[None] | None: """Memory now, the Redis SET on the request's post-call pipeline; None when no pipeline is open, so the caller takes its direct path.""" batch: Final = None if self.redis_cache is None else active_post_call_redis_batch(self.redis_cache) + return None if batch is None else await self._set_on_batch(batch, key, value, ttl) + + async def async_delete_cache_pre_call(self, key: str) -> BatchResult[None] | None: + """Memory now, the Redis DEL on the request's pipeline; None when no pipeline is open, so the caller + takes its direct path.""" + batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache) if batch is None: return None + if self.in_memory_cache is not None: + self.in_memory_cache.delete_cache(key) + return batch.delete(key) + + async def _set_on_batch(self, batch: RedisBatch, key: str, value: object, ttl: float | None) -> BatchResult[None]: effective_ttl: Final = self.default_in_memory_ttl if ttl is None else ttl if self.in_memory_cache is not None: await self.in_memory_cache.async_set_cache(key, value, ttl=effective_ttl) diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index f6052192685..d408aac8cda 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -47,6 +47,7 @@ class _RedisPipeline(Protocol): def incrbyfloat(self, name: str, amount: float) -> object: ... def expire(self, name: str, time: timedelta) -> object: ... def set(self, name: str, value: str, ex: timedelta | None = None) -> object: ... + def delete(self, *names: str) -> object: ... async def execute(self, raise_on_error: bool = True) -> list[object]: ... @@ -233,6 +234,27 @@ class _Set(_Op[None]): await self._redis_cache.async_set_cache_pipeline_with_ttls(((self._key, self._value, self._ttl),)) +class _Delete(_Op[None]): + """DEL of one key, the pipelined twin of ``async_delete_cache``.""" + + __slots__ = ("_key", "_redis_cache") + + def __init__(self, redis_cache: RedisCache, key: str) -> None: + super().__init__() + self._redis_cache: Final = redis_cache + self._key: Final = key + + def enqueue(self, pipe: _RedisPipeline) -> int: + pipe.delete(self._redis_cache.check_and_fix_namespace(key=self._key)) + return 1 + + def resolve(self, replies: Sequence[object]) -> None: + return None + + async def run_alone(self) -> None: + await self._redis_cache.async_delete_cache(self._key) + + class BatchResult(Generic[_T]): """Awaitable handle for one declared operation; awaiting it flushes the batch it belongs to.""" @@ -269,10 +291,23 @@ class RedisBatch: _pending: list[_Op[object]] = field(default_factory=list) # mutable-ok: drained by flush _flush_hooks: list[Callable[[], None]] = field(default_factory=list) # mutable-ok: append-only registry _lock: asyncio.Lock = field(default_factory=asyncio.Lock) + _misses: set[str] = field(default_factory=set) # mutable-ok: keys an MGET of this request read as absent flushes: int = 0 def mget(self, keys: Sequence[str]) -> BatchResult[Mapping[str, object]]: - return self._declare(_MGet(self.redis_cache, keys)) + op: Final = _MGet(self.redis_cache, keys) + op.future.add_done_callback(self._note_misses) + return self._declare(op) + + def _note_misses(self, future: asyncio.Future[Mapping[str, object]]) -> None: + if future.cancelled() or future.exception() is not None: + return + self._misses.update(key for key, value in future.result().items() if value is None) + + def read_as_missing(self, key: str) -> bool: + """True when an MGET on this batch already found no value under ``key`` and nothing has set it since, + so a per-key GET later in the same request can be answered without another round trip.""" + return key in self._misses def script( self, source: str, run: RegisteredScript, keys: Sequence[str], args: Sequence[_ScriptArg] @@ -283,8 +318,13 @@ class RedisBatch: return self._declare(_Increment(self.redis_cache, key, value, ttl)) def set(self, key: str, value: object, ttl: float | None = None) -> BatchResult[None]: + self._misses.discard(key) return self._declare(_Set(self.redis_cache, key, value, ttl)) + def delete(self, key: str) -> BatchResult[None]: + self._misses.add(key) + return self._declare(_Delete(self.redis_cache, key)) + def add_flush_hook(self, hook: Callable[[], None]) -> None: """Called at the start of every flush so lazily bound readers can declare their keys into the same trip.""" self._flush_hooks.append(hook) @@ -449,10 +489,6 @@ class RequestRedisBatches: def batches(self) -> tuple[RedisBatch, ...]: return tuple(self._batches.values()) - @property - def post_call_batches(self) -> tuple[RedisBatch, ...]: - return tuple(self._post_call.values()) - _active_request_batches: Final[ContextVar[RequestRedisBatches | None]] = ContextVar( "request_redis_batches", default=None diff --git a/litellm/constants.py b/litellm/constants.py index 39c10d71709..530d678457d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1908,6 +1908,7 @@ SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600 SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30 SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000 SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000 +SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE: Final = 100 # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding # staleness so a transient DB error (which surfaces as an empty list) cannot diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 6b3e739ac4f..238b7cc3fdd 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2008,7 +2008,6 @@ def response_cost_calculator( else: if isinstance(response_object, BaseModel): if hasattr(response_object, "_hidden_params"): - response_object._hidden_params["optional_params"] = optional_params provider_response_cost: Final = get_response_cost_from_hidden_params(response_object._hidden_params) if provider_response_cost is not None: return provider_response_cost diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index e555d7e8ec0..3b6827375f3 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -2015,7 +2015,11 @@ def is_unsignable_thinking_block(block: object) -> bool: return not (isinstance(thinking_text, str) and len(thinking_text.strip()) > 0) -def strip_encrypted_reasoning_from_messages(messages: object) -> None: +def strip_encrypted_reasoning_from_messages( + messages: object, + *, + should_strip: Callable[[Mapping[str, object]], bool] | None = None, +) -> None: """Drop the bridge-tagged reasoning blocks a routed deployment cannot decrypt from Anthropic-shaped history. @@ -2030,7 +2034,7 @@ def strip_encrypted_reasoning_from_messages(messages: object) -> None: if not isinstance(messages, list): return for content in anthropic_content_lists(cast(list[object], messages)): # cast-ok: untyped client json - _strip_encrypted_reasoning_from_blocks(content) + _strip_encrypted_reasoning_from_blocks(content, should_strip=should_strip) def anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: @@ -2043,9 +2047,18 @@ def anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: ) -def _strip_encrypted_reasoning_from_blocks(content: object) -> None: +def _strip_encrypted_reasoning_from_blocks( + content: object, + *, + should_strip: Callable[[Mapping[str, object]], bool] | None = None, +) -> None: blocks: Final = cast(list[object], content) # cast-ok: narrowed by the caller's isinstance - kept: Final = tuple(block for block in blocks if not is_encrypted_reasoning_block(block)) + kept: Final = tuple( + block + for block in blocks + if not is_encrypted_reasoning_block(block) + or (should_strip is not None and not should_strip(cast(Mapping[str, object], block))) + ) blocks[:] = kept diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index bf2d588dd3a..3e61a0caa90 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -31,6 +31,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.anthropic import ( ANTHROPIC_HOSTED_TOOLS, + ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER, ANTHROPIC_OAUTH_BETA_HEADER, ANTHROPIC_OAUTH_TOKEN_PREFIX, AllAnthropicToolsValues, @@ -326,6 +327,12 @@ class AnthropicModelInfo(BaseLLMModelInfo): file_ids: Final = get_file_ids_from_messages(messages) return len(file_ids) > 0 + def is_mid_conversation_output_config_used(self, messages: list[AllMessageValues]) -> bool: + """ + Return if "output_config" is in a message + """ + return any("output_config" in message for message in messages) + def is_mcp_server_used(self, mcp_servers: list[AnthropicMcpServerTool] | None) -> bool: if mcp_servers is None: return False @@ -851,6 +858,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): mcp_server_used: bool = False, *, custom_llm_provider: str, + is_mid_conversation_output_config_used: bool = False, ) -> list[str]: """ Get list of common beta headers based on the features that are active. @@ -883,6 +891,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): if mcp_server_used: betas.append("mcp-client-2025-04-04") + if is_mid_conversation_output_config_used: + betas.append(ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER) + return list(set(betas)) @staticmethod @@ -915,6 +926,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): container_with_skills_used: bool = False, api_base: str | None = None, use_bearer_for_custom_base: bool = False, + is_mid_conversation_output_config_used: bool = False, ) -> dict: betas: Final = set() # Anthropic no longer requires the prompt-caching beta header @@ -950,6 +962,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): if container_with_skills_used: betas.add("skills-2025-10-02") + if is_mid_conversation_output_config_used: + betas.add(ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER) + _is_oauth: Final = api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX) headers: Final = { "anthropic-version": anthropic_version or "2023-06-01", @@ -1015,6 +1030,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): mcp_server_used: Final = self.is_mcp_server_used(mcp_servers=optional_params.get("mcp_servers")) pdf_used: Final = self.is_pdf_used(messages=messages) file_id_used: Final = self.is_file_id_used(messages=messages) + is_mid_conversation_output_config_used: Final = self.is_mid_conversation_output_config_used(messages=messages) web_search_tool_used: Final = self.is_web_search_tool_used(tools=tools) tool_search_used: Final = self.is_tool_search_used(tools=tools) programmatic_tool_calling_used: Final = self.is_programmatic_tool_calling_used(tools=tools) @@ -1032,6 +1048,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): api_key=api_key, auth_token=auth_token, file_id_used=file_id_used, + is_mid_conversation_output_config_used=is_mid_conversation_output_config_used, web_search_tool_used=web_search_tool_used, is_vertex_request=optional_params.get("is_vertex_request", False), user_anthropic_beta_headers=user_anthropic_beta_headers, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index c5abb5e9a1c..4d758640d72 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -255,6 +255,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): tool_search_used: Final = self.is_tool_search_used(tools) programmatic_tool_calling_used: Final = self.is_programmatic_tool_calling_used(tools) input_examples_used: Final = self.is_input_examples_used(tools) + is_mid_conversation_output_config_used: Final = self.is_mid_conversation_output_config_used(messages) user_beta_set: Final = set(get_anthropic_beta_from_headers(headers)) beta_set: Final = set(user_beta_set) @@ -266,6 +267,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): file_id_used=self.is_file_id_used(messages), mcp_server_used=self.is_mcp_server_used(optional_params.get("mcp_servers")), custom_llm_provider="bedrock", + is_mid_conversation_output_config_used=is_mid_conversation_output_config_used, ) beta_set.update(auto_betas) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 94eb0c92e40..f018077ebcd 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -515,7 +515,13 @@ class AmazonAnthropicClaudeMessagesConfig( tool_search_used: Final = anthropic_model_info.is_tool_search_used(tools) programmatic_tool_calling_used: Final = anthropic_model_info.is_programmatic_tool_calling_used(tools) input_examples_used: Final = anthropic_model_info.is_input_examples_used(tools) - + outgoing_messages_typed: Final = cast( + list[AllMessageValues], + anthropic_messages_request["messages"], + ) + is_mid_conversation_output_config_used: Final = anthropic_model_info.is_mid_conversation_output_config_used( + outgoing_messages_typed + ) user_beta_set: Final = set(get_anthropic_beta_from_headers(headers)) beta_set: Final = set(user_beta_set) auto_betas: Final = anthropic_model_info.get_anthropic_beta_list( @@ -528,6 +534,7 @@ class AmazonAnthropicClaudeMessagesConfig( anthropic_messages_optional_request_params.get("mcp_servers") ), custom_llm_provider="bedrock", + is_mid_conversation_output_config_used=is_mid_conversation_output_config_used, ) beta_set.update(auto_betas) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9989846b7e0..136399b557e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3358,7 +3358,7 @@ "supports_function_calling": true }, "azure_ai/claude-haiku-4-5": { - "deprecation_date": "2026-10-19", + "deprecation_date": "2026-11-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -3378,10 +3378,11 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule" }, "azure_ai/claude-opus-4-5": { - "deprecation_date": "2026-10-19", + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -3402,7 +3403,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule" }, "azure_ai/claude-opus-4-6": { "deprecation_date": "2027-02-02", @@ -3640,7 +3642,7 @@ "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-4-5": { - "deprecation_date": "2026-10-19", + "deprecation_date": "2026-11-15", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -3660,7 +3662,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule" }, "azure_ai/claude-sonnet-5": { "deprecation_date": "2027-06-30", @@ -30721,6 +30724,7 @@ "output_cost_per_image": 0.08 }, "gemini/veo-3.1-fast-generate-preview": { + "deprecation_date": "2026-10-22", "litellm_provider": "gemini", "max_input_tokens": 1024, "max_tokens": 1024, @@ -30737,6 +30741,7 @@ ] }, "gemini/veo-3.1-generate-preview": { + "deprecation_date": "2026-10-22", "litellm_provider": "gemini", "max_input_tokens": 1024, "max_tokens": 1024, @@ -30752,6 +30757,7 @@ ] }, "gemini/veo-3.1-lite-generate-preview": { + "deprecation_date": "2026-10-22", "litellm_provider": "gemini", "max_input_tokens": 1024, "max_tokens": 1024, @@ -32912,10 +32918,13 @@ "gpt-image-2.5-flare": { "cache_read_input_image_token_cost": 2e-06, "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_batches": 6.25e-07, "input_cost_per_token": 5e-06, "litellm_provider": "openai", "mode": "image_generation", "input_cost_per_image_token": 8e-06, + "input_cost_per_image_token_batches": 4e-06, + "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, "supported_endpoints": [ "/v1/images/generations", @@ -32944,10 +32953,13 @@ "gpt-image-2.5-sunburst": { "cache_read_input_image_token_cost": 2e-06, "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_batches": 6.25e-07, "input_cost_per_token": 5e-06, "litellm_provider": "openai", "mode": "image_generation", "input_cost_per_image_token": 8e-06, + "input_cost_per_image_token_batches": 4e-06, + "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, "supported_endpoints": [ "/v1/images/generations", @@ -38611,6 +38623,7 @@ }, "mistral/zai-glm-5-2": { "cache_read_input_token_cost": 1.4e-07, + "deprecation_date": "2026-10-31", "input_cost_per_token": 1.4e-06, "litellm_provider": "mistral", "max_input_tokens": 1048576, @@ -38741,6 +38754,7 @@ "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-0": { + "deprecation_date": "2026-09-30", "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, "ocr_cost_per_page_batches": 0.002, @@ -60616,6 +60630,7 @@ "supports_vision": true }, "mistral/labs-leanstral-1-5": { + "deprecation_date": "2026-09-30", "input_cost_per_token": 0.0, "litellm_provider": "mistral", "max_input_tokens": 262144, @@ -61337,13 +61352,16 @@ }, "fireworks_ai/nemotron-lightning-3p5-30b-a3b": { "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_priority": 1.25e-08, "input_cost_per_token": 5e-08, + "input_cost_per_token_priority": 6.25e-08, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2e-07, + "output_cost_per_token_priority": 2.5e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, @@ -61353,13 +61371,16 @@ }, "fireworks_ai/nemotron-3-ultra-nvfp4": { "cache_read_input_token_cost": 1.2e-07, + "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 6e-07, + "input_cost_per_token_priority": 7.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2.4e-06, + "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, @@ -61389,13 +61410,16 @@ }, "fireworks_ai/accounts/fireworks/models/nemotron-lightning-3p5-30b-a3b": { "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_priority": 1.25e-08, "input_cost_per_token": 5e-08, + "input_cost_per_token_priority": 6.25e-08, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2e-07, + "output_cost_per_token_priority": 2.5e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, @@ -61405,13 +61429,16 @@ }, "fireworks_ai/accounts/fireworks/models/nemotron-3-ultra-nvfp4": { "cache_read_input_token_cost": 1.2e-07, + "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 6e-07, + "input_cost_per_token_priority": 7.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2.4e-06, + "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, @@ -64321,13 +64348,16 @@ }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-us": { "cache_read_input_token_cost": 3.9e-07, + "cache_read_input_token_cost_priority": 4.875e-07, "input_cost_per_token": 2.1e-06, + "input_cost_per_token_priority": 2.625e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.6e-06, + "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, @@ -64356,13 +64386,16 @@ }, "fireworks_ai/glm-5p3-us": { "cache_read_input_token_cost": 3.9e-07, + "cache_read_input_token_cost_priority": 4.875e-07, "input_cost_per_token": 2.1e-06, + "input_cost_per_token_priority": 2.625e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.6e-06, + "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, @@ -64446,12 +64479,15 @@ }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-flash-us": { "cache_read_input_token_cost": 4.5e-08, + "cache_read_input_token_cost_priority": 5.625e-08, "input_cost_per_token": 2.25e-07, + "input_cost_per_token_priority": 2.8125e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 7.5e-07, + "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_response_schema": true, @@ -64477,12 +64513,15 @@ }, "fireworks_ai/glm-5p3-flash-us": { "cache_read_input_token_cost": 4.5e-08, + "cache_read_input_token_cost_priority": 5.625e-08, "input_cost_per_token": 2.25e-07, + "input_cost_per_token_priority": 2.8125e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 7.5e-07, + "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_response_schema": true, @@ -64571,6 +64610,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 6e-08, "output_cost_per_token": 2.5e-07, "litellm_provider": "together_ai", @@ -70451,6 +70491,7 @@ }, "together_ai/nvidia/nemotron-3-ultra-550b-a55b": { "cache_read_input_token_cost": 2e-07, + "deprecation_date": "2026-08-27", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_input_tokens": 512288, @@ -77476,11 +77517,14 @@ }, "fireworks_ai/accounts/fireworks/models/ember-1": { "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_priority": 3.75e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_priority": 3.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_priority": 1.875e-05, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, @@ -79218,12 +79262,12 @@ "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1050000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "source": "https://developers.openai.com/api/docs/pricing", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_endpoints": [ "/v1/responses" ], @@ -79253,12 +79297,12 @@ "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1050000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "source": "https://developers.openai.com/api/docs/pricing", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_modalities": [ "text", "image" @@ -79285,12 +79329,12 @@ "input_cost_per_token_above_272k_tokens": 4.4e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "responses", "output_cost_per_token": 1.1e-05, "output_cost_per_token_above_272k_tokens": 1.65e-05, - "source": "https://developers.openai.com/api/docs/models/gpt-6.1-sol", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -79323,12 +79367,12 @@ "input_cost_per_token_above_272k_tokens": 4.4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1050000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.1e-05, "output_cost_per_token_above_272k_tokens": 1.65e-05, - "source": "https://developers.openai.com/api/docs/pricing", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_endpoints": [ "/v1/responses" ], @@ -79348,5 +79392,33 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true + }, + "vertex_ai/gemini-3.8-flash-tts": { + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "audio_speech", + "output_cost_per_audio_token": 9e-06, + "output_cost_per_token": 9e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "vertex_ai/gemini-3.8-flash-lite-tts": { + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "audio_speech", + "output_cost_per_audio_token": 6e-06, + "output_cost_per_token": 6e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/audio/speech" + ] } } diff --git a/litellm/proxy/_experimental/mcp_server/auth/managed_agent_access.py b/litellm/proxy/_experimental/mcp_server/auth/managed_agent_access.py new file mode 100644 index 00000000000..096b7eb3c77 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/auth/managed_agent_access.py @@ -0,0 +1,74 @@ +from types import MappingProxyType +from typing import Final + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings +from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure +from litellm.types.proxy.agent_identity import AgentIdentityFailure + + +async def _delegated_resource_subject(user_id: str) -> UserAPIKeyAuth: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + + human: Final = await MCPRequestHandler.reload_admitted_user(user_id, requires_fresh_policy=True) + return human.model_copy(update=MappingProxyType({"mcp_explicit_grants_only": True})) + + +async def managed_agent_servers(auth: UserAPIKeyAuth) -> tuple[str, ...]: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + agent: Final = auth.managed_agent_policy + if agent is None: + return () + + try: + base: Final = frozenset(await MCPRequestHandler.get_allowed_mcp_servers_for_agent(auth)) + ceilings: Final = await resolve_managed_agent_ceilings(agent) + expanded: Final = tuple( + frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids))) + for ceiling in ceilings + ) + grouped: Final = frozenset(server for server in base if all(server in ceiling for ceiling in expanded)) + caller_capped, _ = await MCPRequestHandler.apply_agent_caller_ceiling(sorted(grouped), auth) + own: Final = frozenset(caller_capped) + context: Final = auth.managed_agent_context + if context is None or context.mode == "autonomous": + return tuple(sorted(own)) + if context.user_id is None: + return () + human: Final = await _delegated_resource_subject(context.user_id) + allowed: Final = await MCPRequestHandler.resolve_admitted_subject_servers( + human, allowed_team_ids=frozenset((auth.team_id,)) if auth.team_id else frozenset() + ) + return tuple(sorted(own.intersection(allowed))) + except Exception: # noqa: BLE001 # Authorization boundary: every unresolved policy must deny access + raise_identity_failure( + AgentIdentityFailure(code="policy_unavailable", message="Agent MCP policy is unavailable") + ) + + +async def managed_agent_tools(server_id: str, auth: UserAPIKeyAuth) -> list[str] | None: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + + if server_id not in await managed_agent_servers(auth): + return [] + try: + granted: Final = await MCPRequestHandler.get_agent_tool_permissions_for_server(server_id, auth) + own: Final = await MCPRequestHandler.apply_agent_caller_tool_ceiling(granted, server_id, auth) + context: Final = auth.managed_agent_context + if context is None or context.mode == "autonomous": + return None if own is None else sorted(own) + if context.user_id is None: + return [] + human: Final = await _delegated_resource_subject(context.user_id) + human_tools: Final = await MCPRequestHandler.resolve_admitted_subject_tools( + server_id, human, allowed_team_ids=frozenset((auth.team_id,)) if auth.team_id else frozenset() + ) + if own is None: + return human_tools + return sorted(own) if human_tools is None else sorted(frozenset(own).intersection(human_tools)) + except Exception: # noqa: BLE001 # Authorization boundary: every unresolved policy must deny access + raise_identity_failure( + AgentIdentityFailure(code="policy_unavailable", message="Agent tool policy is unavailable") + ) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index a93ffaeac9f..457c9b1680b 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -51,6 +51,7 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( resolve_agent_access_group_ceiling, ) from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth +from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.user_api_key_auth import ( _get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth @@ -67,7 +68,6 @@ from litellm.repositories.table_repositories import ( AgentsRepository, MCPServerRepository, ) -from litellm.repositories.user_repository import UserRepository from litellm.types.mcp_server.mcp_server_manager import MCPServer if TYPE_CHECKING: @@ -1086,7 +1086,7 @@ class MCPRequestHandler: assert_never(identity.subject_type) @staticmethod - async def reload_admitted_user(user_id: str) -> UserAPIKeyAuth: + async def reload_admitted_user(user_id: str, *, requires_fresh_policy: bool = False) -> UserAPIKeyAuth: """Reload the live user an interactively-minted envelope references and admit them as themselves. The user's own object permission and ``org_id`` ride on the returned ``UserAPIKeyAuth``, and the @@ -1111,6 +1111,7 @@ class MCPRequestHandler: prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=False, + check_db_only=requires_fresh_policy, ) # Resolve the user's own MCP object permission (get_user_object does not load it) so the shared # get_allowed_mcp_servers can grant the user their litellm-granted servers. Reuses the same @@ -1119,6 +1120,7 @@ class MCPRequestHandler: if user_object is not None and object_permission is None and user_object.object_permission_id: object_permission = await get_object_permission( object_permission_id=user_object.object_permission_id, + check_db_only=requires_fresh_policy, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, ) @@ -1147,6 +1149,7 @@ class MCPRequestHandler: # Server-only marker, set AFTER construction: the before-validator strips it from any validated # input, so caller-supplied data (key metadata, JWT claims) can never forge it. admitted.mcp_admitted_user_subject = True + admitted.requires_fresh_policy = requires_fresh_policy # Carry each granting team's per-server mcp_rpm_limit: this subject reaches servers through # several teams under its own identity, so without this a cross-team user outruns every team's # limit. Resolved from the same roster-checked sources as the grant union, so a team throttles @@ -1202,7 +1205,7 @@ class MCPRequestHandler: return None @staticmethod - async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth: + async def _reload_admitted_key(key_hash: str, *, check_db_only: bool = False) -> UserAPIKeyAuth: """Reload the live key record an admitted envelope references and re-check live policy. Resolving the current ``UserAPIKeyAuth`` (cache first, then DB) is what stops the @@ -1234,6 +1237,7 @@ class MCPRequestHandler: hashed_token=key_hash, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + check_db_only=check_db_only, ) except (ProxyException, HTTPException): raise HTTPException(status_code=401, detail="Invalid or expired credential") from None @@ -1597,6 +1601,11 @@ class MCPRequestHandler: """ from litellm.proxy.proxy_server import general_settings + if managed_agent_policy(user_api_key_auth) is not None: + from litellm.proxy._experimental.mcp_server.auth.managed_agent_access import managed_agent_servers + + return MCPServerAccess(server_ids=await managed_agent_servers(user_api_key_auth), scope="scoped") + key_object_permission: Final = MCPRequestHandler._get_key_object_permission(user_api_key_auth) try: @@ -1606,7 +1615,7 @@ class MCPRequestHandler: # independent; an opt-out silences only its own source, inside the recursive call). if _is_mcp_admitted_user_subject(user_api_key_auth) and user_api_key_auth is not None: return MCPServerAccess( - server_ids=tuple(await MCPRequestHandler._resolve_admitted_subject_servers(user_api_key_auth)), + server_ids=tuple(await MCPRequestHandler.resolve_admitted_subject_servers(user_api_key_auth)), ) # Get allowed servers from key and team @@ -1703,7 +1712,7 @@ class MCPRequestHandler: if user_api_key_auth and user_api_key_auth.agent_id: agent_capped: Final = _agent_capped_servers( allowed_mcp_servers, - await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth), + await MCPRequestHandler.get_allowed_mcp_servers_for_agent(user_api_key_auth), await MCPRequestHandler._get_agent_access_group_server_ceiling(user_api_key_auth), ) if agent_capped is not None: @@ -1716,7 +1725,7 @@ class MCPRequestHandler: ######################################################### # Cap an agent key at what the user and team that invoked the agent may reach ######################################################### - caller_capped, caller_restricts = await MCPRequestHandler._apply_agent_caller_ceiling( + caller_capped, caller_restricts = await MCPRequestHandler.apply_agent_caller_ceiling( allowed_mcp_servers, user_api_key_auth ) @@ -1829,10 +1838,14 @@ class MCPRequestHandler: scoped.object_permission = auth.object_permission scoped.object_permission_id = auth.object_permission_id scoped.access_group_ids = auth.access_group_ids + scoped.requires_fresh_policy = auth.requires_fresh_policy + scoped.mcp_explicit_grants_only = auth.mcp_explicit_grants_only return scoped @staticmethod - async def _admitted_subject_sources(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + async def admitted_subject_sources( + auth: UserAPIKeyAuth, *, allowed_team_ids: frozenset[str] | None = None + ) -> list[UserAPIKeyAuth]: """The independent sources a keyless admitted subject reaches MCP servers through: their own direct grants, plus every team they are a live roster member of. @@ -1849,6 +1862,8 @@ class MCPRequestHandler: if not auth.user_id or prisma_client is None: return sources for team_id in await MCPRequestHandler._resolve_user_team_ids(auth.user_id, auth): + if allowed_team_ids is not None and team_id not in allowed_team_ids: + continue team_obj = await MCPRequestHandler._roster_team_object(team_id, auth) if team_obj is None: continue @@ -1886,6 +1901,7 @@ class MCPRequestHandler: user_api_key_cache=user_api_key_cache, parent_otel_span=auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(auth and auth.requires_fresh_policy), ) except Exception as e: # noqa: BLE001 # per-source isolation: one team's blip must not deny the others # Fault isolation is per SOURCE: an unresolvable team contributes nothing (fail closed for @@ -1932,7 +1948,9 @@ class MCPRequestHandler: return team_obj @staticmethod - async def admitted_source_grants(auth: UserAPIKeyAuth) -> list[tuple[UserAPIKeyAuth, set[str]]]: + async def admitted_source_grants( + auth: UserAPIKeyAuth, *, allowed_team_ids: frozenset[str] | None = None + ) -> list[tuple[UserAPIKeyAuth, set[str]]]: """``(source, the servers that source grants)`` for every source of an admitted subject. THE owner of "which source reaches which server". The reachable union, the per-team throttle @@ -1941,15 +1959,17 @@ class MCPRequestHandler: roster instead of by grant charged unrelated teams' buckets).""" return [ (source, set(await MCPRequestHandler.get_allowed_mcp_servers(source, keyless_source=True))) - for source in await MCPRequestHandler._admitted_subject_sources(auth) + for source in await MCPRequestHandler.admitted_subject_sources(auth, allowed_team_ids=allowed_team_ids) ] @staticmethod - async def _resolve_admitted_subject_servers(auth: UserAPIKeyAuth) -> list[str]: + async def resolve_admitted_subject_servers( + auth: UserAPIKeyAuth, *, allowed_team_ids: frozenset[str] | None = None + ) -> list[str]: """Union of what each of the admitted subject's sources reaches, each answered by the canonical resolver so no rule is reimplemented for this caller shape.""" reachable: Final[set[str]] = set() - for _source, granted in await MCPRequestHandler.admitted_source_grants(auth): + for _source, granted in await MCPRequestHandler.admitted_source_grants(auth, allowed_team_ids=allowed_team_ids): reachable.update(granted) return list(reachable) @@ -2007,7 +2027,9 @@ class MCPRequestHandler: return min((source for source, _ in granting), key=lambda s: s.team_id or "") @staticmethod - async def _resolve_admitted_subject_tools(server_id: str, auth: UserAPIKeyAuth) -> list[str] | None: + async def resolve_admitted_subject_tools( + server_id: str, auth: UserAPIKeyAuth, *, allowed_team_ids: frozenset[str] | None = None + ) -> list[str] | None: """Effective tool allowlist on ``server_id`` for an admitted subject, as the union over the sources that actually grant that server. @@ -2029,7 +2051,7 @@ class MCPRequestHandler: ) or await MCPRequestHandler.admin_view_unscoped(auth) allowed: Final[set[str]] = set() - for source, granted in await MCPRequestHandler.admitted_source_grants(auth): + for source, granted in await MCPRequestHandler.admitted_source_grants(auth, allowed_team_ids=allowed_team_ids): # The open channel is evaluated against the user's OWN source (team_id is None), so that # source's restrictions apply to it; a team's rules never ride an open-channel server. if server_id not in granted and not (reachable_via_open_channel and source.team_id is None): @@ -2088,6 +2110,7 @@ class MCPRequestHandler: user_api_key_cache=user_api_key_cache, parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) if not team_obj: @@ -2098,6 +2121,8 @@ class MCPRequestHandler: @staticmethod async def _toolset_tool_permissions( object_permission: LiteLLM_ObjectPermissionTable | None, + *, + requires_fresh_policy: bool = False, ) -> Mapping[str, Sequence[str]]: """The ``server_id -> tool names`` grants of this permission row's toolsets, empty when it declares none. The shared resolver for the team, org, and internal-user levels, so a toolset @@ -2114,7 +2139,8 @@ class MCPRequestHandler: if object_permission is None or not object_permission.mcp_toolsets: return _EMPTY_TOOLSET_GRANTS resolved: Final = await global_mcp_server_manager.resolve_toolset_tool_permissions( - toolset_ids=object_permission.mcp_toolsets + toolset_ids=object_permission.mcp_toolsets, + requires_fresh_policy=requires_fresh_policy, ) if not resolved: raise UnloadableEntitlementError( @@ -2126,10 +2152,15 @@ class MCPRequestHandler: async def _toolset_tools_for_server( object_permission: LiteLLM_ObjectPermissionTable | None, server_id: str, + *, + requires_fresh_policy: bool = False, ) -> Sequence[str] | None: """Tool names this row's toolsets grant on ``server_id``, ``None`` when its toolsets place no restriction on that server (it declares no toolsets, or none of them name it).""" - return (await MCPRequestHandler._toolset_tool_permissions(object_permission)).get(server_id) + grants: Final = await MCPRequestHandler._toolset_tool_permissions( + object_permission, requires_fresh_policy=requires_fresh_policy + ) + return grants.get(server_id) @staticmethod def _union_tool_grants( @@ -2171,6 +2202,7 @@ class MCPRequestHandler: user_api_key_cache=user_api_key_cache, parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) @staticmethod @@ -2219,12 +2251,17 @@ class MCPRequestHandler: if not user_api_key_auth: return None + if managed_agent_policy(user_api_key_auth) is not None: + from litellm.proxy._experimental.mcp_server.auth.managed_agent_access import managed_agent_tools + + return await managed_agent_tools(server_id, user_api_key_auth) + try: # FIRST statement, mirroring get_allowed_mcp_servers: a keyless admitted subject resolves per # source and shares nothing with the single-credential prelude below. Ordering is the invariant: # sat after the prelude, a fault in a lookup the subject never uses denied tools its teams grant. if _is_mcp_admitted_user_subject(user_api_key_auth): - return await MCPRequestHandler._resolve_admitted_subject_tools(server_id, user_api_key_auth) + return await MCPRequestHandler.resolve_admitted_subject_tools(server_id, user_api_key_auth) # Get key and team object permissions (already loaded in main auth flow) key_obj_perm: Final = MCPRequestHandler._get_key_object_permission(user_api_key_auth) @@ -2249,9 +2286,12 @@ class MCPRequestHandler: # tool-level check sees the key's full effective tool scope key_toolset_ids: Final = (key_obj_perm.mcp_toolsets or []) if key_obj_perm else [] key_toolset_tools: Final = ( - (await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=key_toolset_ids)).get( - server_id - ) + ( + await global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=key_toolset_ids, + requires_fresh_policy=user_api_key_auth.requires_fresh_policy, + ) + ).get(server_id) if key_toolset_ids else None ) @@ -2265,7 +2305,9 @@ class MCPRequestHandler: # Tools granted through the team's toolsets restrict this server exactly # as the team's direct tool permissions do, mirroring the key path above - team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(team_obj_perm, server_id) + team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server( + team_obj_perm, server_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + ) team_tools: Final = MCPRequestHandler._union_tool_grants(team_direct_tools, team_toolset_tools) # Apply same inheritance logic as get_allowed_mcp_servers @@ -2291,7 +2333,7 @@ class MCPRequestHandler: ) allowed_tools = _as_list( - await MCPRequestHandler._apply_agent_caller_tool_ceiling(allowed_tools, server_id, user_api_key_auth) + await MCPRequestHandler.apply_agent_caller_tool_ceiling(allowed_tools, server_id, user_api_key_auth) ) return await MCPRequestHandler._apply_agent_and_org_tool_ceilings( @@ -2334,7 +2376,7 @@ class MCPRequestHandler: if user_api_key_auth.agent_id: # Pre-fetch agent object_permission once to avoid a duplicate DB query. agent_obj_perm: Final = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) - agent_tools: Final = await MCPRequestHandler._get_agent_tool_permissions_for_server( + agent_tools: Final = await MCPRequestHandler.get_agent_tool_permissions_for_server( server_id=server_id, user_api_key_auth=user_api_key_auth, agent_object_permission=agent_obj_perm, @@ -2365,7 +2407,9 @@ class MCPRequestHandler: if org_obj_perm and org_obj_perm.mcp_tool_permissions else None ) - org_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(org_obj_perm, server_id) + org_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server( + org_obj_perm, server_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + ) org_tools: Final = MCPRequestHandler._union_tool_grants(org_direct_tools, org_toolset_tools) if org_tools is not None: allowed_tools = ( @@ -2456,6 +2500,7 @@ class MCPRequestHandler: prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) if not raw_server_ids: return [] @@ -2502,6 +2547,7 @@ class MCPRequestHandler: user_api_key_cache=user_api_key_cache, parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) if key_object_permission is None: return [] @@ -2518,7 +2564,8 @@ class MCPRequestHandler: # Get MCP servers from access groups access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups( - key_object_permission.mcp_access_groups or [] + key_object_permission.mcp_access_groups or [], + requires_fresh_policy=user_api_key_auth.requires_fresh_policy, ) # servers referenced in tool permissions should also be accessible @@ -2531,7 +2578,14 @@ class MCPRequestHandler: # ceilings as any other key-level grant toolset_ids: Final = key_object_permission.mcp_toolsets or [] toolset_servers: Final = ( - list((await global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=toolset_ids)).keys()) + list( + ( + await global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=toolset_ids, + requires_fresh_policy=user_api_key_auth.requires_fresh_policy, + ) + ).keys() + ) if toolset_ids else [] ) @@ -2550,7 +2604,7 @@ class MCPRequestHandler: """Get allowed MCP servers a caller inherits from the team it is pinned to. Exactly one team, or none. A subject that reaches servers through SEVERAL teams does not - fan out here: it is resolved one source per team in ``_resolve_admitted_subject_servers``, + fan out here: it is resolved one source per team in ``resolve_admitted_subject_servers``, and each of those sources pins a single ``team_id`` before reaching this point. Keeping the fan-out here as well would be a second multi-team path to drift from that one. """ @@ -2568,7 +2622,7 @@ class MCPRequestHandler: which must NOT silently gain the union across every team the user belongs to), and it covers each single-source auth an admitted subject fans out into — those pin a team_id, so they land on the first branch. The admitted subject itself never reaches here: it resolves per source - in ``_resolve_admitted_subject_servers`` before this point. The ``UI_TEAM_ID`` sentinel + in ``resolve_admitted_subject_servers`` before this point. The ``UI_TEAM_ID`` sentinel resolves to no teams exactly as before.""" if user_api_key_auth is None or not user_api_key_auth.team_id: return [] @@ -2596,6 +2650,7 @@ class MCPRequestHandler: user_id_upsert=False, parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) except Exception as e: # noqa: BLE001 # a team-resolution blip narrows access, never raises verbose_logger.warning("Failed to resolve user teams for MCP grant: %s", e) @@ -2605,7 +2660,12 @@ class MCPRequestHandler: return list(dict.fromkeys(t for t in user_object.teams if t and t != UI_TEAM_ID)) @staticmethod - async def _team_granted_servers(team_obj: LiteLLM_TeamTable, team_access_group_servers: list[str]) -> set[str]: + async def _team_granted_servers( + team_obj: LiteLLM_TeamTable, + team_access_group_servers: list[str], + *, + requires_fresh_policy: bool = False, + ) -> set[str]: """The raw MCP-server set a team grants (before any org ceiling): its object_permission (direct ``mcp_servers``, the ``all_proxy_servers`` sentinel → the full registry, legacy access groups, tool-perm-referenced servers, toolset-referenced servers) unioned with its unified @@ -2620,13 +2680,17 @@ class MCPRequestHandler: if SpecialMCPServerName.all_proxy_servers.value in (object_permissions.mcp_servers or []): return set(global_mcp_server_manager.get_registry().keys()) legacy_access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups( - object_permissions.mcp_access_groups or [] + object_permissions.mcp_access_groups or [], + requires_fresh_policy=requires_fresh_policy, + ) + toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions( + object_permissions, requires_fresh_policy=requires_fresh_policy ) return ( set(global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or [])) | set(legacy_access_group_servers) | set(global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys()) - | (await MCPRequestHandler._toolset_tool_permissions(object_permissions)).keys() + | toolset_grants.keys() | set(team_access_group_servers) ) @@ -2667,6 +2731,7 @@ class MCPRequestHandler: user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) if team_obj is None: return [] @@ -2680,12 +2745,19 @@ class MCPRequestHandler: prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) - servers: Final = await MCPRequestHandler._team_granted_servers(team_obj, team_access_group_servers) + servers: Final = await MCPRequestHandler._team_granted_servers( + team_obj, + team_access_group_servers, + requires_fresh_policy=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), + ) return list(servers) except Exception as e: - if isinstance(e, UnloadableEntitlementError): + if isinstance(e, UnloadableEntitlementError) or ( + user_api_key_auth is not None and user_api_key_auth.requires_fresh_policy + ): raise verbose_logger.warning("Failed to get allowed MCP servers for team: %s", e) return [] @@ -2716,6 +2788,7 @@ class MCPRequestHandler: user_api_key_cache=user_api_key_cache, parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) except Exception as e: # noqa: BLE001 # a named entitlement we cannot read denies, whatever the read failed with raise unloadable from e @@ -2811,7 +2884,8 @@ class MCPRequestHandler: direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or []) access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups( - object_permissions.mcp_access_groups or [] + object_permissions.mcp_access_groups or [], + requires_fresh_policy=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) tool_perm_servers: Final = list( @@ -2820,7 +2894,10 @@ class MCPRequestHandler: # servers referenced by the org's toolset grants are part of the org ceiling, # exactly as servers referenced by its inline tool permissions are - toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(object_permissions) + toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions( + object_permissions, + requires_fresh_policy=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), + ) all_servers: Final = tuple( {*direct_mcp_servers, *access_group_servers, *tool_perm_servers, *toolset_grants} @@ -2912,7 +2989,8 @@ class MCPRequestHandler: # Get MCP servers from access groups access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups( - object_permission.mcp_access_groups or [] + object_permission.mcp_access_groups or [], + requires_fresh_policy=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) # servers referenced in tool permissions should also be accessible @@ -2961,7 +3039,9 @@ class MCPRequestHandler: return None user_id: Final = user_api_key_auth.user_id - object_permission_id: Final = await MCPRequestHandler._user_object_permission_id(user_id, prisma_client) + object_permission_id: Final = await MCPRequestHandler._user_object_permission_id( + user_id, prisma_client, check_db_only=user_api_key_auth.requires_fresh_policy + ) if object_permission_id is None: return None @@ -2971,6 +3051,7 @@ class MCPRequestHandler: user_api_key_cache=user_api_key_cache, parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) if object_permission is None: raise ValueError( @@ -2979,7 +3060,9 @@ class MCPRequestHandler: return object_permission @staticmethod - async def _user_object_permission_id(user_id: str, prisma_client: "PrismaClient") -> str | None: + async def _user_object_permission_id( + user_id: str, prisma_client: "PrismaClient", *, check_db_only: bool = False + ) -> str | None: """The permission row this human's user row links to, or None when they link none. Caches the link (with a sentinel for "links none") so a human without an entitlement costs no @@ -2988,16 +3071,23 @@ class MCPRequestHandler: whether someone is entitled is the state that existed before this level, so it places no ceiling. Only a link we DID resolve can make the caller deny. """ + from litellm.proxy.auth.auth_checks import get_user_object from litellm.proxy.proxy_server import user_api_key_cache cache_key: Final = user_object_permission_id_cache_key(user_id) try: - cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) + cached: Final[object] = None if check_db_only else await user_api_key_cache.async_get_cache(key=cache_key) if cached == USER_NO_MCP_PERMISSION_SENTINEL: return None if isinstance(cached, str) and cached: return cached - user_row: Final = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + user_row: Final = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + check_db_only=check_db_only, + ) linked: Final[object] = getattr(user_row, "object_permission_id", None) if user_row is not None else None object_permission_id: Final = linked if isinstance(linked, str) and linked else None await user_api_key_cache.async_set_cache( @@ -3006,7 +3096,9 @@ class MCPRequestHandler: ttl=get_management_object_ttl(user_api_key_cache), ) return object_permission_id - except Exception as e: # noqa: BLE001 # unknown whether entitled at all: no ceiling, as before + except Exception as e: # noqa: BLE001 # Legacy callers retain their existing optional user-ceiling behavior + if check_db_only: + raise HTTPException(503, "User policy is unavailable") from e verbose_logger.warning("MCP user entitlement: link for %r unresolved, no ceiling: %s", user_id, e) return None @@ -3031,13 +3123,17 @@ class MCPRequestHandler: return [] direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or []) + fresh: Final = bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy) access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups( - object_permissions.mcp_access_groups or [] + object_permissions.mcp_access_groups or [], + requires_fresh_policy=fresh, ) tool_perm_servers: Final = list( global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys() ) - toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(object_permissions) + toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions( + object_permissions, requires_fresh_policy=fresh + ) return tuple({*direct_mcp_servers, *access_group_servers, *tool_perm_servers, *toolset_grants}) except Exception as e: # noqa: BLE001 # any resolution fault is an unresolved ceiling, never "no ceiling" verbose_logger.warning("Failed to get allowed MCP servers for user: %s", e) @@ -3075,7 +3171,7 @@ class MCPRequestHandler: return capped, True @staticmethod - async def _apply_agent_caller_ceiling( + async def apply_agent_caller_ceiling( allowed_mcp_servers: Sequence[str], user_api_key_auth: UserAPIKeyAuth | None = None, ) -> tuple[tuple[str, ...], bool]: @@ -3119,9 +3215,13 @@ class MCPRequestHandler: (any non-empty entitlement, or an unresolved one, disqualifies), exactly as ``operator_open_server_ids`` reads the same row. The one owner of this predicate: the server-axis registry resolution in ``get_allowed_mcp_servers`` and the tools-axis open - channel in ``_resolve_admitted_subject_tools`` both consult it, so the two axes cannot + channel in ``resolve_admitted_subject_tools`` both consult it, so the two axes cannot disagree.""" - if user_api_key_auth is None or not user_api_key_has_admin_view(user_api_key_auth): + if ( + user_api_key_auth is None + or user_api_key_auth.mcp_explicit_grants_only + or not user_api_key_has_admin_view(user_api_key_auth) + ): return False object_permission: Final = user_api_key_auth.object_permission credential_scoped: Final = ( @@ -3167,7 +3267,11 @@ class MCPRequestHandler: user_direct_tools: Final = global_mcp_server_manager.expand_tool_permissions( object_permissions.mcp_tool_permissions ).get(server_id) - user_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(object_permissions, server_id) + user_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server( + object_permissions, + server_id, + requires_fresh_policy=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), + ) user_tools: Final = MCPRequestHandler._union_tool_grants(user_direct_tools, user_toolset_tools) if user_tools is None: return allowed_tools @@ -3176,7 +3280,7 @@ class MCPRequestHandler: return list(set(allowed_tools) & set(user_tools)) @staticmethod - async def _apply_agent_caller_tool_ceiling( + async def apply_agent_caller_tool_ceiling( allowed_tools: Sequence[str] | None, server_id: str, user_api_key_auth: UserAPIKeyAuth | None = None, @@ -3184,7 +3288,7 @@ class MCPRequestHandler: """Narrow an agent key's tools on ``server_id`` to those the invoking user and team (echoed back by the agent as ``x-litellm-user-id`` / ``x-litellm-team-id``) may call: the echoed team's tool grants when it names any on this server, then the echoed user's own tool entitlement. The tools - axis twin of ``_apply_agent_caller_ceiling``, so the headers only ever narrow. Denies every tool + axis twin of ``apply_agent_caller_ceiling``, so the headers only ever narrow. Denies every tool on the server when the caller's team cannot be loaded, since a caller we cannot resolve must not read as unrestricted.""" from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -3196,7 +3300,9 @@ class MCPRequestHandler: return allowed_tools try: team_obj_perm: Final = await MCPRequestHandler._get_team_object_permission(caller_auth) - team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(team_obj_perm, server_id) + team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server( + team_obj_perm, server_id, requires_fresh_policy=caller_auth.requires_fresh_policy + ) except Exception as e: # noqa: BLE001 # an unresolved caller team must deny, not widen verbose_logger.warning( "MCP agent caller team tool ceiling unresolvable, denying tools on %r: %s", server_id, e @@ -3241,7 +3347,11 @@ class MCPRequestHandler: end_user_direct_tools: Final = global_mcp_server_manager.expand_tool_permissions( object_permissions.mcp_tool_permissions ).get(server_id) - end_user_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(object_permissions, server_id) + end_user_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server( + object_permissions, + server_id, + requires_fresh_policy=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), + ) end_user_tools: Final = MCPRequestHandler._union_tool_grants(end_user_direct_tools, end_user_toolset_tools) if end_user_tools is None: return allowed_tools @@ -3302,6 +3412,11 @@ class MCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.agent_id: return None + managed: Final = managed_agent_policy(user_api_key_auth) + if managed is not None: + permission: Final = managed.object_permission + return LiteLLM_ObjectPermissionTable.model_validate(permission) if permission is not None else None + if prisma_client is None: verbose_logger.debug("prisma_client is None") return None @@ -3319,7 +3434,7 @@ class MCPRequestHandler: ) @staticmethod - async def _get_allowed_mcp_servers_for_agent( + async def get_allowed_mcp_servers_for_agent( user_api_key_auth: UserAPIKeyAuth | None = None, agent_object_permission: LiteLLM_ObjectPermissionTable | None = None, ) -> list[str]: @@ -3358,12 +3473,16 @@ class MCPRequestHandler: obj_perm.mcp_servers or [] ) access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups( - obj_perm.mcp_access_groups or [] + obj_perm.mcp_access_groups or [], + requires_fresh_policy=user_api_key_auth.requires_fresh_policy, ) - toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(obj_perm) - return list({*expanded_direct_servers, *access_group_servers, *toolset_grants}) + toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions( + obj_perm, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + ) + inline_tools: Final = global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions) + return list({*expanded_direct_servers, *access_group_servers, *toolset_grants, *inline_tools}) except Exception as e: - if isinstance(e, UnloadableEntitlementError): + if managed_agent_policy(user_api_key_auth) is not None or isinstance(e, UnloadableEntitlementError): raise verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e) return [] @@ -3390,7 +3509,7 @@ class MCPRequestHandler: return frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids))) @staticmethod - async def _get_agent_tool_permissions_for_server( + async def get_agent_tool_permissions_for_server( server_id: str, user_api_key_auth: UserAPIKeyAuth | None = None, agent_object_permission: LiteLLM_ObjectPermissionTable | None = None, @@ -3430,11 +3549,13 @@ class MCPRequestHandler: if obj_perm.mcp_tool_permissions else None ) - toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(obj_perm, server_id) + toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server( + obj_perm, server_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + ) agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools) - return list(agent_tools) if agent_tools else None + return list(agent_tools) if agent_tools is not None else None except Exception as e: - if isinstance(e, UnloadableEntitlementError): + if managed_agent_policy(user_api_key_auth) is not None or isinstance(e, UnloadableEntitlementError): raise verbose_logger.warning("Failed to get agent tool permissions for server: %s", e) return None @@ -3452,28 +3573,38 @@ class MCPRequestHandler: return server_ids @staticmethod - async def _get_db_server_ids_for_access_groups(prisma_client, access_groups: list[str]) -> set[str]: + async def _get_db_server_ids_for_access_groups( + prisma_client, + access_groups: list[str], + *, + use_writer: bool = False, + ) -> set[str]: """ Helper to get server_ids from DB servers that match any of the given access groups. """ server_ids: Final[set[str]] = set() if access_groups and prisma_client is not None: try: - mcp_servers: Final = await MCPServerRepository(prisma_client).table.find_many( + mcp_servers: Final = await MCPServerRepository(prisma_client, use_writer=use_writer).table.find_many( where={"mcp_access_groups": {"hasSome": access_groups}} ) for server in mcp_servers: server_ids.add(server.server_id) except Exception as e: + if use_writer: + raise verbose_logger.debug("Error getting MCP servers from access groups: %s", e) return server_ids @staticmethod async def _get_mcp_servers_from_access_groups( access_groups: list[str], + *, + requires_fresh_policy: bool = False, ) -> list[str]: """ - Resolve MCP access groups to server IDs by querying BOTH the MCP server table (DB) AND config-loaded servers + Resolve MCP access groups to server IDs by querying BOTH the MCP server table (DB) AND config-loaded servers. + ``requires_fresh_policy`` reads the writer and propagates a read fault instead of resolving to no servers. """ from litellm.proxy.proxy_server import prisma_client @@ -3489,11 +3620,15 @@ class MCPRequestHandler: ) # Use the new helper for DB servers - db_server_ids = await MCPRequestHandler._get_db_server_ids_for_access_groups(prisma_client, access_groups) + db_server_ids = await MCPRequestHandler._get_db_server_ids_for_access_groups( + prisma_client, access_groups, use_writer=requires_fresh_policy + ) server_ids.update(db_server_ids) return list(server_ids) except Exception as e: + if requires_fresh_policy: + raise verbose_logger.warning("Failed to get MCP servers from access groups: %s", e) return [] @@ -3548,6 +3683,7 @@ class MCPRequestHandler: user_api_key_cache=user_api_key_cache, parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) if key_object_permission is None: return [] @@ -3591,6 +3727,7 @@ class MCPRequestHandler: user_api_key_cache=user_api_key_cache, parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) if team_obj is None: verbose_logger.debug("team_obj is None") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 6341f7972f0..7509f62f7f4 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -181,6 +181,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, is_per_server_oauth_discovery_eligible, ) +from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl @@ -3494,7 +3495,9 @@ class MCPServerManager: ``allow_all_server_ids`` / ``submitted_server_ids`` are injectable so the server union, which precomputes both for its fallback path, does not compute them twice.""" - if user_api_key_auth is not None and user_api_key_auth.mcp_toolset_id is not None: + if user_api_key_auth is not None and ( + user_api_key_auth.mcp_toolset_id is not None or user_api_key_auth.mcp_explicit_grants_only + ): return set() if allow_all_server_ids is None: allow_all_server_ids = self.get_allow_all_keys_server_ids() @@ -3543,9 +3546,14 @@ class MCPServerManager: 2. If admin and no object_permission, return all servers 3. Otherwise, use standard permission checks """ + if managed_agent_policy(user_api_key_auth) is not None: + managed: Final = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + return managed if access is None else [server for server in managed if server in access.server_ids] + from litellm.proxy.proxy_server import general_settings as proxy_general_settings resolved_general_settings: Final = proxy_general_settings if general_settings is None else general_settings + explicit_grants_only: Final = bool(user_api_key_auth and user_api_key_auth.mcp_explicit_grants_only) allow_all_server_ids: Final = self.get_allow_all_keys_server_ids() # A keyless admitted subject is resolved per grant source, and channel decisions that are @@ -3577,7 +3585,7 @@ class MCPServerManager: # only keys without their own mcp_servers list get submitted servers unioned in. submitted_server_ids: Final = ( [] - if has_explicit_object_permission + if has_explicit_object_permission or explicit_grants_only else await self._get_active_submitted_mcp_server_ids_for_user(user_api_key_auth) ) @@ -3646,12 +3654,14 @@ class MCPServerManager: return [ server_id for server_id in dict.fromkeys(allow_all_server_ids + submitted_server_ids) - if scope is None or server_id == scope + if not explicit_grants_only and (scope is None or server_id == scope) ] async def resolve_toolset_tool_permissions( self, toolset_ids: list[str], + *, + requires_fresh_policy: bool = False, ) -> dict[str, list[str]]: """ Resolve a list of toolset IDs into a mcp_tool_permissions dict. @@ -3661,6 +3671,10 @@ class MCPServerManager: Redis-backed ``DualCache`` in production) so that cache entries are shared across workers and cold-cache DB hits are minimised. + ``requires_fresh_policy`` bypasses the cache and reads the writer so a + revocation is honoured on the very next request; a read fault then + propagates instead of resolving to no grants. + A row names a tool on the server identified by ``server_id``, so the stored name is the tool's own name and is used as written. It is never reduced by the server's wire prefix: that prefix is added on the way out @@ -3675,12 +3689,16 @@ class MCPServerManager: return {} cache_key: Final = "toolset_perms:" + ",".join(sorted(toolset_ids)) - cached: Final[dict[str, list[str]] | None] = await user_api_key_cache.async_get_cache(key=cache_key) + cached: Final[dict[str, list[str]] | None] = ( + None if requires_fresh_policy else await user_api_key_cache.async_get_cache(key=cache_key) + ) if cached is not None: return cached try: - toolsets: Final = await list_mcp_toolsets(prisma_client, toolset_ids=toolset_ids) + toolsets: Final = await list_mcp_toolsets( + prisma_client, toolset_ids=toolset_ids, use_writer=requires_fresh_policy + ) tool_permissions: Final[dict[str, list[str]]] = {} for toolset in toolsets: for tool in toolset.tools: @@ -3694,6 +3712,8 @@ class MCPServerManager: ) return tool_permissions except Exception as e: + if requires_fresh_policy: + raise verbose_logger.warning("Failed to resolve toolset permissions: %s", e) return {} diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index 48bad178927..dcbd0064514 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -65,9 +65,9 @@ class MCPToolsetTable(Protocol): async def delete(self, where: Mapping[str, object]) -> MCPToolsetRow: ... -def _toolset_table(prisma_client: PrismaClient) -> MCPToolsetTable: +def _toolset_table(prisma_client: PrismaClient, *, use_writer: bool = False) -> MCPToolsetTable: """The toolset table actions of the prisma client.""" - return MCPToolsetRepository(prisma_client).table + return MCPToolsetRepository(prisma_client, use_writer=use_writer).table def _toolset_from_row(row: MCPToolsetRow) -> MCPToolset: @@ -107,12 +107,16 @@ async def get_mcp_toolset( async def list_mcp_toolsets( prisma_client: PrismaClient, toolset_ids: Sequence[str] | None = None, + *, + use_writer: bool = False, ) -> Sequence[MCPToolset]: try: where: Final[Mapping[str, object]] = {} if toolset_ids is None else {"toolset_id": {"in": toolset_ids}} - rows: Final = await _toolset_table(prisma_client).find_many(where=where) + rows: Final = await _toolset_table(prisma_client, use_writer=use_writer).find_many(where=where) return [_toolset_from_row(r) for r in rows] except Exception as e: + if use_writer: + raise verbose_proxy_logger.warning("litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - %s", e) return [] diff --git a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py b/litellm/proxy/_experimental/mcp_server/ui_session_utils.py index 107a4818de1..901259c18ad 100644 --- a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py +++ b/litellm/proxy/_experimental/mcp_server/ui_session_utils.py @@ -58,6 +58,7 @@ async def resolve_ui_session_team_ids( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=False, + check_db_only=user_api_key_auth.requires_fresh_policy, parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) @@ -92,7 +93,9 @@ async def admitted_user_context(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKey ) try: - admitted: Final = await MCPRequestHandler.reload_admitted_user(user_id) + admitted: Final = await MCPRequestHandler.reload_admitted_user( + user_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + ) except HTTPException as e: verbose_logger.warning("MCP dashboard session: admitted-subject reload failed for %s: %s", user_id, e.detail) return None diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index bb063f8f77a..78b7b729375 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -2521,6 +2521,91 @@ "title": "AgentExtension", "type": "object" }, + "AgentIdentityBinding": { + "properties": { + "active": { + "default": true, + "title": "Active", + "type": "boolean" + }, + "agent_id": { + "title": "Agent Id", + "type": "string" + }, + "client_id": { + "title": "Client Id", + "type": "string" + }, + "issuer": { + "title": "Issuer", + "type": "string" + }, + "last_authenticated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Last Authenticated At" + }, + "provider": { + "const": "microsoft_entra", + "title": "Provider", + "type": "string" + }, + "required_roles": { + "default": [], + "items": { + "type": "string" + }, + "title": "Required Roles", + "type": "array" + }, + "required_scopes": { + "default": [ + "user_impersonation" + ], + "items": { + "type": "string" + }, + "title": "Required Scopes", + "type": "array" + }, + "revision": { + "title": "Revision", + "type": "string" + }, + "service_principal_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Service Principal Id" + }, + "tenant_id": { + "title": "Tenant Id", + "type": "string" + } + }, + "required": [ + "agent_id", + "provider", + "tenant_id", + "client_id", + "issuer", + "revision" + ], + "title": "AgentIdentityBinding", + "type": "object" + }, "AgentInterface": { "description": "Declares a combination of a target URL and a transport protocol.", "properties": { @@ -2972,6 +3057,21 @@ ], "title": "Created By" }, + "enabled": { + "default": true, + "title": "Enabled", + "type": "boolean" + }, + "execution_mode": { + "default": "autonomous", + "enum": [ + "autonomous", + "delegated", + "both" + ], + "title": "Execution Mode", + "type": "string" + }, "extra_headers": { "anyOf": [ { @@ -2986,6 +3086,26 @@ ], "title": "Extra Headers" }, + "identity": { + "anyOf": [ + { + "$ref": "#/components/schemas/AgentIdentityBinding" + }, + { + "type": "null" + } + ] + }, + "identity_managed": { + "default": false, + "title": "Identity Managed", + "type": "boolean" + }, + "jwt_auth_configured": { + "default": false, + "title": "Jwt Auth Configured", + "type": "boolean" + }, "keys": { "anyOf": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d9fb053035b..be06d2e7321 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -15,6 +15,7 @@ from pydantic import ( Json, JsonValue, PositiveInt, + PrivateAttr, field_validator, model_validator, ) @@ -27,7 +28,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( validate_langfuse_span_scope_value, validate_no_callback_env_reference, ) -from litellm.types.agents import AgentCaller +from litellm.types.agents import AgentCaller, AgentResponse from litellm.types.integrations.compression_interception import ( CompressionSavingsMetadata, ) @@ -46,6 +47,7 @@ from litellm.types.mcp import ( MCPTransportType, ) from litellm.types.mcp_server.mcp_server_manager import MCPInfo +from litellm.types.proxy.agent_identity import ManagedAgentContext from litellm.types.proxy.carried_budget_state import ( OrgBudgetSnapshot, TeamBudgetSnapshot, @@ -567,6 +569,7 @@ class LiteLLMRoutes(enum.Enum): "/agents", "/a2a/{agent_id}", "/a2a/{agent_id}/message/send", + "/v1/a2a/{agent_id}/message/send", "/a2a/{agent_id}/message/stream", "/a2a/{agent_id}/.well-known/agent-card.json", ) @@ -3302,6 +3305,8 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob # metadata or JWT claims, so it cannot be forged to gain the team-inherited MCP grant union # or to escape the caller-Authorization egress scrub. exclude=True keeps it out of serialization. mcp_admitted_user_subject: bool = Field(default=False, exclude=True) + requires_fresh_policy: bool = Field(default=False, exclude=True) + mcp_explicit_grants_only: bool = Field(default=False, exclude=True) # team_id -> that team's mcp_rpm_limit map, for a keyless admitted subject that reaches MCP # servers through several teams at once and therefore has no single team_id for the limiter to # key off. Server-only and stripped from validated input for the same reason as the marker @@ -3326,6 +3331,13 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob "user id." ), ) + invoked_agent_id: str | None = Field(default=None, exclude=True) + invoked_agent_policy: AgentResponse | None = Field(default=None, exclude=True) + agent_invocation_cost: float | None = Field(default=None, exclude=True) + billing_agent_policy: AgentResponse | None = Field(default=None, exclude=True) + _managed_delegation_verified: bool = PrivateAttr(default=False) + managed_agent_policy: AgentResponse | None = Field(default=None, exclude=True) + managed_agent_context: ManagedAgentContext | None = Field(default=None, exclude=True) agent_caller: AgentCaller | None = Field( default=None, exclude=True, @@ -3363,11 +3375,19 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob # path via post-construction assignment. Strip it from any validated input (constructor # kwargs, model_validate, a JWT/key claim splat) so it can never be forged from caller data. values.pop("mcp_admitted_user_subject", None) + values.pop("requires_fresh_policy", None) + values.pop("mcp_explicit_grants_only", None) values.pop("mcp_source_team_rpm_limits", None) values.pop("mcp_session_resource_server_id", None) values.pop("mcp_toolset_id", None) values.pop("via_virtual_key", None) values.pop("agent_caller", None) + values.pop("managed_agent_context", None) + values.pop("managed_agent_policy", None) + values.pop("invoked_agent_id", None) + values.pop("invoked_agent_policy", None) + values.pop("agent_invocation_cost", None) + values.pop("billing_agent_policy", None) if values.get("api_key") is not None: values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))}) if isinstance(values.get("api_key"), str): @@ -4063,6 +4083,11 @@ class SpendLogsRouterMetadata(TypedDict): class SpendLogsMetadata(TypedDict): + actor_agent_id: ReadOnly[NotRequired[str | None]] + target_agent_id: ReadOnly[NotRequired[str | None]] + billing_agent_id: ReadOnly[NotRequired[str | None]] + agent_execution_mode: ReadOnly[NotRequired[str | None]] + verified_human_user_id: ReadOnly[NotRequired[str | None]] autorouter_baseline_observation: ReadOnly[str | None] """ Specific metadata k,v pairs logged to spendlogs for easier cost tracking @@ -4126,6 +4151,7 @@ class SpendLogsPayload(TypedDict): model_id: str | None model_group: str | None mcp_namespaced_tool_name: str | None + billing_agent_id: ReadOnly[NotRequired[str | None]] agent_id: str | None api_base: str user: str @@ -5048,6 +5074,7 @@ class JWTAuthBuilderResult(TypedDict): org_id: str | None team_membership: LiteLLM_TeamMembership | None jwt_claims: dict # Decoded JWT token claims (avoids re-decoding) + managed_agent_context: ReadOnly[NotRequired[ManagedAgentContext | None]] agent_id: ReadOnly[str | None] diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 2a189a76545..e94d5e7ea78 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -597,7 +597,6 @@ async def get_agent_card( if agent is None: raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found") - # Check agent permission (skip for admin users) is_allowed: Final = await AgentRequestHandler.is_agent_allowed( agent_id=agent.agent_id, user_api_key_auth=user_api_key_dict, diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py index 49e5407ff88..67547e82f24 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -1,13 +1,16 @@ import asyncio from collections.abc import Awaitable, Callable from dataclasses import dataclass -from typing import Final, TypeAlias +from typing import TYPE_CHECKING, Final, TypeAlias from fastapi import HTTPException from litellm._logging import verbose_proxy_logger from litellm.proxy._types import LiteLLM_AccessGroupTable +if TYPE_CHECKING: + from litellm.types.agents import AgentResponse + AccessGroupIds: TypeAlias = tuple[str, ...] AccessGroupIdsLoader: TypeAlias = Callable[[str], Awaitable[AccessGroupIds]] # mutable-ok: Callable params LoadedAccessGroup: TypeAlias = LiteLLM_AccessGroupTable | None @@ -34,7 +37,7 @@ async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: return tuple(agent.access_group_ids or ()) if agent is not None else () -async def _load_access_group(access_group_id: str) -> LoadedAccessGroup: +async def _load_access_group(access_group_id: str, *, check_db_only: bool = False) -> LoadedAccessGroup: from litellm.proxy.auth.auth_checks import get_access_object from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache @@ -47,8 +50,11 @@ async def _load_access_group(access_group_id: str) -> LoadedAccessGroup: prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ) except HTTPException as e: + if check_db_only: + raise verbose_proxy_logger.warning( "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, e.detail ) @@ -59,13 +65,20 @@ async def resolve_agent_access_group_ceiling( agent_id: str, load_access_group_ids: AccessGroupIdsLoader = _registry_access_group_ids, load_access_group: AccessGroupLoader = _load_access_group, + *, + check_db_only: bool = False, ) -> AgentAccessGroupCeiling | None: """``None`` when the agent has no access groups attached, so nothing is capped.""" access_group_ids: Final = await load_access_group_ids(agent_id) if not access_group_ids: return None - loaded: Final = await asyncio.gather(*(load_access_group(group_id) for group_id in access_group_ids)) + loaded: Final = await asyncio.gather( + *( + _load_access_group(group_id, check_db_only=True) if check_db_only else load_access_group(group_id) + for group_id in access_group_ids + ) + ) groups: Final = tuple(group for group in loaded if group is not None) return AgentAccessGroupCeiling( access_group_ids=access_group_ids, @@ -73,3 +86,16 @@ async def resolve_agent_access_group_ceiling( mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids), agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids), ) + + +async def resolve_managed_agent_ceilings(agent: "AgentResponse") -> tuple[AgentAccessGroupCeiling, ...]: + async def authoritative_group(group_id: str) -> LoadedAccessGroup: + return await _load_access_group(group_id, check_db_only=True) + + async def manual_ids(_agent_id: str) -> AccessGroupIds: + return tuple(agent.access_group_ids or ()) + + manual: Final = await resolve_agent_access_group_ceiling( + agent.agent_id, load_access_group_ids=manual_ids, load_access_group=authoritative_group + ) + return (manual,) if manual is not None else () diff --git a/litellm/proxy/agent_endpoints/auth/agent_caller.py b/litellm/proxy/agent_endpoints/auth/agent_caller.py index 47d43e8f71b..1ff6f1ffe04 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_caller.py +++ b/litellm/proxy/agent_endpoints/auth/agent_caller.py @@ -8,6 +8,7 @@ can only narrow access and need no trust. """ from collections.abc import Mapping +from types import MappingProxyType from typing import Final from litellm._logging import verbose_proxy_logger @@ -45,7 +46,7 @@ def agent_caller_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth | Non user_id=caller.user_id, team_id=caller.team_id, parent_otel_span=user_api_key_auth.parent_otel_span, - ) + ).model_copy(update=MappingProxyType({"requires_fresh_policy": user_api_key_auth.requires_fresh_policy})) async def load_agent_caller_team(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 9fe74bfee3f..75b99b0ab79 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -8,8 +8,11 @@ Follows the same pattern as MCP permission handling. import asyncio from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import Final, TypeAlias +from fastapi import HTTPException + from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.ui_session_utils import build_effective_auth_contexts from litellm.proxy._types import ( @@ -24,6 +27,7 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( resolve_agent_access_group_ceiling, ) from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth +from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy from litellm.repositories.table_repositories import AgentsRepository from litellm.types.agents import AgentResponse @@ -83,13 +87,23 @@ class AgentRequestHandler: async def resolve_agent_access( user_api_key_auth: UserAPIKeyAuth | None = None, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, + *, + strict: bool = False, ) -> AgentAccess: """Agents the key may reach: key and team grants, intersected with the agent's access group ceiling and, for an agent key acting on behalf of an invoking user, with that user's team grants.""" - key_team_access: Final = await AgentRequestHandler._resolve_key_team_agent_access(user_api_key_auth) - caller_access: Final = await AgentRequestHandler._agent_caller_access(user_api_key_auth) + if managed_agent_policy(user_api_key_auth) is not None: + return await _managed_actor_agent_access(user_api_key_auth) + key_team_access: Final = await AgentRequestHandler.resolve_key_team_agent_access( + user_api_key_auth, strict=strict + ) + if strict and isinstance(key_team_access, UnrestrictedAgentAccess): + return RestrictedAgentAccess(frozenset()) + caller_access: Final = await AgentRequestHandler.agent_caller_access(user_api_key_auth, strict=strict) own_access: Final = _intersect_agent_access(key_team_access, caller_access) - agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth, resolve_ceiling) + agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling( + user_api_key_auth, resolve_ceiling, strict=strict + ) if agent_ceiling is None: return own_access if isinstance(own_access, UnrestrictedAgentAccess): @@ -97,20 +111,26 @@ class AgentRequestHandler: return RestrictedAgentAccess(own_access.agent_ids & agent_ceiling) @staticmethod - async def _agent_caller_access(user_api_key_auth: UserAPIKeyAuth | None) -> AgentAccess: + async def agent_caller_access(user_api_key_auth: UserAPIKeyAuth | None, *, strict: bool = False) -> AgentAccess: caller_auth: Final = agent_caller_auth(user_api_key_auth) if user_api_key_auth else None if caller_auth is None: return UnrestrictedAgentAccess() - return await AgentRequestHandler._get_allowed_agents_for_team(caller_auth) + return await AgentRequestHandler._get_allowed_agents_for_team(caller_auth, strict=strict) @staticmethod - async def _resolve_key_team_agent_access( + async def resolve_key_team_agent_access( user_api_key_auth: UserAPIKeyAuth | None, + *, + strict: bool = False, ) -> AgentAccess: try: - key_access: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) - team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) + key_access: Final = await AgentRequestHandler.get_allowed_agents_for_key(user_api_key_auth, strict=strict) + team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team( + user_api_key_auth, strict=strict + ) except Exception as e: + if strict: + raise HTTPException(503, "Agent invocation policy is unavailable") from e verbose_logger.warning("Failed to get allowed agents: %s", e) return UnrestrictedAgentAccess() return _intersect_agent_access(key_access, team_access) @@ -119,10 +139,16 @@ class AgentRequestHandler: async def _agent_access_group_ceiling( user_api_key_auth: UserAPIKeyAuth | None, resolve_ceiling: CeilingResolver, + *, + strict: bool = False, ) -> frozenset[str] | None: if user_api_key_auth is None or not user_api_key_auth.agent_id: return None - ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id) + ceiling: Final = ( + await resolve_agent_access_group_ceiling(user_api_key_auth.agent_id, check_db_only=True) + if strict + else await resolve_ceiling(user_api_key_auth.agent_id) + ) if ceiling is None: return None return _to_stable_ids(ceiling.agent_ids) @@ -144,6 +170,49 @@ class AgentRequestHandler: bool: True if agent is allowed, False otherwise """ from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore + from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure + from litellm.proxy.proxy_server import prisma_client + from litellm.types.proxy.agent_identity import AgentIdentityFailure + + registered: Final = global_agent_registry.get_agent_by_id(agent_id) + registry_managed: Final = isinstance(registered, AgentResponse) and registered.identity_managed + if registry_managed or (registered is None and prisma_client is not None): + target: Final = await AgentIdentityStore.from_client(prisma_client).agent(agent_id) + if isinstance(target, AgentIdentityFailure): + if registry_managed: + raise_identity_failure(target) + elif target is None and registry_managed: + return False + elif isinstance(target, AgentResponse) and target.identity_managed: + if ( + not target.enabled + or target.identity is None + or not target.identity.active + or user_api_key_auth is None + ): + return False + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + + key_hash: Final = user_api_key_auth.api_key or user_api_key_auth.token + authority: Final = ( + await MCPRequestHandler._reload_admitted_key(key_hash, check_db_only=True) # pyright: ignore[reportPrivateUsage] # the authoritative key reload has no public seam + if key_hash + and managed_agent_policy(user_api_key_auth) is None + and not user_api_key_auth.is_session_token + else user_api_key_auth + ) + fresh_auth: Final = authority.model_copy( + update=MappingProxyType( + {"requires_fresh_policy": True, "agent_caller": user_api_key_auth.agent_caller} + ) + ) + explicit: Final = await _granted_agent_ids( + fresh_auth, + _strict_agent_access, + build_effective_auth_contexts, + ) + return target.agent_id in explicit match await AgentRequestHandler.resolve_agent_access(user_api_key_auth, resolve_ceiling): case UnrestrictedAgentAccess(): @@ -202,8 +271,10 @@ class AgentRequestHandler: return team_obj.object_permission @staticmethod - async def _get_allowed_agents_for_key( + async def get_allowed_agents_for_key( user_api_key_auth: UserAPIKeyAuth | None = None, + *, + strict: bool = False, ) -> AgentAccess: """ Get allowed agents for a key. @@ -237,24 +308,36 @@ class AgentRequestHandler: return UnrestrictedAgentAccess() access_group_agents: Final = ( - tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups))) + tuple( + await AgentRequestHandler._get_agents_from_access_groups( + declared_access_groups, check_db_only=strict + ) + ) if declared_access_groups else () ) unified_agents: Final = ( - tuple(await AgentRequestHandler._get_unified_access_group_agents(list(key_access_group_ids))) + tuple( + await AgentRequestHandler._get_unified_access_group_agents( + key_access_group_ids, check_db_only=strict + ) + ) if key_access_group_ids else () ) return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) except Exception as e: + if strict: + raise HTTPException(503, "Agent invocation policy is unavailable") from e verbose_logger.warning("Failed to get allowed agents for key: %s", e) return UnrestrictedAgentAccess() @staticmethod async def _get_allowed_agents_for_team( user_api_key_auth: UserAPIKeyAuth | None = None, + *, + strict: bool = False, ) -> AgentAccess: """ Get allowed agents for a team. @@ -263,7 +346,7 @@ class AgentRequestHandler: 2. Also includes agents from team's access_group_ids (unified access groups) Fetches the team object once and reuses it for both permission sources. - Declared-but-empty grants stay restricted; see `_get_allowed_agents_for_key`. + Declared-but-empty grants stay restricted; see `get_allowed_agents_for_key`. """ if user_api_key_auth is None: return UnrestrictedAgentAccess() @@ -280,7 +363,7 @@ class AgentRequestHandler: ) if not prisma_client: - return UnrestrictedAgentAccess() + return RestrictedAgentAccess(frozenset()) if strict else UnrestrictedAgentAccess() # Fetch the team object once for both permission sources team_obj: Final = await get_team_object( @@ -289,10 +372,11 @@ class AgentRequestHandler: user_api_key_cache=user_api_key_cache, parent_otel_span=user_api_key_auth.parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=strict, ) if team_obj is None: - return UnrestrictedAgentAccess() + return RestrictedAgentAccess(frozenset()) if strict else UnrestrictedAgentAccess() # 1. Get agents from object_permission (native permissions) object_permissions: Final = team_obj.object_permission @@ -307,18 +391,28 @@ class AgentRequestHandler: return UnrestrictedAgentAccess() access_group_agents: Final = ( - tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups))) + tuple( + await AgentRequestHandler._get_agents_from_access_groups( + declared_access_groups, check_db_only=strict + ) + ) if declared_access_groups else () ) unified_agents: Final = ( - tuple(await AgentRequestHandler._get_unified_access_group_agents(list(team_access_group_ids))) + tuple( + await AgentRequestHandler._get_unified_access_group_agents( + team_access_group_ids, check_db_only=strict + ) + ) if team_access_group_ids else () ) return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) except Exception as e: + if strict: + raise HTTPException(503, "Agent invocation policy is unavailable") from e # litellm-dashboard is the default UI team and will never have agents; # skip noisy warnings for it. if user_api_key_auth.team_id != UI_TEAM_ID: @@ -326,7 +420,9 @@ class AgentRequestHandler: return UnrestrictedAgentAccess() @staticmethod - def _get_config_agent_ids_for_access_groups(config_agents: list, access_groups: list[str]) -> set[str]: + def _get_config_agent_ids_for_access_groups( + config_agents: Sequence[AgentResponse], access_groups: Sequence[str] + ) -> set[str]: """ Helper to get agent_ids from config-loaded agents that match any of the given access groups. """ @@ -339,7 +435,9 @@ class AgentRequestHandler: return server_ids @staticmethod - async def _get_db_agent_ids_for_access_groups(prisma_client, access_groups: list[str]) -> set[str]: + async def _get_db_agent_ids_for_access_groups( + prisma_client, access_groups: Sequence[str], *, check_db_only: bool = False + ) -> set[str]: """ Helper to get agent_ids from DB agents that match any of the given access groups. @@ -349,23 +447,27 @@ class AgentRequestHandler: if not access_groups or prisma_client is None: return set() - agents: Final = await AgentsRepository(prisma_client).table.find_many( + agents: Final = await AgentsRepository(prisma_client, use_writer=check_db_only).table.find_many( where={"agent_access_groups": {"hasSome": access_groups}} ) return {agent.agent_id for agent in agents} @staticmethod - async def _get_unified_access_group_agents(access_group_ids: list[str]) -> list[str]: + async def _get_unified_access_group_agents( + access_group_ids: Sequence[str], *, check_db_only: bool = False + ) -> list[str]: """ Resolve unified access group ids to agent IDs. """ from litellm.proxy.auth.auth_checks import _get_agent_ids_from_access_groups - return await _get_agent_ids_from_access_groups(access_group_ids=access_group_ids) + return await _get_agent_ids_from_access_groups(access_group_ids=access_group_ids, check_db_only=check_db_only) @staticmethod async def _get_agents_from_access_groups( - access_groups: list[str], + access_groups: Sequence[str], + *, + check_db_only: bool = False, ) -> list[str]: """ Resolve agent access groups to agent IDs by querying BOTH the agent table (DB) AND config-loaded agents. @@ -373,14 +475,13 @@ class AgentRequestHandler: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.proxy_server import prisma_client - # Use the helper for config-loaded agents config_agent_ids: Final = AgentRequestHandler._get_config_agent_ids_for_access_groups( global_agent_registry.agent_list, access_groups ) # Use the helper for DB agents db_agent_ids: Final = await AgentRequestHandler._get_db_agent_ids_for_access_groups( - prisma_client, access_groups + prisma_client, access_groups, check_db_only=check_db_only ) return list(config_agent_ids | db_agent_ids) @@ -531,4 +632,60 @@ async def accessible_agents( AgentRequestHandler.resolve_agent_access if resolve_access is None else resolve_access, effective_contexts, ) - return tuple(agent for agent in agents if agent.agent_id in allowed_agent_ids) + allowed: Final = await asyncio.gather( + *( + AgentRequestHandler.is_agent_allowed(agent.agent_id, user_api_key_auth) + for agent in agents + if agent.identity_managed + ) + ) + managed_ids: Final = frozenset( + agent.agent_id + for agent, permitted in zip((agent for agent in agents if agent.identity_managed), allowed) + if permitted + ) + return tuple( + agent + for agent in agents + if (agent.agent_id in managed_ids if agent.identity_managed else agent.agent_id in allowed_agent_ids) + ) + + +async def _strict_agent_access(auth: UserAPIKeyAuth) -> AgentAccess: + return await AgentRequestHandler.resolve_agent_access(auth, strict=True) + + +async def _managed_actor_agent_access(auth: UserAPIKeyAuth) -> AgentAccess: + agent: Final = managed_agent_policy(auth) + if agent is None or not agent.object_permission: + return RestrictedAgentAccess(frozenset()) + permission: Final = LiteLLM_ObjectPermissionTable.model_validate(agent.object_permission or MappingProxyType({})) + own_auth: Final = UserAPIKeyAuth(object_permission=permission) + own: Final = _granted_ids(await AgentRequestHandler.get_allowed_agents_for_key(own_auth, strict=True)) + + from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings + + ceilings: Final = await resolve_managed_agent_ceilings(agent) + grouped: Final = frozenset(target for target in own if all(target in ceiling.agent_ids for ceiling in ceilings)) + caller: Final = await AgentRequestHandler.agent_caller_access(auth, strict=True) + capped: Final = grouped if isinstance(caller, UnrestrictedAgentAccess) else grouped & caller.agent_ids + context: Final = auth.managed_agent_context + if context is None or context.mode == "autonomous": + return RestrictedAgentAccess(capped) + if context.user_id is None: + return RestrictedAgentAccess(frozenset()) + human_ids: Final = await verified_human_agent_grants(context.user_id, auth.team_id) + return RestrictedAgentAccess(capped.intersection(human_ids)) + + +async def verified_human_agent_grants(user_id: str | None, team_id: str | None = None) -> frozenset[str]: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + + if user_id is None: + return frozenset() + human: Final = await MCPRequestHandler.reload_admitted_user(user_id, requires_fresh_policy=True) + sources: Final = await MCPRequestHandler.admitted_subject_sources( + human, allowed_team_ids=frozenset((team_id,)) if team_id else frozenset() + ) + human_access: Final = await asyncio.gather(*(_strict_agent_access(source) for source in sources)) + return frozenset().union(*(_granted_ids(access) for access in human_access)) diff --git a/litellm/proxy/agent_endpoints/auth/managed_authorization.py b/litellm/proxy/agent_endpoints/auth/managed_authorization.py new file mode 100644 index 00000000000..a15cf074ad5 --- /dev/null +++ b/litellm/proxy/agent_endpoints/auth/managed_authorization.py @@ -0,0 +1,84 @@ +from typing import Final + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore +from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import AgentIdentityFailure, ManagedAgentContext + + +def managed_agent_policy(auth: "UserAPIKeyAuth | None") -> AgentResponse | None: + """The admitted managed policy, or ``None`` when the subject was never admitted as a managed agent. + + ``admit_managed_actor`` only assigns ``managed_agent_policy`` after ``actor_admission_failure`` + has verified the bound context, so an ``AgentResponse`` here means admission succeeded. + """ + policy: Final = auth.managed_agent_policy if auth is not None else None + return policy if isinstance(policy, AgentResponse) else None + + +async def admit_managed_actor(auth: UserAPIKeyAuth, store: AgentIdentityStore | None) -> None: + delegation_verified: Final = auth._managed_delegation_verified # pyright: ignore[reportPrivateUsage] # the one-shot delegation marker is a PrivateAttr by design + auth._managed_delegation_verified = False # pyright: ignore[reportPrivateUsage] # consumed here so a replayed token cannot reuse it + if auth.agent_id is None: + return + if store is None: + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + registered: Final = global_agent_registry.get_agent_by_id(auth.agent_id) + if auth.managed_agent_context is not None or ( + registered is not None and (registered.identity_managed or registered.identity is not None) + ): + raise_identity_failure( + AgentIdentityFailure(code="policy_unavailable", message="Managed agent policy requires a database") + ) + return + agent: Final = await store.agent(auth.agent_id) + if isinstance(agent, AgentIdentityFailure): + raise_identity_failure(agent) + if agent is None: + retired: Final = await store.retired_agent(auth.agent_id) + if isinstance(retired, AgentIdentityFailure): + raise_identity_failure(retired) + if auth.managed_agent_context is not None or retired: + raise_identity_failure(AgentIdentityFailure(message="Agent no longer exists")) + return + if not agent.identity_managed: + return + if auth.jwt_claims and auth.managed_agent_context is None: + raise_identity_failure(AgentIdentityFailure(message="A managed agent requires a matching verified identity")) + failure: Final = actor_admission_failure(agent, auth.managed_agent_context) + if failure is not None: + raise_identity_failure(failure) + auth.managed_agent_policy = agent + auth.billing_agent_policy = agent + auth.requires_fresh_policy = True + if ( + auth.managed_agent_context is not None + and auth.managed_agent_context.mode == "delegated" + and not delegation_verified + ): + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import verified_human_agent_grants + + grants: Final = await verified_human_agent_grants(auth.managed_agent_context.user_id, auth.team_id) + if agent.agent_id not in grants: + raise_identity_failure( + AgentIdentityFailure(message="The delegated user is not permitted to invoke this agent") + ) + + +def actor_admission_failure( + agent: AgentResponse, + context: ManagedAgentContext | None, +) -> AgentIdentityFailure | None: + if not agent.enabled or agent.identity is None or not agent.identity.active: + return AgentIdentityFailure(message="Agent execution is disabled") + if context is None: + return AgentIdentityFailure(message="This agent requires its bound identity provider token") + if context.agent_id != agent.agent_id or context.binding_revision != agent.identity.revision: + return AgentIdentityFailure(message="Agent identity changed during authentication; retry") + if agent.execution_mode not in (context.mode, "both"): + return AgentIdentityFailure(message="Agent is not enabled for this execution mode") + if context.mode == "delegated" and not context.user_id: + return AgentIdentityFailure(message="A verified human subject is required") + return None diff --git a/litellm/proxy/agent_endpoints/identity.py b/litellm/proxy/agent_endpoints/identity.py new file mode 100644 index 00000000000..c0e5a748144 --- /dev/null +++ b/litellm/proxy/agent_endpoints/identity.py @@ -0,0 +1,17 @@ +from collections.abc import Mapping +from typing import Final + +from fastapi import HTTPException + +LEGACY_IDENTITY_MESSAGE: Final = ( + "litellm_params.identity is not supported: bind an Entra application through the top-level identity field" +) + + +def has_legacy_identity(params: Mapping[str, object] | None) -> bool: + return params is not None and "identity" in params + + +def reject_legacy_identity(params: Mapping[str, object] | None) -> None: + if has_legacy_identity(params): + raise HTTPException(400, LEGACY_IDENTITY_MESSAGE) diff --git a/litellm/proxy/agent_endpoints/identity_store.py b/litellm/proxy/agent_endpoints/identity_store.py new file mode 100644 index 00000000000..3c8163a8838 --- /dev/null +++ b/litellm/proxy/agent_endpoints/identity_store.py @@ -0,0 +1,252 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Final + +from litellm.proxy.agent_endpoints.managed_identity import classify_agent_subject +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl +from litellm.repositories.table_repositories import ( + AgentIdentityRepository, + AgentsRepository, + RetiredAgentIdentityRepository, + RetiredAgentRepository, + VerifiedSubjectRepository, +) +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import ( + AgentIdentityFailure, + ManagedAgentContext, + MicrosoftInteractiveSubject, + VerifiedHumanSubject, +) + +if TYPE_CHECKING: + from prisma.models import LiteLLM_VerifiedSubject + from prisma.types import ( + LiteLLM_AgentIdentityUpdateManyMutationInput, + LiteLLM_AgentIdentityWhereInput, + LiteLLM_AgentIdentityWhereUniqueInput, + LiteLLM_AgentsTableInclude, + LiteLLM_AgentsTableWhereUniqueInput, + LiteLLM_RetiredAgentWhereUniqueInput, + LiteLLM_VerifiedSubjectCreateInput, + LiteLLM_VerifiedSubjectUpsertInput, + LiteLLM_VerifiedSubjectWhereUniqueInput, + ) + + +class AgentIdentityStore: + @classmethod + def from_client(cls, client: object, *, cache: UserApiKeyCache | None = None) -> "AgentIdentityStore": + return cls( + AgentsRepository(client, use_writer=True), + AgentIdentityRepository(client, use_writer=True), + VerifiedSubjectRepository(client, use_writer=True), + RetiredAgentIdentityRepository(client, use_writer=True), + RetiredAgentRepository(client, use_writer=True), + cache=cache, + ) + + def __init__( + self, + agents: AgentsRepository, + identities: AgentIdentityRepository, + humans: VerifiedSubjectRepository, + retired: RetiredAgentIdentityRepository | None = None, + retired_agents: RetiredAgentRepository | None = None, + *, + cache: UserApiKeyCache | None = None, + ) -> None: + self.agents = agents + self.identities = identities + self.humans = humans + self.retired = retired + self.retired_agents = retired_agents + self.cache = cache + + async def agent(self, agent_id: str) -> AgentResponse | AgentIdentityFailure | None: + try: + where: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id} + include: Final[LiteLLM_AgentsTableInclude] = { + "identity": True, + "object_permission": True, + } + row: Final = await self.agents.table.find_unique(where=where, include=include) + if row is None: + return None + return AgentResponse.model_validate(row.model_dump()) + except Exception: + return AgentIdentityFailure(code="policy_unavailable", message="Agent policy could not be loaded") + + async def unbound_client(self, where: "LiteLLM_AgentIdentityWhereUniqueInput") -> AgentIdentityFailure | None: + if self.retired is not None: + try: + retired: Final = await self.retired.table.find_unique(where=where) + except Exception: + return AgentIdentityFailure( + code="policy_unavailable", message="Retired agent identity could not be checked" + ) + if retired is not None: + return AgentIdentityFailure(message="This agent identity binding has been retired") + return None + + async def _bound_agent_id(self, tenant_id: str, client_id: str) -> str | AgentIdentityFailure | None: + cache_key: Final = f"agent_identity:{json.dumps((tenant_id, client_id))}" + cached: Final[object] = await self.cache.async_get_cache(key=cache_key) if self.cache is not None else None + if isinstance(cached, str): + return cached + where: Final[LiteLLM_AgentIdentityWhereUniqueInput] = { + "provider_tenant_id_client_id": { + "provider": "microsoft_entra", + "tenant_id": tenant_id, + "client_id": client_id, + } + } + try: + row: Final = await self.identities.table.find_unique(where=where) + except Exception: + return AgentIdentityFailure(code="policy_unavailable", message="Agent identity could not be loaded") + if row is None: + return await self.unbound_client(where) + if self.cache is not None: + await self.cache.async_set_cache( + key=cache_key, value=row.agent_id, ttl=get_management_object_ttl(self.cache) + ) + return row.agent_id + + async def resolve_verified_claims( + self, claims: Mapping[str, object] + ) -> ManagedAgentContext | AgentIdentityFailure | None: + issuer: Final = claims.get("iss") + tenant: Final = claims.get("tid") + client: Final = claims.get("azp") + if not isinstance(issuer, str) or not isinstance(tenant, str) or not isinstance(client, str): + return None + agent_id: Final = await self._bound_agent_id(tenant, client) + if agent_id is None or isinstance(agent_id, AgentIdentityFailure): + return agent_id + agent: Final = await self.agent(agent_id) + if isinstance(agent, AgentIdentityFailure): + return agent + if ( + agent is None + or not agent.identity_managed + or not agent.enabled + or agent.identity is None + or not agent.identity.active + ): + return AgentIdentityFailure(message="Agent is disabled or no longer bound to an identity") + subject: Final = classify_agent_subject(agent.identity, claims, agent.execution_mode) + if isinstance(subject, AgentIdentityFailure): + return subject + if subject.kind == "application": + return ManagedAgentContext( + agent_id=agent.agent_id, + binding_revision=agent.identity.revision, + mode=subject.mode, + subject_oid=subject.oid, + ) + proven: Final = await self.subject(issuer, tenant, claims.get("oid")) + if isinstance(proven, AgentIdentityFailure): + return proven + human: Final = ( + VerifiedHumanSubject.model_validate(proven.model_dump()) + if proven is not None + and proven.kind == "human" + and proven.verified_via == "sso_interactive" + and proven.user_id is not None + else None + ) + if human is None: + return AgentIdentityFailure(message="The delegated user must first sign in through trusted Microsoft SSO") + return ManagedAgentContext( + agent_id=agent.agent_id, + binding_revision=agent.identity.revision, + mode=subject.mode, + user_id=human.user_id, + subject_oid=subject.oid, + ) + + async def subject( + self, issuer: str, tenant_id: str, oid: object + ) -> "LiteLLM_VerifiedSubject | AgentIdentityFailure | None": + if not isinstance(oid, str): + return None + try: + where: Final[LiteLLM_VerifiedSubjectWhereUniqueInput] = { + "issuer_tenant_id_oid": {"issuer": issuer, "tenant_id": tenant_id, "oid": oid} + } + return await self.humans.table.find_unique(where=where) + except Exception: + return AgentIdentityFailure(code="policy_unavailable", message="Subject classification is unavailable") + + async def retired_agent(self, agent_id: str) -> bool | AgentIdentityFailure: + if self.retired_agents is None: + return AgentIdentityFailure(code="policy_unavailable", message="Agent history is unavailable") + try: + where: Final[LiteLLM_RetiredAgentWhereUniqueInput] = {"original_agent_id": agent_id} + return await self.retired_agents.table.find_unique(where=where) is not None + except Exception: + return AgentIdentityFailure(code="policy_unavailable", message="Agent history is unavailable") + + async def record_authentication(self, context: ManagedAgentContext) -> AgentIdentityFailure | None: + try: + if context.binding_revision is None: + return AgentIdentityFailure(message="Agent authentication requires a binding revision") + where: Final[LiteLLM_AgentIdentityWhereInput] = { + "agent_id": context.agent_id, + "revision": context.binding_revision, + "active": True, + "agent": {"is": {"enabled": True, "identity_managed": True}}, + } + data: Final[LiteLLM_AgentIdentityUpdateManyMutationInput] = { + "last_authenticated_at": datetime.now(timezone.utc) + } + count: Final = await self.identities.table.update_many(where=where, data=data) + if count != 1: + return AgentIdentityFailure(message="Agent identity changed during authentication; retry") + return None + except Exception: + return AgentIdentityFailure(code="policy_unavailable", message="Agent authentication could not be recorded") + + async def enroll_interactive_human( + self, + subject: MicrosoftInteractiveSubject, + user_id: str, + ) -> AgentIdentityFailure | None: + try: + where: Final[LiteLLM_VerifiedSubjectWhereUniqueInput] = { + "issuer_tenant_id_oid": {"issuer": subject.issuer, "tenant_id": subject.tenant_id, "oid": subject.oid} + } + create_data: Final[LiteLLM_VerifiedSubjectCreateInput] = { + "issuer": subject.issuer, + "tenant_id": subject.tenant_id, + "oid": subject.oid, + "user_id": user_id, + "verified_via": "sso_interactive", + } + data: Final[LiteLLM_VerifiedSubjectUpsertInput] = {"create": create_data, "update": {}} + row: Final = await self.humans.table.upsert(where=where, data=data) + if row.kind != "human" or row.user_id != user_id or row.verified_via != "sso_interactive": + return AgentIdentityFailure(message="Microsoft subject is already bound to another local identity") + return None + except Exception: + return AgentIdentityFailure( + code="policy_unavailable", message="Microsoft subject enrollment is unavailable" + ) + + +async def resolve_managed_agent( + claims: Mapping[str, object], + client: object, + *, + cache: UserApiKeyCache | None = None, +) -> ManagedAgentContext | None: + from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure + + if client is None: + return None + result: Final = await AgentIdentityStore.from_client(client, cache=cache).resolve_verified_claims(claims) + if isinstance(result, AgentIdentityFailure): + raise_identity_failure(result) + return result diff --git a/litellm/proxy/agent_endpoints/managed_identity.py b/litellm/proxy/agent_endpoints/managed_identity.py new file mode 100644 index 00000000000..260b74fcbd1 --- /dev/null +++ b/litellm/proxy/agent_endpoints/managed_identity.py @@ -0,0 +1,220 @@ +from collections.abc import Mapping +from datetime import datetime +from typing import Final, NoReturn, TypedDict +from uuid import uuid4 + +from fastapi import HTTPException +from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly + +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import ( + AgentExecutionMode, + AgentIdentityBinding, + AgentIdentityFailure, + AgentSubject, + EntraIdentityConfig, +) + +_MODE: Final = TypeAdapter(AgentExecutionMode) + + +class IdentityFields(TypedDict, total=False): + provider: ReadOnly[str] + tenant_id: ReadOnly[str] + client_id: ReadOnly[str] + issuer: ReadOnly[str] + service_principal_id: ReadOnly[str | None] + required_roles: ReadOnly[tuple[str, ...]] + required_scopes: ReadOnly[tuple[str, ...]] + active: ReadOnly[bool] + revision: ReadOnly[str] + last_authenticated_at: ReadOnly[datetime | None] + + +class IdentityUpsert(TypedDict): + create: ReadOnly[IdentityFields] + update: ReadOnly[IdentityFields] + + +class IdentityRelationWrite(TypedDict, total=False): + create: ReadOnly[IdentityFields] + update: ReadOnly[IdentityFields] + upsert: ReadOnly[IdentityUpsert] + + +class IdentityHistoryKey(TypedDict): + provider: ReadOnly[str] + tenant_id: ReadOnly[str] + client_id: ReadOnly[str] + + +class IdentityHistoryWhere(TypedDict): + provider_tenant_id_client_id: ReadOnly[IdentityHistoryKey] + + +class IdentityHistoryEntry(IdentityHistoryKey): + issuer: ReadOnly[str] + + +class IdentityHistoryConnect(TypedDict): + where: ReadOnly[IdentityHistoryWhere] + create: ReadOnly[IdentityHistoryEntry] + + +class IdentityHistoryWrite(TypedDict): + connectOrCreate: ReadOnly[IdentityHistoryConnect] + + +class ManagedWriteFields(TypedDict, total=False): + enabled: ReadOnly[bool] + execution_mode: ReadOnly[AgentExecutionMode] + identity_managed: ReadOnly[bool] + identity: ReadOnly[IdentityRelationWrite] + retired_identities: ReadOnly[IdentityHistoryWrite] + + +def raise_identity_failure(failure: AgentIdentityFailure, status_code: int = 403) -> NoReturn: + raise HTTPException(503 if failure.code == "policy_unavailable" else status_code, failure.message) + + +def _configuration_failure( + identity: EntraIdentityConfig | AgentIdentityBinding | None, + mode: AgentExecutionMode, + enabling_without_binding: bool, +) -> AgentIdentityFailure | None: + if identity is not None and mode != "delegated" and not identity.service_principal_id: + return AgentIdentityFailure( + message="Autonomous mode requires the Enterprise application service-principal object ID" + ) + if enabling_without_binding and ( + identity is None or isinstance(identity, AgentIdentityBinding) and not identity.active + ): + return AgentIdentityFailure(message="Bind an identity before enabling this managed agent") + return None + + +def managed_write_fields( + incoming: Mapping[str, object], + existing: AgentResponse | None, + updated_by: str, +) -> ManagedWriteFields | AgentIdentityFailure: + try: + identity: Final = ( + EntraIdentityConfig.model_validate(incoming["identity"]) if incoming.get("identity") is not None else None + ) + mode: Final = _MODE.validate_python( + incoming.get("execution_mode", existing.execution_mode if existing else "autonomous") + ) + current_identity: Final = identity if "identity" in incoming else existing.identity if existing else None + failure: Final = _configuration_failure( + current_identity, + mode, + incoming.get("enabled") is True + and "identity" not in incoming + and bool(existing and existing.identity_managed), + ) + if failure is not None: + return failure + empty: Final[ManagedWriteFields] = {} + identity_fields: Final = _identity_write(identity, existing) if "identity" in incoming else empty + result: Final[ManagedWriteFields] = { + **({"enabled": incoming["enabled"] is True} if "enabled" in incoming else {}), + **({"execution_mode": mode} if "execution_mode" in incoming else {}), + **identity_fields, + } + return result + except (ValidationError, ValueError) as exc: + return AgentIdentityFailure(message=f"Invalid agent identity configuration: {exc}") + + +def _identity_write(identity: EntraIdentityConfig | None, existing: AgentResponse | None) -> ManagedWriteFields: + if identity is None: + unbind: Final[ManagedWriteFields] = { + **( + {"identity": {"update": {"active": False, "revision": str(uuid4()), "last_authenticated_at": None}}} + if existing and existing.identity + else {} + ), + **({"identity_managed": True, "enabled": False} if existing and existing.identity_managed else {}), + } + return unbind + if ( + existing + and existing.identity + and existing.identity.active + and all(getattr(existing.identity, name) == value for name, value in identity.model_dump().items()) + ): + unchanged: Final[ManagedWriteFields] = {} + return unchanged + binding: Final[IdentityFields] = { + "provider": identity.provider, + "tenant_id": identity.tenant_id, + "client_id": identity.client_id, + "service_principal_id": identity.service_principal_id, + "required_roles": identity.required_roles, + "required_scopes": identity.required_scopes, + "issuer": identity.issuer, + "active": True, + "revision": str(uuid4()), + "last_authenticated_at": None, + } + result: Final[ManagedWriteFields] = { + "retired_identities": { + "connectOrCreate": { + "where": { + "provider_tenant_id_client_id": { + "provider": identity.provider, + "tenant_id": identity.tenant_id, + "client_id": identity.client_id, + } + }, + "create": { + "provider": identity.provider, + "issuer": identity.issuer, + "tenant_id": identity.tenant_id, + "client_id": identity.client_id, + }, + } + }, + "identity_managed": True, + "identity": {"upsert": {"create": binding, "update": binding}} if existing else {"create": binding}, + } + return result + + +def classify_agent_subject( + binding: AgentIdentityBinding, + claims: Mapping[str, object], + allowed_mode: AgentExecutionMode, +) -> AgentSubject | AgentIdentityFailure: + if (claims.get("iss"), claims.get("tid"), claims.get("azp")) != ( + binding.issuer, + binding.tenant_id, + binding.client_id, + ): + return AgentIdentityFailure(message="Token does not match the registered Entra application") + oid: Final = claims.get("oid") + if not isinstance(oid, str) or not oid: + return AgentIdentityFailure(message="Entra token must identify its object subject") + scope: Final = claims.get("scp") + facets: Final = claims.get("xms_sub_fct") + if facets is not None and (not isinstance(facets, str) or "13" in facets.split()): + return AgentIdentityFailure(message="Native agent-user authentication is not supported by this binding") + if scope is not None and not isinstance(scope, str): + return AgentIdentityFailure(message="Invalid delegated scope claim") + if isinstance(scope, str) and scope: + if allowed_mode == "autonomous" or oid == binding.service_principal_id or claims.get("idtyp") == "app": + return AgentIdentityFailure(message="Delegated token contradicts the configured agent identity or mode") + granted_scopes: Final = frozenset(scope.split()) + if not granted_scopes or not frozenset(binding.required_scopes).issubset(granted_scopes): + return AgentIdentityFailure(message="Token lacks the required delegated scopes") + return AgentSubject(kind="delegated_subject", oid=oid, mode="delegated") + if allowed_mode == "delegated" or oid != binding.service_principal_id or claims.get("idtyp") == "user": + return AgentIdentityFailure(message="Application token contradicts the configured agent identity or mode") + roles: Final = claims.get("roles", ()) + if not isinstance(roles, (list, tuple)) or any(not isinstance(role, str) for role in roles): + return AgentIdentityFailure(message="Invalid application roles claim") + if not frozenset(binding.required_roles).issubset(roles): + return AgentIdentityFailure(message="Token lacks the required application roles") + return AgentSubject(kind="application", oid=oid, mode="autonomous") diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8fbeaf18460..3ec430332ee 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -14,6 +14,7 @@ import math import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, Protocol, TypeAlias @@ -23,7 +24,7 @@ from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, Unpack import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching.dual_cache import LimitedSizeOrderedDict +from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict from litellm.constants import ( CLI_JWT_EXPIRATION_HOURS, CLI_SESSION_KEY_PREFIX, @@ -77,6 +78,7 @@ from litellm.proxy.agent_endpoints.auth.agent_caller import ( load_agent_caller_team, load_agent_caller_user, ) +from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy from litellm.proxy.auth.budget_throttle import ( budget_throttle_percentage, should_throttle_budget_exceeded, @@ -1057,6 +1059,20 @@ async def common_checks( code=status.HTTP_400_BAD_REQUEST, ) + managed_policy: Final = managed_agent_policy(valid_token) + if _model and valid_token is not None and managed_policy is not None: + managed_models: Final = (managed_policy.object_permission or MappingProxyType({})).get("models", ()) + if not isinstance(managed_models, (list, tuple)) or not managed_models: + raise HTTPException(403, "This agent has no model grants") + _can_object_call_model( + model=_resolve_team_alias(_model, valid_token.team_model_aliases, valid_token.team_id, llm_router), + llm_router=llm_router, + models=list(managed_models), + team_id=valid_token.team_id, + object_type="agent", + key_model_aliases=key_model_aliases_for_auth_check(valid_token), + ) + await _check_agent_access_group_model_access(model=_model, valid_token=valid_token, llm_router=llm_router) await _check_agent_caller_model_access( model=_model, @@ -1784,11 +1800,12 @@ async def _load_bounded_registry( if not isinstance(cached, _RegistryNotCached): return cached + waited_for_another_load: Final = load_lock.locked() async with load_lock: - # The request that held the lock has since cached an answer for everyone waiting on it. - cached_after_wait: Final = await _cached_registry(cache_key, overflow_sentinel, user_api_key_cache) - if not isinstance(cached_after_wait, _RegistryNotCached): - return cached_after_wait + if waited_for_another_load: + cached_after_wait: Final = await _cached_registry(cache_key, overflow_sentinel, user_api_key_cache) + if not isinstance(cached_after_wait, _RegistryNotCached): + return cached_after_wait return await _fetch_and_cache_registry( cache_key=cache_key, @@ -2641,7 +2658,7 @@ async def get_user_object( ) if should_check_db: - response = await _user_table(UserRepository(prisma_client)).find_unique( + response = await _user_table(UserRepository(prisma_client, use_writer=bool(check_db_only))).find_unique( where={"user_id": user_id}, include={"organization_memberships": True} ) @@ -2679,7 +2696,7 @@ async def get_user_object( budget_duration=new_user_params["budget_duration"] ) - response = await _user_table(UserRepository(prisma_client)).create( + response = await _user_table(UserRepository(prisma_client, use_writer=bool(check_db_only))).create( data=new_user_params, include={"organization_memberships": True}, ) @@ -2782,17 +2799,12 @@ async def _cache_team_object( team_table.last_refreshed_at = time.time() key: Final = f"team_id:{team_id}" + usage_cache: Final = None if proxy_logging_obj is None else proxy_logging_obj.internal_usage_cache.dual_cache + # On a shared Redis the write below replaces the team entry and the alias DEL below removes the alias entry + # for both caches, so the usage cache only has its own memory to clear. + redis_shared: Final = usage_cache is not None and usage_cache.redis_cache is user_api_key_cache.redis_cache - if proxy_logging_obj is not None: - try: - await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) - except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not fail the write - verbose_proxy_logger.warning( - "Failed to invalidate internal usage cache entry %s; " - "a stale team object may be served until its TTL expires: %s", - key, - e, - ) + await _invalidate_usage_cache_entry(usage_cache, key, redis_shared=redis_shared, stale="team object") # team_id is the table primary key — guaranteed unique, safe to write. await _cache_management_object( @@ -2819,9 +2831,11 @@ async def _cache_team_object( if team_table.team_alias: alias_key: Final = f"team_alias:{team_table.team_alias}" try: - user_api_key_cache.delete_cache(key=alias_key) - if proxy_logging_obj is not None: - await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=alias_key) + pipelined_delete: Final = await user_api_key_cache.async_delete_cache_pre_call(alias_key) + if pipelined_delete is None: + await user_api_key_cache.async_delete_cache(key=alias_key) + else: + await pipelined_delete except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not fail the mutation verbose_proxy_logger.warning( "Failed to invalidate cached team alias entry %s; " @@ -2829,6 +2843,30 @@ async def _cache_team_object( alias_key, e, ) + await _invalidate_usage_cache_entry(usage_cache, alias_key, redis_shared=redis_shared, stale="team alias") + + +async def _invalidate_usage_cache_entry( + usage_cache: DualCache | None, + key: str, + *, + redis_shared: bool, + stale: str, +) -> None: + if usage_cache is None: + return + try: + if redis_shared: + usage_cache.in_memory_cache.delete_cache(key) + else: + await usage_cache.async_delete_cache(key=key) + except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not fail the write + verbose_proxy_logger.warning( + "Failed to invalidate internal usage cache entry %s; a stale %s may be served until its TTL expires: %s", + key.replace("\r", "").replace("\n", ""), + stale, + e, + ) async def invalidate_team_member_spend_state( @@ -3104,9 +3142,9 @@ class TeamNotFoundError(HTTPException): @log_db_metrics async def _get_team_db_check( - team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None + team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None, *, use_writer: bool = False ) -> "_PrismaTeamRow | None": - response = await _team_table(TeamRepository(prisma_client)).find_unique( + response = await _team_table(TeamRepository(prisma_client, use_writer=use_writer)).find_unique( where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS ) @@ -3140,6 +3178,7 @@ async def _get_team_object_from_user_api_key_cache( proxy_logging_obj: ProxyLogging | None, key: str, team_id_upsert: bool | None = None, + use_writer: bool = False, ) -> LiteLLM_TeamTableCachedObj: db_access_time_key: Final = key should_check_db: Final = _should_check_db( @@ -3148,7 +3187,9 @@ async def _get_team_object_from_user_api_key_cache( db_cache_expiry=db_cache_expiry, ) if should_check_db: - response = await _get_team_db_check(team_id=team_id, prisma_client=prisma_client, team_id_upsert=team_id_upsert) + response = await _get_team_db_check( + team_id=team_id, prisma_client=prisma_client, team_id_upsert=team_id_upsert, use_writer=use_writer + ) # The database answered and the row is not there. Distinct from every # other failure here, which leaves the team's grant unknown. if response is None: @@ -3170,8 +3211,11 @@ async def _get_team_object_from_user_api_key_cache( user_api_key_cache=user_api_key_cache, parent_otel_span=None, proxy_logging_obj=proxy_logging_obj, + check_db_only=use_writer, ) except Exception as e: + if use_writer: + raise verbose_proxy_logger.debug( "Failed to load object_permission for team %s with object_permission_id=%s: %s", team_id, @@ -3261,6 +3305,7 @@ async def get_team_object( db_cache_expiry=db_cache_expiry, key=key, team_id_upsert=team_id_upsert, + use_writer=bool(check_db_only), ) except TeamNotFoundError: raise @@ -3306,16 +3351,15 @@ async def get_access_object( prisma_client: DatabaseClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging | None = None, + *, + check_db_only: bool = False, ) -> LiteLLM_AccessGroupTable: """ - Check if access_group_id in proxy AccessGroupTable - - Always checks cache first, then DB only when not found in cache + - Checks cache first unless authoritative writer admission is requested - if valid, return LiteLLM_AccessGroupTable object - if not, then raise an error - Unlike get_team_object, this has no check_cache_only or check_db_only flags; - it always follows cache-first-then-db semantics. - Raises: - HTTPException: If access group doesn't exist in db or cache (status_code=404) """ @@ -3324,18 +3368,19 @@ async def get_access_object( key: Final = f"access_group_id:{access_group_id}" - cached_access_obj: Final = await user_api_key_cache.async_get_cache( - key=key, - model_type=LiteLLM_AccessGroupTable, + cached_access_obj: Final = ( + None + if check_db_only + else await user_api_key_cache.async_get_cache(key=key, model_type=LiteLLM_AccessGroupTable) ) if cached_access_obj is not None: return cached_access_obj # Not in cache - fetch from DB try: - response: Final = await _dictable_table(AccessGroupRepository(prisma_client), "access_group").find_unique( - where={"access_group_id": access_group_id} - ) + response: Final = await _dictable_table( + AccessGroupRepository(prisma_client, use_writer=check_db_only), "access_group" + ).find_unique(where={"access_group_id": access_group_id}) if response is None: raise HTTPException( @@ -3362,8 +3407,12 @@ async def get_access_object( access_group_id, ) raise HTTPException( - status_code=404, - detail={"error": f"Access group doesn't exist in db. Access group={access_group_id}. Error: {e}"}, + status_code=503 if check_db_only else 404, + detail=( + "Access group policy is unavailable" + if check_db_only + else {"error": f"Access group doesn't exist in db. Access group={access_group_id}. Error: {e}"} + ), ) @@ -3565,13 +3614,16 @@ async def get_org_object_by_alias( ) +LITELLM_SESSION_TOKEN_PREFIX: Final = "litellm_login_" + + class ExperimentalUIJWTToken: @staticmethod def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str: from datetime import timedelta from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - encrypt_value_helper, + encrypt_bearer_token, ) if user_info.user_role is None: @@ -3597,7 +3649,7 @@ class ExperimentalUIJWTToken: user_role=LitellmUserRoles(user_info.user_role), ) - return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True)) + return encrypt_bearer_token(valid_token.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX) @staticmethod def get_cli_jwt_auth_token( @@ -3628,7 +3680,7 @@ class ExperimentalUIJWTToken: from datetime import timedelta from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - encrypt_value_helper, + encrypt_bearer_token, ) if user_info.user_role is None: @@ -3666,7 +3718,7 @@ class ExperimentalUIJWTToken: is_session_token=True, ) - return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True)) + return encrypt_bearer_token(valid_token.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX) @staticmethod def get_key_object_from_ui_hash_key( @@ -3676,10 +3728,10 @@ class ExperimentalUIJWTToken: from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - decrypt_value_helper, + decrypt_bearer_token, ) - decrypted_token: Final = decrypt_value_helper(hashed_token, key="ui_hash_key", exception_type="debug") + decrypted_token: Final = decrypt_bearer_token(hashed_token, prefix=LITELLM_SESSION_TOKEN_PREFIX) if decrypted_token is None: return None try: @@ -3694,6 +3746,8 @@ async def _fetch_key_object_from_db_with_reconnect( parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging | None, deadline_seconds: float | None = None, + *, + check_db_only: bool = False, ) -> BaseModel | None: """ Fetch key object from DB and retry once if a DB connection error can be healed. @@ -3707,6 +3761,7 @@ async def _fetch_key_object_from_db_with_reconnect( prisma_client=prisma_client, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ), name="key", deadline_seconds=deadline_seconds, @@ -3718,10 +3773,13 @@ async def _fetch_key_object_from_db_unbounded( prisma_client: PrismaClient, parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, ) -> BaseModel | None: + fetch: Final = partial(prisma_client.get_data, use_writer=True) if check_db_only else prisma_client.get_data async with db_lookup_gate.current(): try: - return await prisma_client.get_data( + return await fetch( token=hashed_token, table_name="combined_view", parent_otel_span=parent_otel_span, @@ -3743,7 +3801,7 @@ async def _fetch_key_object_from_db_unbounded( lock_timeout_seconds=auth_reconnect_lock_timeout, ) if did_reconnect: - return await prisma_client.get_data( + return await fetch( token=hashed_token, table_name="combined_view", parent_otel_span=parent_otel_span, @@ -3831,6 +3889,8 @@ async def get_key_object( parent_otel_span: Span | None = None, proxy_logging_obj: ProxyLogging | None = None, check_cache_only: bool | None = None, + *, + check_db_only: bool = False, ) -> UserAPIKeyAuth: """ - Check if team id in proxy Team Table @@ -3845,9 +3905,8 @@ async def get_key_object( # Same flow as before: use cache only when we have a hit we can turn into UserAPIKeyAuth # (dict from Redis / model_dump, or UserAPIKeyAuth from in-memory). Otherwise fall through to DB. - user_api_key_auth: Final = await user_api_key_cache.async_get_cache( - key=key, - model_type=UserAPIKeyAuth, + user_api_key_auth: Final = ( + None if check_db_only else await user_api_key_cache.async_get_cache(key=key, model_type=UserAPIKeyAuth) ) if user_api_key_auth is not None: return _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_auth) @@ -3861,6 +3920,7 @@ async def get_key_object( prisma_client=prisma_client, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ) if _valid_token is None: @@ -3874,7 +3934,7 @@ async def get_key_object( _response: Final = UserAPIKeyAuth.model_validate(_valid_token.model_dump(exclude_none=True)) # Load object_permission if object_permission_id exists but object_permission is not loaded - if _response.object_permission_id and not _response.object_permission: + if _response.object_permission_id and (check_db_only or not _response.object_permission): try: _response.object_permission = await get_object_permission( object_permission_id=_response.object_permission_id, @@ -3882,14 +3942,20 @@ async def get_key_object( user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ) except Exception as e: + if check_db_only: + raise verbose_proxy_logger.debug( "Failed to load object_permission for key with object_permission_id=%s: %s", _response.object_permission_id, e, ) + if check_db_only: + return _response + # save the key object to cache await _cache_key_object( hashed_token=hashed_token, @@ -3919,6 +3985,7 @@ async def get_object_permission( user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None = None, proxy_logging_obj: ProxyLogging | None = None, + check_db_only: bool = False, ) -> LiteLLM_ObjectPermissionTable | None: """ - Check if object permission id in proxy ObjectPermissionTable @@ -3930,9 +3997,13 @@ async def get_object_permission( # check if in cache key: Final = object_permission_cache_key(object_permission_id) - deserialized_perm: Final = await user_api_key_cache.async_get_cache( - key=key, - model_type=LiteLLM_ObjectPermissionTable, + deserialized_perm: Final = ( + None + if check_db_only + else await user_api_key_cache.async_get_cache( + key=key, + model_type=LiteLLM_ObjectPermissionTable, + ) ) if deserialized_perm is not None: return deserialized_perm @@ -3940,10 +4011,12 @@ async def get_object_permission( # else, check db try: response: Final = await _dictable_table( - ObjectPermissionRepository(prisma_client), "object_permission" + ObjectPermissionRepository(prisma_client, use_writer=check_db_only), "object_permission" ).find_unique(where={"object_permission_id": object_permission_id}) if response is None: + if check_db_only: + raise HTTPException(status_code=403, detail="Referenced object permission does not exist") return None _perm_obj: Final = LiteLLM_ObjectPermissionTable.model_validate(response.dict()) @@ -3956,6 +4029,8 @@ async def get_object_permission( return _perm_obj except Exception: + if check_db_only: + raise return None @@ -4165,6 +4240,7 @@ async def _get_resources_from_access_groups( prisma_client: DatabaseClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, + check_db_only: bool = False, ) -> list[str]: """ Fetch access groups by their IDs (from cache or DB) and collect @@ -4207,9 +4283,12 @@ async def _get_resources_from_access_groups( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ) resources.extend(getattr(ag, resource_field, [])) except Exception: + if check_db_only: + raise verbose_proxy_logger.debug( "Could not fetch access group %s for resource field %s", ag_id, @@ -4242,6 +4321,7 @@ async def _get_mcp_server_ids_from_access_groups( prisma_client: PrismaClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, + check_db_only: bool = False, ) -> list[str]: """ Collect MCP server IDs from unified access groups. @@ -4253,6 +4333,7 @@ async def _get_mcp_server_ids_from_access_groups( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ) @@ -4261,6 +4342,7 @@ async def _get_agent_ids_from_access_groups( prisma_client: PrismaClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, + check_db_only: bool = False, ) -> list[str]: """ Collect agent IDs from unified access groups. @@ -4272,6 +4354,7 @@ async def _get_agent_ids_from_access_groups( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, ) @@ -4471,26 +4554,37 @@ async def _check_agent_access_group_model_access( """Attached groups naming no model deny every model; the empty allowlist in ``_can_object_call_model`` allows.""" if not model or valid_token is None or not valid_token.agent_id: return True - ceiling: Final = await resolve_ceiling(valid_token.agent_id) - if ceiling is None: - return True - if not ceiling.models: - raise ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=model), - internal_message=f"agent {valid_token.agent_id} access groups {ceiling.access_group_ids} grant no models", - type=ProxyErrorTypes.agent_model_access_denied, - param="model", - code=status.HTTP_403_FORBIDDEN, - ) - dispatched: Final = _resolve_team_alias(model, valid_token.team_model_aliases, valid_token.team_id, llm_router) - return _can_object_call_model( - model=dispatched, - llm_router=llm_router, - models=sorted(ceiling.models), - team_id=valid_token.team_id, - object_type="agent", - key_model_aliases=key_model_aliases_for_auth_check(valid_token), + + from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_managed_agent_ceilings + + managed: Final = managed_agent_policy(valid_token) + unmanaged: Final = await resolve_ceiling(valid_token.agent_id) if managed is None else None + ceilings: Final = ( + await resolve_managed_agent_ceilings(managed) + if managed is not None + else (unmanaged,) + if unmanaged is not None + else () ) + dispatched: Final = _resolve_team_alias(model, valid_token.team_model_aliases, valid_token.team_id, llm_router) + for ceiling in ceilings: + if not ceiling.models: + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=model), + internal_message=f"agent {valid_token.agent_id} access groups grant no models", + type=ProxyErrorTypes.agent_model_access_denied, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + _can_object_call_model( + model=dispatched, + llm_router=llm_router, + models=sorted(ceiling.models), + team_id=valid_token.team_id, + object_type="agent", + key_model_aliases=key_model_aliases_for_auth_check(valid_token), + ) + return True LoadedCallerTeam: TypeAlias = LiteLLM_TeamTable | None diff --git a/litellm/proxy/auth/auth_object_prefetch.py b/litellm/proxy/auth/auth_object_prefetch.py index ce55190aa02..14d3e2c07dc 100644 --- a/litellm/proxy/auth/auth_object_prefetch.py +++ b/litellm/proxy/auth/auth_object_prefetch.py @@ -15,7 +15,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.caching.redis_batch import active_request_redis_batch from litellm.caching.redis_cache import RedisCache -from litellm.constants import DEFAULT_IN_MEMORY_TTL +from litellm.constants import DEFAULT_IN_MEMORY_TTL, REGISTRY_ERROR_NEGATIVE_CACHE_TTL from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.models.team_membership import LiteLLM_TeamMembership @@ -324,3 +324,35 @@ async def prefetch_auth_objects( await _fill_from_db(refs, _missing_in_memory(missing, memory), user_api_key_cache, prisma_client) except Exception as e: # noqa: BLE001 # warm-up only; the getters enforce and fail closed on their own verbose_proxy_logger.warning("auth prefetch skipped, falling back to per-object lookups: %s", e) + + +def _identity_memory_ttl(value: object, management_ttl: float) -> float: + """A registry stored as a string is a sentinel, written with the shorter of the two registry TTLs.""" + return min(REGISTRY_ERROR_NEGATIVE_CACHE_TTL, management_ttl) if isinstance(value, str) else management_ttl + + +async def prefetch_identity_keys(cache_keys: Sequence[str], user_api_key_cache: UserApiKeyCache) -> None: + """Warm the entries auth reads before it knows the key's owners (the key object, the end user and the two + registries) in one MGET on the request pipeline. Keys the MGET finds absent stay noted on the pipeline, so the + per-key getters that follow go to the database without a GET of their own. Best effort, like the + owner prefetch: the getters read and enforce on their own.""" + try: + redis_cache: Final = user_api_key_cache.redis_cache + if redis_cache is None: + return + missing: Final = tuple( + key + for key in dict.fromkeys(cache_keys) + if user_api_key_cache.in_memory_cache_for(key).get_cache(key=key) is None + ) + if not missing: + return + found: Final = _RowValues.validate_python(await _read_redis_rows(sorted(missing), redis_cache)) + management_ttl: Final = get_management_object_ttl(user_api_key_cache) + except Exception as e: # noqa: BLE001 # warm-up only; the getters read Redis and the database on their own + verbose_proxy_logger.warning("auth identity prefetch skipped, falling back to per-key lookups: %s", e) + return + for key, value in ((key, found.get(key)) for key in missing): + if value is not None: + memory: _InMemoryCache = user_api_key_cache.in_memory_cache_for(key) + _set_in_memory(memory, key, value, _identity_memory_ttl(value, management_ttl)) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ed4d63fb9fc..36c6a4c476b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -73,7 +73,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler from litellm.proxy.auth.auth_method import AuthMethod -from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects +from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects, prefetch_identity_keys from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, get_end_user_id_from_request_body, @@ -120,6 +120,9 @@ from litellm.proxy.common_utils.model_listing_utils import claude_code_requested from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, + end_user_cache_key, + end_user_restricted_registry_cache_key, + model_access_group_registry_cache_key, team_membership_auth_cache_key, ) from litellm.proxy.db.db_lookup_gate import bounded_db_lookup @@ -1892,6 +1895,11 @@ async def _user_api_key_auth_builder( proxy_logging_obj=proxy_logging_obj, route=route, ) + if prisma_client is not None: + await prefetch_identity_keys( + _identity_cache_keys(api_key, end_user_id=end_user_id, key_is_resolved=valid_token is not None), + user_api_key_cache=user_api_key_cache, + ) if end_user_id: try: end_user_params["end_user_id"] = end_user_id @@ -3196,6 +3204,15 @@ async def _authorize_authenticated_request( # admin-only-route / model-access / budget checks) surface as # ProxyException consistently with pre-refactor behavior. try: + from litellm.proxy.agent_endpoints.auth.managed_authorization import admit_managed_actor + from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore + from litellm.proxy.proxy_server import prisma_client + + if user_api_key_auth_obj.agent_id is not None: + await admit_managed_actor( + user_api_key_auth_obj, + AgentIdentityStore.from_client(prisma_client) if prisma_client is not None else None, + ) await _run_centralized_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, request=request, @@ -3248,6 +3265,21 @@ def _spend_counter_redis_cache() -> RedisCache | None: return spend_counter_cache.redis_cache +def _identity_cache_keys(api_key: str, *, end_user_id: str | None, key_is_resolved: bool) -> tuple[str, ...]: + """Cache keys auth reads before it knows the key's owners, all known from the request alone. A key object is + cached under the hash of the bearer, so the bearer itself never reaches Redis.""" + return tuple( + key + for key in ( + None if key_is_resolved else hash_token(api_key), + None if not end_user_id else end_user_cache_key(end_user_id), + None if not end_user_id else end_user_restricted_registry_cache_key(), + model_access_group_registry_cache_key(), + ) + if key is not None + ) + + async def _prefetch_referenced_auth_objects( valid_token: UserAPIKeyAuth, end_user_id: str | None, diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 3584aaaf833..ae7240b8a7f 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -72,26 +72,55 @@ def _derive_key(signing_key: str) -> bytes: return hashlib.sha256(signing_key.encode()).digest() -def _encrypt_aes_gcm(value: str, signing_key: str) -> str: - """Encrypt under AES-256-GCM and return the versioned ``v2:gcm:`` string.""" +def _seal_aes_gcm(value: str, signing_key: str, aad: bytes | None) -> bytes: from cryptography.hazmat.primitives.ciphers.aead import AESGCM nonce: Final = os.urandom(12) # AESGCM.encrypt returns ciphertext || tag(16); wire format is nonce || that. - blob: Final = AESGCM(_derive_key(signing_key)).encrypt(nonce, value.encode("utf-8"), None) - return _V2_GCM_PREFIX + base64.urlsafe_b64encode(nonce + blob).decode("utf-8") + return nonce + AESGCM(_derive_key(signing_key)).encrypt(nonce, value.encode("utf-8"), aad) + + +def _open_aes_gcm(sealed: bytes, signing_key: str, aad: bytes | None) -> str: + from cryptography.hazmat.primitives.ciphers.aead import AESGCM + + # An empty plaintext still serializes to nonce(12) || tag(16) = 28 bytes, so a + # short/empty buffer here is a corrupt value: let AESGCM.decrypt raise and be + # swallowed by the caller (returns None/original), same as legacy. + return AESGCM(_derive_key(signing_key)).decrypt(sealed[:12], sealed[12:], aad).decode("utf-8") + + +def _encrypt_aes_gcm(value: str, signing_key: str) -> str: + """Encrypt under AES-256-GCM and return the versioned ``v2:gcm:`` string.""" + sealed: Final = _seal_aes_gcm(value=value, signing_key=signing_key, aad=None) + return _V2_GCM_PREFIX + base64.urlsafe_b64encode(sealed).decode("utf-8") def _decrypt_aes_gcm(value: str, signing_key: str) -> str: """Decrypt a versioned ``v2:gcm:`` string produced by :func:`_encrypt_aes_gcm`.""" - from cryptography.hazmat.primitives.ciphers.aead import AESGCM + sealed: Final = base64.urlsafe_b64decode(value[len(_V2_GCM_PREFIX) :]) + return _open_aes_gcm(sealed=sealed, signing_key=signing_key, aad=None) - raw: Final = base64.urlsafe_b64decode(value[len(_V2_GCM_PREFIX) :]) - # An empty plaintext still serializes to nonce(12) || tag(16) = 28 bytes, so a - # short/empty buffer here is a corrupt value: let AESGCM.decrypt raise and be - # swallowed by decrypt_value_helper (returns None/original), same as legacy. - nonce, blob = raw[:12], raw[12:] - return AESGCM(_derive_key(signing_key)).decrypt(nonce, blob, None).decode("utf-8") + +def encrypt_bearer_token(value: str, prefix: str) -> str: + """AES-256-GCM as unpadded base64url behind ``prefix``, which is also the AAD so a token can't change kind.""" + salt_key: Final = _get_salt_key() + if not isinstance(salt_key, str): + raise ValueError("Set LITELLM_SALT_KEY or a master key to mint bearer tokens") + sealed: Final = _seal_aes_gcm(value=value, signing_key=salt_key, aad=prefix.encode("utf-8")) + return prefix + base64.urlsafe_b64encode(sealed).decode("ascii").rstrip("=") + + +def decrypt_bearer_token(token: str, prefix: str) -> str | None: + """None unless ``token`` came from :func:`encrypt_bearer_token` with the same ``prefix``.""" + salt_key: Final = _get_salt_key() + if not isinstance(salt_key, str) or not token.startswith(prefix): + return None + encoded: Final = token.removeprefix(prefix) + try: + sealed: Final = base64.b64decode(encoded + "=" * (-len(encoded) % 4), altchars=b"-_", validate=True) + return _open_aes_gcm(sealed=sealed, signing_key=salt_key, aad=prefix.encode("utf-8")) + except Exception: # noqa: BLE001 # base64 and AES-GCM each raise their own "not a token" type + return None def encrypt_value_helper(value: str, new_encryption_key: str | None = None): diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 61d7078ae4c..c99665986dd 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -17,6 +17,8 @@ from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec if TYPE_CHECKING: from opentelemetry.trace import Span + from litellm.caching.redis_batch import BatchResult + T = TypeVar("T", bound=BaseModel) _HASHED_TOKEN_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}") @@ -27,6 +29,9 @@ def is_user_key_cache_key(key: str) -> bool: return _HASHED_TOKEN_CACHE_KEY.fullmatch(key) is not None +_PIPELINED_SET_OPTIONS: Final = frozenset(("ttl",)) + + class UserApiKeyCache(DualCache): """ DualCache wrapper for UserAPIKeyAuth-like payloads. @@ -208,10 +213,23 @@ class UserApiKeyCache(DualCache): return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs) async def async_set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object): + """Inside a request the Redis SET rides the request's pipeline (memory is written at once); anywhere + else, or with options the pipeline does not carry, it goes to Redis directly as before.""" model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None)) payload: Final[object] = CacheCodec.serialize(value, model_type=model_type) + ttl: Final = kwargs.get("ttl") + pipelined: Final = ( + key is not None + and not local_only + and kwargs.keys() <= _PIPELINED_SET_OPTIONS + and (ttl is None or isinstance(ttl, (int, float))) + ) if key is not None and is_user_key_cache_key(key): + if pipelined and await self.key_object_cache.async_set_cache_pre_call(key, payload, ttl) is not None: + return None return await self.key_object_cache.async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) + if pipelined and await super().async_set_cache_pre_call(key, payload, ttl) is not None: + return None return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs) def delete_cache(self, key: str) -> None: @@ -226,6 +244,11 @@ class UserApiKeyCache(DualCache): return await super().async_delete_cache(key) + async def async_delete_cache_pre_call(self, key: str) -> BatchResult[None] | None: + if is_user_key_cache_key(key): + return await self.key_object_cache.async_delete_cache_pre_call(key) + return await super().async_delete_cache_pre_call(key) + async def async_delete_cache_keys(self, keys: Sequence[str]) -> None: """Batch twin of ``async_delete_cache``, partitioned like ``async_set_cache_pipeline``. diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7e159ec90e7..d37dfe87ad5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2442,11 +2442,13 @@ async def _update_key_row_with_soft_budget( existing_key_row=existing_key_row, changed_by=changed_by, ) + include_object_permission: Final[prisma.types.LiteLLM_VerificationTokenInclude] = {"object_permission": True} updated_row: Final = await tx.litellm_verificationtoken.update( where=key_where, data=with_settings_updated_at( prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) ), + include=include_object_permission, ) updated_data: Final[Mapping[str, object]] = ( updated_row.model_dump() if updated_row is not None else MappingProxyType({}) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 536c58df65a..42ac74cae33 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -142,6 +142,7 @@ ROUTE_ENDPOINT_MAPPING: Final = { "acancel_run": "/evals/{eval_id}/runs/{run_id}/cancel", "adelete_run": "/evals/{eval_id}/runs/{run_id}", "acreate_batch": "/batches", + "aretrieve_batch": "/batches", } diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 03e59257f76..f29caa9ceb7 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -78,6 +78,11 @@ model LiteLLM_AgentsTable { object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) + identity_managed Boolean @default(false) + enabled Boolean @default(true) + execution_mode String @default("autonomous") + identity LiteLLM_AgentIdentity? + retired_identities LiteLLM_RetiredAgentIdentity[] tpm_limit Int? rpm_limit Int? session_tpm_limit Int? @@ -88,6 +93,56 @@ model LiteLLM_AgentsTable { updated_by String } +model LiteLLM_AgentIdentity { + agent_id String @id + active Boolean @default(true) + agent LiteLLM_AgentsTable @relation(fields: [agent_id], references: [agent_id], onDelete: Cascade) + provider String + issuer String + tenant_id String + client_id String + service_principal_id String? + required_roles String[] @default([]) + required_scopes String[] @default(["user_impersonation"]) + revision String @default(uuid()) + last_authenticated_at DateTime? + @@unique([provider, tenant_id, client_id]) + @@unique([issuer, service_principal_id]) +} + +model LiteLLM_RetiredAgentIdentity { + binding_id String @id @default(uuid()) + agent_id String? + agent LiteLLM_AgentsTable? @relation(fields: [agent_id], references: [agent_id], onDelete: SetNull) + provider String + issuer String + tenant_id String + client_id String + @@unique([provider, tenant_id, client_id]) +} + +model LiteLLM_RetiredAgent { + original_agent_id String @id + retired_at DateTime @default(now()) +} + +model LiteLLM_VerifiedSubject { + subject_id String @id @default(uuid()) + issuer String + tenant_id String + oid String + kind String @default("human") + user_id String? + user LiteLLM_UserTable? @relation(fields: [user_id], references: [user_id], onDelete: Cascade) + verified_via String @default("sso_interactive") + verified_at DateTime @default(now()) + @@unique([issuer, tenant_id, oid]) + @@index([user_id]) +} + + + + model LiteLLM_OrganizationTable { organization_id String @id @default(uuid()) organization_alias String @@ -241,6 +296,7 @@ model LiteLLM_DeletedTeamTable { // Track spend, rate limit, budget Users model LiteLLM_UserTable { + verified_subjects LiteLLM_VerifiedSubject[] user_id String @id user_alias String? team_id String? @@ -675,6 +731,7 @@ model LiteLLM_SpendLogs { session_id String? status String? mcp_namespaced_tool_name String? + billing_agent_id String? agent_id String? proxy_server_request Json? @default("{}") litellm_call_id String? diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 0e3e0598d17..560363ca7d7 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -4,7 +4,7 @@ from collections.abc import Set as AbstractSet from dataclasses import dataclass from datetime import datetime, timedelta from types import MappingProxyType -from typing import Final, TypeVar +from typing import Final, Literal, TypeVar from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -17,6 +17,7 @@ from litellm.constants import ( SPEND_LOG_KEY_METADATA_CACHE_TTL, SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS, + SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE, ) from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.proxy.utils import PrismaClient @@ -39,26 +40,58 @@ WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[]) ORDER BY token, deleted_at DESC """ -_SPEND_LOG_ALIAS_SQL: Final = """ -SELECT api_key AS digest, - MIN(key_alias) AS first_alias, - MAX(key_alias) AS last_alias, - MIN(team_id) AS first_team, - MAX(team_id) AS last_team, - MIN(user_id) AS first_owner, - MAX(user_id) AS last_owner -FROM ( - SELECT api_key, - NULLIF(metadata->>'user_api_key_alias', '') AS key_alias, - COALESCE(NULLIF(team_id, ''), NULLIF(metadata->>'user_api_key_team_id', '')) AS team_id, - COALESCE(NULLIF("user", ''), NULLIF(metadata->>'user_api_key_user_id', '')) AS user_id - FROM "LiteLLM_SpendLogs" - WHERE api_key = ANY($1::text[]) - AND "startTime" >= $2::timestamp - AND "startTime" < $3::timestamp -) named -WHERE COALESCE(key_alias, user_id, team_id) IS NOT NULL -GROUP BY api_key + +def _named_spend_log_edge_row_sql( + direction: Literal["ASC", "DESC"], since: Literal["$2::timestamp", "oldest_probe.stopped_at"] +) -> str: + return f""" + SELECT "startTime", key_alias, team_id, user_id + FROM ( + SELECT "startTime", + NULLIF(metadata->>'user_api_key_alias', '') AS key_alias, + COALESCE(NULLIF(team_id, ''), NULLIF(metadata->>'user_api_key_team_id', '')) AS team_id, + COALESCE(NULLIF("user", ''), NULLIF(metadata->>'user_api_key_user_id', '')) AS user_id + FROM ( + SELECT "startTime", metadata, team_id, "user" + FROM "LiteLLM_SpendLogs" + WHERE api_key = keys.digest + AND "startTime" >= {since} + AND "startTime" < $3::timestamp + ORDER BY "startTime" {direction} + LIMIT {SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE} + ) edge + ) named + WHERE COALESCE(key_alias, user_id, team_id) IS NOT NULL + ORDER BY "startTime" {direction} + LIMIT 1 + """ + + +_OLDEST_PROBE_STOPPED_AT_SQL: Final = f""" + SELECT COALESCE(first_row."startTime", ( + SELECT "startTime" + FROM "LiteLLM_SpendLogs" + WHERE api_key = keys.digest + AND "startTime" >= $2::timestamp + AND "startTime" < $3::timestamp + ORDER BY "startTime" ASC + OFFSET {SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE - 1} + LIMIT 1 + )) AS stopped_at +""" + +_SPEND_LOG_ALIAS_SQL: Final = f""" +SELECT keys.digest, + first_row.key_alias AS first_alias, + last_row.key_alias AS last_alias, + first_row.team_id AS first_team, + last_row.team_id AS last_team, + first_row.user_id AS first_owner, + last_row.user_id AS last_owner +FROM unnest($1::text[]) AS keys(digest) +LEFT JOIN LATERAL ({_named_spend_log_edge_row_sql("ASC", "$2::timestamp")}) first_row ON true +LEFT JOIN LATERAL ({_OLDEST_PROBE_STOPPED_AT_SQL}) oldest_probe ON true +LEFT JOIN LATERAL ({_named_spend_log_edge_row_sql("DESC", "oldest_probe.stopped_at")}) last_row ON true """ _DAILY_USER_SPEND_OWNER_SQL: Final = """ @@ -69,6 +102,7 @@ GROUP BY api_key """ _SPEND_LOG_STATEMENT_TIMEOUT_SQL: Final = f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}" +_SPEND_LOG_NO_BITMAP_SCAN_SQL: Final = "SET LOCAL enable_bitmapscan = off" _SPEND_LOG_TRANSACTION_TIMEOUT: Final = timedelta(milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS) _HASHED_JWT_PREFIX: Final = "hashed-jwt-" @@ -91,7 +125,9 @@ class _TokenDigestRow(BaseModel): def _unanimous(first: str | None, last: str | None) -> str | None: - return first if first == last else None + if first is None: + return last + return first if last is None or first == last else None class _SpendLogDigestRow(BaseModel): @@ -148,9 +184,12 @@ async def _rows_within_the_statement_timeout( prisma_client: PrismaClient, sql: str, *params: object, + planner_settings: tuple[str, ...] = (), ) -> Sequence[Mapping[str, object]]: async with prisma_client.db.tx(timeout=_SPEND_LOG_TRANSACTION_TIMEOUT) as transaction: await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL) + for setting in planner_settings: + await transaction.execute_raw(setting) return await transaction.query_raw(sql, *params) @@ -364,7 +403,14 @@ async def _query_spend_log_metadata( ) -> Mapping[str, KeyMetadataDict] | None: start, end = window rows: Final = await _db_or_empty( - lambda: _rows_within_the_statement_timeout(prisma_client, _SPEND_LOG_ALIAS_SQL, sorted(digests), start, end), + lambda: _rows_within_the_statement_timeout( + prisma_client, + _SPEND_LOG_ALIAS_SQL, + sorted(digests), + start, + end, + planner_settings=(_SPEND_LOG_NO_BITMAP_SCAN_SQL,), + ), "Failed spend-log alias recovery for %d missing keys: %s", len(digests), ) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index c4ed8713f95..728579db5fc 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2444,7 +2444,7 @@ def _build_spend_log_search_condition( f"(request_id = {raw} OR (" f"\"startTime\" >= ({window_start}::timestamptz AT TIME ZONE 'UTC') " f"AND \"startTime\" <= ({window_end}::timestamptz AT TIME ZONE 'UTC') " - f'AND (api_key = {raw} OR team_id = {raw} OR "user" = {raw} OR end_user = {raw} ' + f'AND (litellm_call_id = {raw} OR api_key = {raw} OR team_id = {raw} OR "user" = {raw} OR end_user = {raw} ' f"OR session_id = {raw} OR model_id = {raw})))" ) return _SpendLogSearchCondition(sql=sql, params=(search, start_date, end_date)) @@ -2557,7 +2557,7 @@ async def ui_view_spend_logs( search: str | None = fastapi.Query( default=None, description=( - "Match a log whose request_id, api_key (hash), team_id, user, end_user, " + "Match a log whose request_id, litellm_call_id, api_key (hash), team_id, user, end_user, " "session_id, or model_id equals this value. request_id matches across all time; the other columns " "match inside start_date/end_date, which stay required" ), diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 1c51fb21d6e..e69c0c80420 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -44,6 +44,7 @@ from litellm.litellm_core_utils.litellm_logging import ( ) from litellm.litellm_core_utils.ptu_pricing import azure_spillover from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes +from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload, SpendLogsRouterMetadata from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error @@ -1083,6 +1084,11 @@ def _get_messages_for_spend_logs_payload( _SENSITIVE_REQUEST_BODY_KEYS: Final = frozenset({"secret_fields"}) +_REQUEST_BODY_CREDENTIAL_MASKER: Final = SensitiveDataMasker(extra_sensitive_patterns=frozenset({"apikey"})) + + +def _is_request_body_credential(key: str, value: object) -> bool: + return isinstance(value, str) and _REQUEST_BODY_CREDENTIAL_MASKER.is_sensitive_key(key) def _sanitize_request_body_for_spend_logs_payload( @@ -1094,8 +1100,9 @@ def _sanitize_request_body_for_spend_logs_payload( Recursively sanitize request body to prevent logging large base64 strings or other large values. Truncates strings longer than MAX_STRING_LENGTH_PROMPT_IN_DB characters and handles nested dictionaries. - Also strips keys listed in _SENSITIVE_REQUEST_BODY_KEYS (e.g. secret_fields - which contains raw HTTP headers including Authorization tokens). + At every nesting level, also strips keys listed in _SENSITIVE_REQUEST_BODY_KEYS (e.g. secret_fields, + which holds raw HTTP headers including Authorization tokens), and replaces string values under keys + SensitiveDataMasker classifies as credentials with REDACTED_BY_LITELM_STRING. """ from litellm.constants import ( LITELLM_TRUNCATED_PAYLOAD_FIELD, @@ -1152,7 +1159,11 @@ def _sanitize_request_body_for_spend_logs_payload( return value return value - return {k: _sanitize_value(v) for k, v in request_body.items() if k not in _SENSITIVE_REQUEST_BODY_KEYS} + return { + k: REDACTED_BY_LITELM_STRING if _is_request_body_credential(k, v) else _sanitize_value(v) + for k, v in request_body.items() + if k not in _SENSITIVE_REQUEST_BODY_KEYS + } # Quoted-key form: ``"input"`` / ``'messages'`` / ``"prompt"`` followed by diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f980d1e2428..8a325eba699 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -258,7 +258,7 @@ if TYPE_CHECKING: from prisma.actions import LiteLLM_DeprecatedVerificationTokenActions from prisma.client import TransactionManager from prisma.models import LiteLLM_DeprecatedVerificationToken - from prisma.types import HttpConfig + from prisma.types import HttpConfig, LiteLLM_VerificationTokenInclude from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation @@ -4200,6 +4200,8 @@ _PRISMA_DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5) async def _lookup_deprecated_key( db: PrismaWrapper | RoutingPrismaWrapper, hashed_token: str, + *, + check_db_only: bool = False, ) -> str | None: """ Check if a token exists in the deprecated keys table and is still within its grace period. @@ -4211,7 +4213,7 @@ async def _lookup_deprecated_key( now_ts: Final = now.timestamp() # Check cache first - cached: Final = _deprecated_key_cache.get(hashed_token) + cached: Final = None if check_db_only else _deprecated_key_cache.get(hashed_token) if cached is not None: active_token_id, cache_expires_at_ts, revoke_at_ts = cached if now_ts < cache_expires_at_ts and now_ts < revoke_at_ts: @@ -4879,6 +4881,7 @@ class PrismaClient: proxy_logging_obj: ProxyLogging | None = None, budget_id_list: list[str] | None = None, check_deprecated: bool = True, + use_writer: bool = False, ): args_passed_in: Final = locals() start_time: Final = time.time() @@ -5177,12 +5180,20 @@ class PrismaClient: WHERE v.token = $1 """ - response = await self._query_first_with_cached_plan_fallback(sql_query, hashed_token) + response = ( + await self.writer_db.query_first(sql_query, hashed_token) + if use_writer + else await self._query_first_with_cached_plan_fallback(sql_query, hashed_token) + ) # If not found in main table, check deprecated keys (grace period) # check_deprecated=False on the recursive call prevents unbounded chaining if response is None and hashed_token is not None and check_deprecated: - active_token_id: Final = await _lookup_deprecated_key(db=self.db, hashed_token=hashed_token) + active_token_id: Final = await _lookup_deprecated_key( + db=self.writer_db if use_writer else self.db, + hashed_token=hashed_token, + check_db_only=use_writer, + ) if active_token_id: # The recursive call returns a finished # LiteLLM_VerificationTokenView; the dict @@ -5194,6 +5205,7 @@ class PrismaClient: parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, check_deprecated=False, + use_writer=use_writer, ) if deprecated_response is not None: verbose_proxy_logger.debug("Deprecated key used during grace period") @@ -5445,9 +5457,11 @@ class PrismaClient: # check if plain text or hash token = _hash_token_if_needed(token=token) db_data["token"] = token + include_object_permission: Final[LiteLLM_VerificationTokenInclude] = {"object_permission": True} response: Final = await VerificationTokenRepository(self).table.update( where={"token": token}, data=with_settings_updated_at(db_data), + include=include_object_permission, ) verbose_proxy_logger.debug("\033[91m" + f"DB Token Table update succeeded {response}" + "\033[0m") _data: dict = {} diff --git a/litellm/repositories/object_permission_repository.py b/litellm/repositories/object_permission_repository.py index 7736939c696..b732d2ff94c 100644 --- a/litellm/repositories/object_permission_repository.py +++ b/litellm/repositories/object_permission_repository.py @@ -15,9 +15,14 @@ if TYPE_CHECKING: class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): """Repository for object permission database operations.""" + def __init__(self, prisma_client: object, *, use_writer: bool = False) -> None: + super().__init__(prisma_client) + self._use_writer = use_writer + @property def table(self) -> TableActions["prisma_models.LiteLLM_ObjectPermissionTable"]: - return self.prisma_client.db.litellm_objectpermissiontable + database: Final = self.prisma_client.writer_db if self._use_writer else self.prisma_client.db + return database.litellm_objectpermissiontable @property def model_class(self) -> type[LiteLLM_ObjectPermissionTable]: diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index ab68f1a2bc7..4e511a2ec93 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -21,8 +21,9 @@ class PrismaTableRepository(Generic[RowT_co]): table_name: str - def __init__(self, prisma_client: object): + def __init__(self, prisma_client: object, *, use_writer: bool = False) -> None: self._prisma_client = prisma_client + self._use_writer = use_writer @property def prisma_client(self) -> Any: @@ -32,7 +33,9 @@ class PrismaTableRepository(Generic[RowT_co]): @property def table(self) -> TableActions[RowT_co]: - actions: Final[TableActions[RowT_co]] = getattr(self.prisma_client.db, self.table_name) + actions: Final[TableActions[RowT_co]] = getattr( + self.prisma_client.writer_db if self._use_writer else self.prisma_client.db, self.table_name + ) return wrap_table_actions_for_config_sync(actions=actions, table_name=self.table_name) @@ -44,6 +47,18 @@ class AgentsRepository(PrismaTableRepository["prisma_models.LiteLLM_AgentsTable" table_name = "litellm_agentstable" +class AgentIdentityRepository(PrismaTableRepository["prisma_models.LiteLLM_AgentIdentity"]): + table_name = "litellm_agentidentity" + + +class RetiredAgentIdentityRepository(PrismaTableRepository["prisma_models.LiteLLM_RetiredAgentIdentity"]): + table_name = "litellm_retiredagentidentity" + + +class VerifiedSubjectRepository(PrismaTableRepository["prisma_models.LiteLLM_VerifiedSubject"]): + table_name = "litellm_verifiedsubject" + + class ObjectPermissionRepository(PrismaTableRepository["prisma_models.LiteLLM_ObjectPermissionTable"]): table_name = "litellm_objectpermissiontable" @@ -250,3 +265,7 @@ class AuditLogRepository(PrismaTableRepository["prisma_models.LiteLLM_AuditLog"] class AdaptiveRouterSessionRepository(PrismaTableRepository["prisma_models.LiteLLM_AdaptiveRouterSession"]): table_name = "litellm_adaptiveroutersession" + + +class RetiredAgentRepository(PrismaTableRepository["prisma_models.LiteLLM_RetiredAgent"]): + table_name = "litellm_retiredagent" diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index cbe263699c9..57f6fd33c11 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -70,6 +70,9 @@ class _PrismaClientView(Protocol): @property def db(self) -> _PrismaTeamDb: ... + @property + def writer_db(self) -> _PrismaTeamDb: ... + _MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member]) _JSON_ENCODED_TEAM_FIELDS: Final = ( @@ -85,10 +88,14 @@ _JSON_ENCODED_TEAM_FIELDS: Final = ( class TeamRepository(BaseRepository[LiteLLM_TeamTable]): """Repository for team database operations.""" + def __init__(self, prisma_client: object, *, use_writer: bool = False) -> None: + super().__init__(prisma_client) + self._use_writer = use_writer + @property def _db(self) -> _PrismaTeamDb: client: Final[_PrismaClientView] = self.prisma_client - return client.db + return client.writer_db if self._use_writer else client.db @property def table(self) -> TableActions["prisma_models.LiteLLM_TeamTable"]: diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 87eb45f262d..4a2aea46197 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -38,9 +38,14 @@ _PLACEHOLDER_ROWS_ADAPTER: Final = TypeAdapter(tuple[SCIMPlaceholder, ...]) class UserRepository(BaseRepository[LiteLLM_UserTable]): """Repository for user database operations.""" + def __init__(self, prisma_client: object, *, use_writer: bool = False) -> None: + super().__init__(prisma_client) + self._use_writer = use_writer + @property def table(self) -> TableActions["prisma_models.LiteLLM_UserTable"]: - return self.prisma_client.db.litellm_usertable + database: Final = self.prisma_client.writer_db if self._use_writer else self.prisma_client.db + return database.litellm_usertable @property def model_class(self) -> type[LiteLLM_UserTable]: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 12bc9adbac8..10c73071fc7 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -601,7 +601,7 @@ class BaseResponsesAPIStreamingIterator: raw_headers: Final[Mapping[str, object]] = raw if isinstance(raw, Mapping) else EMPTY_MAPPING # rebuild by value and let existing keys win: sharing the source dicts would alias what the proxy # splats into the client's HTTP headers, and copying non-header keys would carry response_cost - target._hidden_params = { # mutable-ok: the cost calculator writes optional_params into _hidden_params + target._hidden_params = { # mutable-ok: logging aliases _hidden_params into request metadata and writes into it "additional_headers": {**headers}, # mutable-ok: fresh copy, logging callbacks may mutate it "headers": {**raw_headers}, # mutable-ok: fresh copy, logging callbacks may mutate it **existing, diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 9b0d259eb8a..c5e6f3995f7 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1,6 +1,6 @@ import base64 import re -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from functools import reduce from typing import Any, Final, Optional, TypeVar, Union, cast, get_type_hints, overload @@ -556,7 +556,11 @@ class ResponsesAPIRequestUtils: return request_input @staticmethod - def strip_encrypted_reasoning_from_input(request_input: object) -> None: + def strip_encrypted_reasoning_from_input( + request_input: object, + *, + should_strip: Callable[[Mapping[str, object]], bool] | None = None, + ) -> None: """Drop reasoning items the routed deployment cannot decrypt, keeping their readable summary. Mutates ``request_input`` in place: the router's fallback snapshot shares this @@ -565,7 +569,12 @@ class ResponsesAPIRequestUtils: if not isinstance(request_input, list): return items: Final = cast(list[object], request_input) # cast-ok: untyped client json - stripped: Final = tuple(ResponsesAPIRequestUtils._without_encrypted_reasoning(item) for item in items) + stripped: Final = tuple( + ResponsesAPIRequestUtils._without_encrypted_reasoning(item) + if should_strip is None or (isinstance(item, Mapping) and should_strip(cast(Mapping[str, object], item))) + else item + for item in items + ) items[:] = (item for item in stripped if item is not None) @staticmethod diff --git a/litellm/router.py b/litellm/router.py index 8aaf58d5a3e..115faad000c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1813,7 +1813,6 @@ class Router: messages: list[dict[str, str]] | None, input: str | list | None, request_kwargs: dict | None, - prefetched_usage: PrefetchedUsage | None = None, ) -> Any | None: """ Asks the strategy selector for a deployment. Caller handles @@ -1839,14 +1838,6 @@ class Router: messages=messages, input=input, ) - case "usage-based-routing-v2" if isinstance(selector, LowestTPMLoggingHandler_v2): - return await selector.async_get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - prefetched_usage=prefetched_usage, - ) case "usage-based-routing-v2" | "cost-based-routing": return await selector.async_get_available_deployments( model_group=model, @@ -12958,7 +12949,6 @@ class Router: specific_deployment: bool | None = False, parent_otel_span: Span | None = None, health_check_probe: bool = False, - routing_read_batch: RoutingReadBatch | None = None, ) -> list[dict] | dict: """ Get the healthy deployments for a model. @@ -13011,6 +13001,7 @@ class Router: health_check_probe=health_check_probe, ) + routing_read_batch: Final = RoutingReadBatch.active() cooldown_deployments: Final = ( await _async_get_cooldown_deployments(litellm_router_instance=self, parent_otel_span=parent_otel_span) if routing_read_batch is None @@ -13298,15 +13289,15 @@ class Router: strategy, strategy_selector = self._get_routing_context(model, request_kwargs) routing_read_batch: Final = RoutingReadBatch.for_strategy(strategy, strategy_selector) - healthy_deployments: Final = await self.async_get_healthy_deployments( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - parent_otel_span=parent_otel_span, - routing_read_batch=routing_read_batch, - ) + with RoutingReadBatch.scoped(routing_read_batch): + healthy_deployments: Final = await self.async_get_healthy_deployments( + model=model, + request_kwargs=request_kwargs, + messages=messages, + input=input, + specific_deployment=specific_deployment, + parent_otel_span=parent_otel_span, + ) if isinstance(healthy_deployments, dict): await self._async_override_selector_pre_call_check( strategy, strategy_selector, healthy_deployments, parent_otel_span @@ -13328,16 +13319,18 @@ class Router: model=model, request_kwargs=request_kwargs, ) - deployment: Final = await self._select_deployment_async( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - prefetched_usage=routing_read_batch.prefetched_usage if routing_read_batch is not None else None, - ) + with PrefetchedUsage.scoped( + routing_read_batch.prefetched_usage if routing_read_batch is not None else None + ): + deployment: Final = await self._select_deployment_async( + strategy=strategy, + selector=strategy_selector, + model=model, + healthy_deployments=healthy_deployments, + messages=messages, + input=input, + request_kwargs=request_kwargs, + ) if deployment is None: exception: Final = await async_raise_no_deployment_exception( litellm_router_instance=self, diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 6e21d5d1f1f..25564a80e0a 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -1,7 +1,9 @@ #### What this does #### # identifies lowest tpm deployment import random -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final @@ -32,6 +34,9 @@ class RoutingArgs(LiteLLMPydanticObjectBase): ttl: int = 1 * 60 # 1min (RPM/TPM expire key) +_active_prefetched_usage: Final[ContextVar["PrefetchedUsage | None"]] = ContextVar("prefetched_usage", default=None) + + @dataclass(frozen=True) class PrefetchedUsage: """ @@ -51,6 +56,19 @@ class PrefetchedUsage: return None return [self.values.get(key) for key in keys] + @staticmethod + @contextmanager + def scoped(usage: "PrefetchedUsage | None") -> Iterator[None]: + token: Final = _active_prefetched_usage.set(usage) + try: + yield + finally: + _active_prefetched_usage.reset(token) + + @staticmethod + def active() -> "PrefetchedUsage | None": + return _active_prefetched_usage.get() + class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): """ @@ -436,18 +454,15 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): def usage_counter_keys(self, healthy_deployments: list) -> tuple[list[str], list[str]]: """The `::tpm:` and `::rpm:` counter keys selection reads.""" current_minute: Final = get_utc_datetime().strftime("%H-%M") - - tpm_keys: Final[list[str]] = [] - rpm_keys: Final[list[str]] = [] - for m in healthy_deployments: - if isinstance(m, dict): - id = m.get("model_info", {}).get( - "id" - ) # a deployment should always have an 'id'. this is set in router.py - deployment_name = m.get("litellm_params", {}).get("model") - tpm_keys.append(f"{id}:{deployment_name}:tpm:{current_minute}") - rpm_keys.append(f"{id}:{deployment_name}:rpm:{current_minute}") - return tpm_keys, rpm_keys + prefixes: Final = tuple( + f"{m.get('model_info', {}).get('id')}:{m.get('litellm_params', {}).get('model')}" + for m in healthy_deployments + if isinstance(m, dict) + ) + return ( + [f"{prefix}:tpm:{current_minute}" for prefix in prefixes], + [f"{prefix}:rpm:{current_minute}" for prefix in prefixes], + ) async def async_get_available_deployments( self, @@ -455,13 +470,13 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): healthy_deployments: list, messages: list[dict[str, str]] | None = None, input: str | list | None = None, - prefetched_usage: PrefetchedUsage | None = None, ): """ Async implementation of get deployments. - Reduces time to retrieve the tpm/rpm values from cache. `prefetched_usage` skips the cache - read when it already holds this request's counters (see `RoutingReadBatch`). + Reduces time to retrieve the tpm/rpm values from cache. A `PrefetchedUsage` scoped + to this request skips the cache read when it already holds its counters (see + `RoutingReadBatch`). """ # get list of potential deployments verbose_router_logger.debug( @@ -473,6 +488,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): tpm_keys, rpm_keys = self.usage_counter_keys(healthy_deployments) combined_tpm_rpm_keys: Final = tpm_keys + rpm_keys + prefetched_usage: Final = PrefetchedUsage.active() if prefetched_usage is not None and prefetched_usage.covers(combined_tpm_rpm_keys): combined_tpm_rpm_values = prefetched_usage.values_for(combined_tpm_rpm_keys) else: diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index cf58f3b3d3c..cf1f18abcba 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -36,7 +36,8 @@ Safe to enable globally: - No cache required. """ -from collections.abc import Iterator, Mapping +from collections.abc import Iterator, Mapping, Sequence +from functools import cache from typing import TYPE_CHECKING, Final, Optional, cast from litellm._logging import verbose_router_logger @@ -114,23 +115,31 @@ class EncryptedContentAffinityCheck(CustomLogger): if not isinstance(request_input, list): return None - for item in request_input: - if not isinstance(item, dict): - continue + return next( + ( + model_id + for item in request_input + if (model_id := EncryptedContentAffinityCheck._model_id_of_input_item(item)) is not None + ), + None, + ) - # First, try to decode from item ID (if present) - item_id = item.get("id") - if item_id and isinstance(item_id, str): - decoded = ResponsesAPIRequestUtils._decode_encrypted_item_id(item_id) - if decoded: - return decoded.get("model_id") + @staticmethod + def _model_id_of_input_item(item: object) -> str | None: + if not isinstance(item, dict): + return None - # If no encoded ID, check if encrypted_content itself is wrapped - encrypted_content = item.get("encrypted_content") - if encrypted_content and isinstance(encrypted_content, str): - model_id = EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content) - if model_id: - return model_id + item_id: Final = item.get("id") + if item_id and isinstance(item_id, str): + decoded: Final = ResponsesAPIRequestUtils._decode_encrypted_item_id(item_id) + if decoded: + return decoded.get("model_id") + + encrypted_content: Final = item.get("encrypted_content") + if encrypted_content and isinstance(encrypted_content, str): + model_id: Final = EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content) + if model_id: + return model_id return None @@ -150,19 +159,20 @@ class EncryptedContentAffinityCheck(CustomLogger): model_id, _ = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content) return model_id or None + @staticmethod + def _model_id_of_anthropic_block(block: Mapping[str, object]) -> str | None: + encrypted_content: Final = encrypted_content_of_block(block) + if encrypted_content is None: + return None + return EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content) + @staticmethod def _extract_model_id_from_anthropic_messages(messages: object) -> str | None: return next( ( model_id for block in EncryptedContentAffinityCheck._anthropic_content_blocks(messages) - if (encrypted_content := encrypted_content_of_block(block)) is not None - if ( - model_id := EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content( - encrypted_content - ) - ) - is not None + if (model_id := EncryptedContentAffinityCheck._model_id_of_anthropic_block(block)) is not None ), None, ) @@ -243,6 +253,50 @@ class EncryptedContentAffinityCheck(CustomLogger): ] return matches, originating + def _strip_reasoning_the_target_cannot_decrypt( + self, + request_input: object, + anthropic_messages: object, + target_deployments: Sequence[Mapping[str, object]], + ) -> None: + target_ids: Final = frozenset( + str(model_info["id"]) + for target in target_deployments + if isinstance((model_info := target.get("model_info")), Mapping) and model_info.get("id") is not None + ) + target_boundaries: Final = frozenset( + boundary + for target in target_deployments + if (boundary := self._encryption_boundary_key(target.get("litellm_params"))) is not None + ) + + @cache + def target_can_decrypt(origin_model_id: str) -> bool: + if origin_model_id in target_ids: + return True + if self.router is None: + return False + origin: Final = self.router.get_deployment(model_id=origin_model_id) + origin_boundary: Final = ( + self._encryption_boundary_key(origin.litellm_params.model_dump(exclude_none=True)) + if origin is not None + else None + ) + return origin_boundary is not None and origin_boundary in target_boundaries + + def should_strip_input_item(item: Mapping[str, object]) -> bool: + origin_model_id: Final = self._model_id_of_input_item(item) + return origin_model_id is not None and not target_can_decrypt(origin_model_id) + + def should_strip_anthropic_block(block: Mapping[str, object]) -> bool: + origin_model_id: Final = self._model_id_of_anthropic_block(block) + return origin_model_id is not None and not target_can_decrypt(origin_model_id) + + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input( + request_input, should_strip=should_strip_input_item + ) + strip_encrypted_reasoning_from_messages(anthropic_messages, should_strip=should_strip_anthropic_block) + # ------------------------------------------------------------------ # Request routing (pre-call filter) # ------------------------------------------------------------------ @@ -303,6 +357,7 @@ class EncryptedContentAffinityCheck(CustomLogger): model_id, ) request_kwargs["_encrypted_content_affinity_pinned"] = True + self._strip_reasoning_the_target_cannot_decrypt(request_input, anthropic_messages, (deployment,)) return [deployment] # Follow-up switched model_name (LIT-2531): pin by Azure resource instead. @@ -318,6 +373,7 @@ class EncryptedContentAffinityCheck(CustomLogger): len(boundary_matches), ) request_kwargs["_encrypted_content_affinity_pinned"] = True + self._strip_reasoning_the_target_cannot_decrypt(request_input, anthropic_messages, boundary_matches) return boundary_matches # The origin cannot serve this turn and no peer shares its encryption boundary, so its diff --git a/litellm/router_utils/routing_read_batch.py b/litellm/router_utils/routing_read_batch.py index 4039d7b1508..adda31312c0 100644 --- a/litellm/router_utils/routing_read_batch.py +++ b/litellm/router_utils/routing_read_batch.py @@ -8,11 +8,14 @@ different objects. `RoutingReadBatch` fetches both key sets in one the usage slice to the strategy, so selection does not read again. """ +import asyncio import itertools -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache @@ -21,31 +24,66 @@ from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 from litellm.router_utils.cooldown_cache import CooldownCache if TYPE_CHECKING: - from opentelemetry.trace import Span as _Span + from opentelemetry.trace import Span - from litellm.router import Router as _Router - - LitellmRouter = _Router - Span = _Span -else: - LitellmRouter = Any - Span = Any + from litellm.router import Router _PREFETCH_SLOT: Final = "routing_read" +async def _backfill_prefetched_cache( + cache: DualCache, + due_keys: tuple[str, ...], + values: Mapping[str, object], +) -> None: + cache_keys: Final = list(due_keys) # mutable-ok: _prepare_batch_get takes a list + prepare_batch_get: Final = cache._prepare_batch_get # pyright: ignore[reportPrivateUsage] # memory backfill + pending: Final = await prepare_batch_get(cache_keys, local_only=True) + redis_values: Final = { # mutable-ok: _apply_batch_get accepts a dictionary + key: values[key] + for key, local in zip(due_keys, pending.result) + if local is None and values.get(key) is not None + } + apply_batch_get: Final = cache._apply_batch_get # pyright: ignore[reportPrivateUsage] # cache backfill + await apply_batch_get(pending, redis_values) + + @dataclass(frozen=True, slots=True) class RoutingPrefetch: """The cooldown and usage keys of a model group, declared on the request's Redis batch before admission flushes it, so the routing read rides the same round trip as the rate limiter's Lua calls.""" keys: frozenset[str] + fetched: frozenset[str] result: BatchResult[Mapping[str, object]] + reservations: tuple[tuple[DualCache, tuple[str, ...], dict[str, float | None]], ...] + + def release(self) -> None: + for cache, _, previous_access_times in self.reservations: + cache._rollback_redis_batch_key_reservations( # pyright: ignore[reportPrivateUsage] # rollback + previous_access_times + ) + + async def _settle(self, future: asyncio.Future[Mapping[str, object]]) -> None: + if future.cancelled(): + self.release() + return + if future.exception() is not None: + self.release() + return + + values: Final = future.result() + try: + for cache, due_keys, _ in self.reservations: + await _backfill_prefetched_cache(cache, due_keys, values) + except Exception: + self.release() + raise @staticmethod def arm( - litellm_router_instance: LitellmRouter, + litellm_router_instance: "Router", usage_selector: LowestTPMLoggingHandler_v2 | None, deployments: list, ) -> None: @@ -60,9 +98,28 @@ class RoutingPrefetch: () if usage_selector is None else tuple(itertools.chain(*usage_selector.usage_counter_keys(deployments))) ) keys: Final = (*cooldown_keys, *usage_keys) - request.prefetched[_PREFETCH_SLOT] = RoutingPrefetch( - keys=frozenset(keys), result=request.batch(redis_cache).mget(keys) + cooldown_store: Final = litellm_router_instance.cooldown_cache.cooldown_store + cooldown_due, cooldown_previous = cooldown_store.reserve_redis_batch_reads(cooldown_keys) + usage_cache: Final = None if usage_selector is None else usage_selector.router_cache + usage_reservation: Final = None if usage_cache is None else usage_cache.reserve_redis_batch_reads(usage_keys) + usage_due: Final = () if usage_reservation is None else tuple(usage_reservation[0]) + due: Final = (*cooldown_due, *usage_due) + reservations: Final = ( + (cooldown_store, tuple(cooldown_due), cooldown_previous), + *( + () + if usage_cache is None or usage_reservation is None + else ((usage_cache, usage_due, usage_reservation[1]),) + ), ) + if not due: + return + result: Final = request.batch(redis_cache).mget(due) + prefetch: Final = RoutingPrefetch( + keys=frozenset(keys), fetched=frozenset(due), result=result, reservations=reservations + ) + result.on_settled(prefetch._settle) + request.prefetched[_PREFETCH_SLOT] = prefetch @staticmethod def armed() -> bool: @@ -78,14 +135,34 @@ class RoutingPrefetch: armed: Final = request.prefetched.pop(_PREFETCH_SLOT, None) if isinstance(armed, RoutingPrefetch) and armed.keys.issuperset(needed): return armed + if isinstance(armed, RoutingPrefetch): + armed.release() return None +_active_routing_read_batch: Final[ContextVar["RoutingReadBatch | None"]] = ContextVar( + "routing_read_batch", default=None +) + + class RoutingReadBatch: def __init__(self, usage_selector: LowestTPMLoggingHandler_v2 | None) -> None: self.usage_selector: Final = usage_selector self.prefetched_usage: PrefetchedUsage | None = None + @staticmethod + @contextmanager + def scoped(batch: "RoutingReadBatch | None") -> Iterator[None]: + token: Final = _active_routing_read_batch.set(batch) + try: + yield + finally: + _active_routing_read_batch.reset(token) + + @staticmethod + def active() -> "RoutingReadBatch | None": + return _active_routing_read_batch.get() + @staticmethod def for_strategy(strategy: str | None, selector: object) -> "RoutingReadBatch | None": """Usage-based routing reads its counters with the cooldown state; every other strategy reads only the @@ -97,9 +174,9 @@ class RoutingReadBatch: async def async_get_cooldown_deployments( self, - litellm_router_instance: LitellmRouter, + litellm_router_instance: "Router", healthy_deployments: list, - parent_otel_span: Span | None, + parent_otel_span: "Span | None", ) -> list[str]: """ `_async_get_cooldown_deployments`, with the strategy's tpm/rpm counters for @@ -107,19 +184,23 @@ class RoutingReadBatch: """ model_ids: Final = litellm_router_instance.get_model_ids() cooldown_keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] - reads: Final[list[tuple[DualCache, list[str]]]] = [ # mutable-ok: the usage read is appended below - (litellm_router_instance.cooldown_cache.cooldown_store, cooldown_keys) - ] - usage_keys: list[str] = [] # mutable-ok: DualCache batch reads take a list - if self.usage_selector is not None: - tpm_keys, rpm_keys = self.usage_selector.usage_counter_keys(healthy_deployments) - usage_keys = tpm_keys + rpm_keys - reads.append((self.usage_selector.router_cache, usage_keys)) + selector: Final = self.usage_selector + usage_keys: Final = ( + () if selector is None else tuple(itertools.chain(*selector.usage_counter_keys(healthy_deployments))) + ) + reads: Final = ( + (litellm_router_instance.cooldown_cache.cooldown_store, cooldown_keys), + *( + () + if selector is None + else ((selector.router_cache, list(usage_keys)),) # mutable-ok: DualCache batch reads take a list + ), + ) results: Final = await self._read_prefetched(reads) or await DualCache.async_batch_get_cache_shared( reads, parent_otel_span=parent_otel_span ) cooldown_results: Final = results[0] - if self.usage_selector is not None: + if selector is not None: usage_values: Final = results[1] self.prefetched_usage = PrefetchedUsage( keys=frozenset(usage_keys), @@ -134,7 +215,7 @@ class RoutingReadBatch: @staticmethod async def _read_prefetched( - reads: list[tuple[DualCache, list[str]]], + reads: Sequence[tuple[DualCache, list[str]]], ) -> list[list[object | None] | None] | None: """Serve the reads from the request's armed `RoutingPrefetch`, backfilling each cache's memory tier as its own batch read would. None when nothing usable was armed or the prefetch failed.""" @@ -149,6 +230,10 @@ class RoutingReadBatch: results: Final[list[list[object | None] | None]] = [] # mutable-ok: filled per read below for cache, keys in reads: pending = await cache._prepare_batch_get(keys, local_only=True) # pyright: ignore[reportPrivateUsage] # same two-step read as async_batch_get_cache_shared + if any( + key not in prefetch.fetched for key, local_value in zip(keys, pending.result) if local_value is None + ): + return None missed = { # mutable-ok: _apply_batch_get takes a dict key: values.get(key) for key, local in zip(keys, pending.result) if local is None } diff --git a/litellm/types/agents.py b/litellm/types/agents.py index f7aef09fa29..3b460bd66c6 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -7,6 +7,10 @@ from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, StrictInt, field from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase +from litellm.types.proxy.agent_identity import ( + AgentExecutionMode, + AgentIdentityBinding, +) if TYPE_CHECKING: from a2a.types import SendMessageResponse @@ -301,6 +305,11 @@ class AgentKeySummary(BaseModel): class AgentResponse(BaseModel): + identity: AgentIdentityBinding | None = None + identity_managed: bool = False + enabled: bool = True + execution_mode: AgentExecutionMode = "autonomous" + jwt_auth_configured: bool = False agent_id: str agent_name: str litellm_params: dict[str, object] | None = None diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 439858ea2b5..20214078852 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -199,6 +199,7 @@ class ObservabilityOptions: logger_fn: Callable[[Mapping[str, object]], None] | None = None verbose: bool | None = None no_log: bool | None = field(default=None, metadata=wire("no-log")) + log_client_error_tracebacks: bool | None = None @dataclass(frozen=True, slots=True, kw_only=True) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index a818daf554d..60d5450a1f2 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -774,6 +774,8 @@ ANTHROPIC_TOOL_SEARCH_TOOL_TYPES: Final = frozenset( # Effort beta header constant ANTHROPIC_EFFORT_BETA_HEADER: Final = "effort-2025-11-24" +ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER: Final = "mid-conversation-output-config-2026-07-01" + ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER: Final = "fine-grained-tool-streaming-2025-05-14" # OAuth constants diff --git a/litellm/types/proxy/agent_identity.py b/litellm/types/proxy/agent_identity.py new file mode 100644 index 00000000000..a7fe0be37e1 --- /dev/null +++ b/litellm/types/proxy/agent_identity.py @@ -0,0 +1,96 @@ +from datetime import datetime +from typing import Literal, TypeAlias +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field, field_validator + +AgentExecutionMode: TypeAlias = Literal["autonomous", "delegated", "both"] + + +class EntraIdentityConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + provider: Literal["microsoft_entra"] + tenant_id: str + client_id: str + service_principal_id: str | None = None + required_roles: tuple[str, ...] = () + required_scopes: tuple[str, ...] = Field( + default=("user_impersonation",), + description="Required delegated scopes. An empty list accepts any nonempty scope granted for this gateway.", + ) + + @field_validator("tenant_id", "client_id", "service_principal_id") + @classmethod + def normalize_identifier(cls, value: str | None) -> str | None: + return str(UUID(value)) if value is not None else None + + @property + def issuer(self) -> str: + return f"https://login.microsoftonline.com/{self.tenant_id}/v2.0" + + +class AgentIdentityBinding(BaseModel): + model_config = ConfigDict(frozen=True) + + agent_id: str + active: bool = True + provider: Literal["microsoft_entra"] + tenant_id: str + client_id: str + service_principal_id: str | None = None + issuer: str + required_roles: tuple[str, ...] = () + required_scopes: tuple[str, ...] = ("user_impersonation",) + revision: str + last_authenticated_at: datetime | None = None + + +class AgentSubject(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["application", "delegated_subject"] + oid: str + mode: Literal["autonomous", "delegated"] + + +class AgentIdentityFailure(BaseModel): + model_config = ConfigDict(frozen=True) + + code: Literal["identity_denied", "policy_unavailable"] = "identity_denied" + message: str + + +class ManagedAgentContext(BaseModel): + model_config = ConfigDict(frozen=True) + + agent_id: str + binding_revision: str | None = None + mode: Literal["autonomous", "delegated"] + user_id: str | None = None + subject_oid: str | None = None + + +class VerifiedHumanSubject(BaseModel): + model_config = ConfigDict(frozen=True) + + issuer: str + tenant_id: str + oid: str + user_id: str + + +class MicrosoftInteractiveSubject(BaseModel): + model_config = ConfigDict(frozen=True) + + issuer: str + tenant_id: str + oid: str + + +class ManagedAgentIdentityStatus(BaseModel): + identity: AgentIdentityBinding | None = None + identity_managed: bool = False + enabled: bool = True + execution_mode: AgentExecutionMode = "autonomous" + last_authenticated_at: datetime | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e32e3b74ec6..c12a4def69a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -284,6 +284,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_above_272k_tokens: float | None cache_creation_input_token_cost_above_272k_tokens_priority: float | None cache_creation_input_token_cost_above_272k_tokens_flex: float | None + cache_creation_input_token_cost_above_272k_tokens_ultrafast: ReadOnly[float | None] cache_creation_input_token_cost_above_1hr: float | None cache_creation_input_token_cost_flex: float | None # OpenAI flex service tier pricing cache_creation_input_token_cost_priority: float | None # OpenAI priority service tier pricing @@ -300,6 +301,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_read_input_token_cost_above_272k_tokens: float | None cache_read_input_token_cost_above_272k_tokens_priority: float | None cache_read_input_token_cost_above_272k_tokens_flex: float | None + cache_read_input_token_cost_above_272k_tokens_ultrafast: ReadOnly[float | None] cache_read_input_token_cost_above_512k_tokens: float | None cache_read_input_token_cost_batches: ReadOnly[float | None] cache_read_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None] @@ -319,6 +321,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_token_above_272k_tokens: float | None # GPT-5.4/5.4-pro: prompts >272K priced at 2x input input_cost_per_token_above_272k_tokens_priority: float | None input_cost_per_token_above_272k_tokens_flex: float | None + input_cost_per_token_above_272k_tokens_ultrafast: ReadOnly[float | None] input_cost_per_token_above_512k_tokens: float | None # MiniMax-M3: prompts >512K priced at 2x input input_cost_per_character_above_128k_tokens: float | None # only for vertex ai models input_cost_per_query: float | None # per-request pricing: rerank, search, and Bedrock Marengo embeddings @@ -360,6 +363,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token_above_272k_tokens: float | None # GPT-5.4/5.4-pro: prompts >272K priced at 1.5x output output_cost_per_token_above_272k_tokens_priority: float | None output_cost_per_token_above_272k_tokens_flex: float | None + output_cost_per_token_above_272k_tokens_ultrafast: ReadOnly[float | None] output_cost_per_token_above_512k_tokens: float | None # MiniMax-M3: prompts >512K priced at 2x output output_cost_per_character_above_128k_tokens: float | None # only for vertex ai models output_cost_per_image: float | None @@ -3737,6 +3741,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_creation_input_token_cost_above_272k_tokens: float | None = None cache_creation_input_token_cost_above_272k_tokens_priority: float | None = None cache_creation_input_token_cost_above_272k_tokens_flex: float | None = None + cache_creation_input_token_cost_above_272k_tokens_ultrafast: float | None = None cache_creation_input_token_cost_flex: float | None = None cache_creation_input_token_cost_priority: float | None = None cache_creation_input_token_cost_ultrafast: float | None = None @@ -3749,6 +3754,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_read_input_token_cost_above_200k_tokens_priority: float | None = None cache_read_input_token_cost_above_272k_tokens_priority: float | None = None cache_read_input_token_cost_above_272k_tokens_flex: float | None = None + cache_read_input_token_cost_above_272k_tokens_ultrafast: float | None = None cache_read_input_token_cost_batches: float | None = None cache_read_input_token_cost_above_200k_tokens_batches: float | None = None cache_read_input_token_cost_above_272k_tokens_batches: float | None = None @@ -3765,6 +3771,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): input_cost_per_token_above_200k_tokens_priority: float | None = None input_cost_per_token_above_272k_tokens_priority: float | None = None input_cost_per_token_above_272k_tokens_flex: float | None = None + input_cost_per_token_above_272k_tokens_ultrafast: float | None = None input_cost_per_token_above_200k_tokens_batches: float | None = None input_cost_per_token_above_272k_tokens_batches: float | None = None input_cost_per_query: float | None = None @@ -3791,6 +3798,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_token_above_200k_tokens_priority: float | None = None output_cost_per_token_above_272k_tokens_priority: float | None = None output_cost_per_token_above_272k_tokens_flex: float | None = None + output_cost_per_token_above_272k_tokens_ultrafast: float | None = None output_cost_per_token_above_200k_tokens_batches: float | None = None output_cost_per_token_above_272k_tokens_batches: float | None = None output_cost_per_character_above_128k_tokens: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index eeccd27c1d8..71186a28be4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6104,6 +6104,9 @@ def _get_model_info_helper( cache_creation_input_token_cost_above_272k_tokens_flex=_model_info.get( "cache_creation_input_token_cost_above_272k_tokens_flex", None ), + cache_creation_input_token_cost_above_272k_tokens_ultrafast=_model_info.get( + "cache_creation_input_token_cost_above_272k_tokens_ultrafast", None + ), cache_creation_input_token_cost_flex=_model_info.get("cache_creation_input_token_cost_flex", None), cache_creation_input_token_cost_priority=_model_info.get( "cache_creation_input_token_cost_priority", None @@ -6129,6 +6132,9 @@ def _get_model_info_helper( cache_read_input_token_cost_above_272k_tokens_flex=_model_info.get( "cache_read_input_token_cost_above_272k_tokens_flex", None ), + cache_read_input_token_cost_above_272k_tokens_ultrafast=_model_info.get( + "cache_read_input_token_cost_above_272k_tokens_ultrafast", None + ), cache_read_input_token_cost_above_512k_tokens=_model_info.get( "cache_read_input_token_cost_above_512k_tokens", None ), @@ -6167,6 +6173,9 @@ def _get_model_info_helper( input_cost_per_token_above_272k_tokens_flex=_model_info.get( "input_cost_per_token_above_272k_tokens_flex", None ), + input_cost_per_token_above_272k_tokens_ultrafast=_model_info.get( + "input_cost_per_token_above_272k_tokens_ultrafast", None + ), input_cost_per_token_above_512k_tokens=_model_info.get("input_cost_per_token_above_512k_tokens", None), input_cost_per_query=_model_info.get("input_cost_per_query", None), cost_per_second=_model_info.get("cost_per_second", None), @@ -6234,6 +6243,9 @@ def _get_model_info_helper( output_cost_per_token_above_272k_tokens_flex=_model_info.get( "output_cost_per_token_above_272k_tokens_flex", None ), + output_cost_per_token_above_272k_tokens_ultrafast=_model_info.get( + "output_cost_per_token_above_272k_tokens_ultrafast", None + ), output_cost_per_token_above_512k_tokens=_model_info.get( "output_cost_per_token_above_512k_tokens", None ), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9989846b7e0..136399b557e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3358,7 +3358,7 @@ "supports_function_calling": true }, "azure_ai/claude-haiku-4-5": { - "deprecation_date": "2026-10-19", + "deprecation_date": "2026-11-15", "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, @@ -3378,10 +3378,11 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule" }, "azure_ai/claude-opus-4-5": { - "deprecation_date": "2026-10-19", + "deprecation_date": "2026-11-24", "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -3402,7 +3403,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_output_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule" }, "azure_ai/claude-opus-4-6": { "deprecation_date": "2027-02-02", @@ -3640,7 +3642,7 @@ "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-4-5": { - "deprecation_date": "2026-10-19", + "deprecation_date": "2026-11-15", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -3660,7 +3662,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule" }, "azure_ai/claude-sonnet-5": { "deprecation_date": "2027-06-30", @@ -30721,6 +30724,7 @@ "output_cost_per_image": 0.08 }, "gemini/veo-3.1-fast-generate-preview": { + "deprecation_date": "2026-10-22", "litellm_provider": "gemini", "max_input_tokens": 1024, "max_tokens": 1024, @@ -30737,6 +30741,7 @@ ] }, "gemini/veo-3.1-generate-preview": { + "deprecation_date": "2026-10-22", "litellm_provider": "gemini", "max_input_tokens": 1024, "max_tokens": 1024, @@ -30752,6 +30757,7 @@ ] }, "gemini/veo-3.1-lite-generate-preview": { + "deprecation_date": "2026-10-22", "litellm_provider": "gemini", "max_input_tokens": 1024, "max_tokens": 1024, @@ -32912,10 +32918,13 @@ "gpt-image-2.5-flare": { "cache_read_input_image_token_cost": 2e-06, "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_batches": 6.25e-07, "input_cost_per_token": 5e-06, "litellm_provider": "openai", "mode": "image_generation", "input_cost_per_image_token": 8e-06, + "input_cost_per_image_token_batches": 4e-06, + "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, "supported_endpoints": [ "/v1/images/generations", @@ -32944,10 +32953,13 @@ "gpt-image-2.5-sunburst": { "cache_read_input_image_token_cost": 2e-06, "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_batches": 6.25e-07, "input_cost_per_token": 5e-06, "litellm_provider": "openai", "mode": "image_generation", "input_cost_per_image_token": 8e-06, + "input_cost_per_image_token_batches": 4e-06, + "input_cost_per_token_batches": 2.5e-06, "output_cost_per_image_token": 3e-05, "supported_endpoints": [ "/v1/images/generations", @@ -38611,6 +38623,7 @@ }, "mistral/zai-glm-5-2": { "cache_read_input_token_cost": 1.4e-07, + "deprecation_date": "2026-10-31", "input_cost_per_token": 1.4e-06, "litellm_provider": "mistral", "max_input_tokens": 1048576, @@ -38741,6 +38754,7 @@ "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-0": { + "deprecation_date": "2026-09-30", "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, "ocr_cost_per_page_batches": 0.002, @@ -60616,6 +60630,7 @@ "supports_vision": true }, "mistral/labs-leanstral-1-5": { + "deprecation_date": "2026-09-30", "input_cost_per_token": 0.0, "litellm_provider": "mistral", "max_input_tokens": 262144, @@ -61337,13 +61352,16 @@ }, "fireworks_ai/nemotron-lightning-3p5-30b-a3b": { "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_priority": 1.25e-08, "input_cost_per_token": 5e-08, + "input_cost_per_token_priority": 6.25e-08, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2e-07, + "output_cost_per_token_priority": 2.5e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, @@ -61353,13 +61371,16 @@ }, "fireworks_ai/nemotron-3-ultra-nvfp4": { "cache_read_input_token_cost": 1.2e-07, + "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 6e-07, + "input_cost_per_token_priority": 7.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2.4e-06, + "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, @@ -61389,13 +61410,16 @@ }, "fireworks_ai/accounts/fireworks/models/nemotron-lightning-3p5-30b-a3b": { "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_priority": 1.25e-08, "input_cost_per_token": 5e-08, + "input_cost_per_token_priority": 6.25e-08, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2e-07, + "output_cost_per_token_priority": 2.5e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, @@ -61405,13 +61429,16 @@ }, "fireworks_ai/accounts/fireworks/models/nemotron-3-ultra-nvfp4": { "cache_read_input_token_cost": 1.2e-07, + "cache_read_input_token_cost_priority": 1.5e-07, "input_cost_per_token": 6e-07, + "input_cost_per_token_priority": 7.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 2.4e-06, + "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, @@ -64321,13 +64348,16 @@ }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-us": { "cache_read_input_token_cost": 3.9e-07, + "cache_read_input_token_cost_priority": 4.875e-07, "input_cost_per_token": 2.1e-06, + "input_cost_per_token_priority": 2.625e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.6e-06, + "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, @@ -64356,13 +64386,16 @@ }, "fireworks_ai/glm-5p3-us": { "cache_read_input_token_cost": 3.9e-07, + "cache_read_input_token_cost_priority": 4.875e-07, "input_cost_per_token": 2.1e-06, + "input_cost_per_token_priority": 2.625e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6.6e-06, + "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_reasoning": true, @@ -64446,12 +64479,15 @@ }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-flash-us": { "cache_read_input_token_cost": 4.5e-08, + "cache_read_input_token_cost_priority": 5.625e-08, "input_cost_per_token": 2.25e-07, + "input_cost_per_token_priority": 2.8125e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 7.5e-07, + "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_response_schema": true, @@ -64477,12 +64513,15 @@ }, "fireworks_ai/glm-5p3-flash-us": { "cache_read_input_token_cost": 4.5e-08, + "cache_read_input_token_cost_priority": 5.625e-08, "input_cost_per_token": 2.25e-07, + "input_cost_per_token_priority": 2.8125e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 7.5e-07, + "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, "supports_response_schema": true, @@ -64571,6 +64610,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 6e-08, "output_cost_per_token": 2.5e-07, "litellm_provider": "together_ai", @@ -70451,6 +70491,7 @@ }, "together_ai/nvidia/nemotron-3-ultra-550b-a55b": { "cache_read_input_token_cost": 2e-07, + "deprecation_date": "2026-08-27", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_input_tokens": 512288, @@ -77476,11 +77517,14 @@ }, "fireworks_ai/accounts/fireworks/models/ember-1": { "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_priority": 3.75e-07, "input_cost_per_token": 3e-06, + "input_cost_per_token_priority": 3.75e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "mode": "chat", "output_cost_per_token": 1.5e-05, + "output_cost_per_token_priority": 1.875e-05, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_reasoning": true, @@ -79218,12 +79262,12 @@ "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1050000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "source": "https://developers.openai.com/api/docs/pricing", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_endpoints": [ "/v1/responses" ], @@ -79253,12 +79297,12 @@ "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1050000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_above_272k_tokens": 1.5e-05, - "source": "https://developers.openai.com/api/docs/pricing", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_modalities": [ "text", "image" @@ -79285,12 +79329,12 @@ "input_cost_per_token_above_272k_tokens": 4.4e-06, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "responses", "output_cost_per_token": 1.1e-05, "output_cost_per_token_above_272k_tokens": 1.65e-05, - "source": "https://developers.openai.com/api/docs/models/gpt-6.1-sol", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -79323,12 +79367,12 @@ "input_cost_per_token_above_272k_tokens": 4.4e-06, "litellm_provider": "bedrock_converse", "max_input_tokens": 1050000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.1e-05, "output_cost_per_token_above_272k_tokens": 1.65e-05, - "source": "https://developers.openai.com/api/docs/pricing", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_endpoints": [ "/v1/responses" ], @@ -79348,5 +79392,33 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true + }, + "vertex_ai/gemini-3.8-flash-tts": { + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "audio_speech", + "output_cost_per_audio_token": 9e-06, + "output_cost_per_token": 9e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, + "vertex_ai/gemini-3.8-flash-lite-tts": { + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "audio_speech", + "output_cost_per_audio_token": 6e-06, + "output_cost_per_token": 6e-06, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/audio/speech" + ] } } diff --git a/osv-scanner.toml b/osv-scanner.toml index 482254d4da6..9bb346a94f9 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -7,3 +7,13 @@ reason = "diskcache has no fixed release published; remove this entry once one e id = "GHSA-h7x2-h6g9-p789" ignoreUntil = 2026-10-14 reason = "mlflow has no fixed release published (3.16.0, 2026-09-04, and master still store gateway secret api_base unvalidated); remove this entry once one exists" + +[[IgnoredVulns]] +id = "GHSA-hj66-6f7g-4r5v" +ignoreUntil = 2026-10-02 +reason = "oauthlib 4.0.0 (the only fixed release, 2026-09-28) is inside the 3-day uv exclude-newer cooldown; bump oauthlib and remove this entry once it clears" + +[[IgnoredVulns]] +id = "GHSA-xpv3-w29h-x7cv" +ignoreUntil = 2026-10-02 +reason = "oauthlib 4.0.0 (the only fixed release, 2026-09-28) is inside the 3-day uv exclude-newer cooldown; bump oauthlib and remove this entry once it clears" diff --git a/pyproject.toml b/pyproject.toml index fb21d8fa23b..77a1a3fdb75 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.104.0" +version = "1.105.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -75,8 +75,8 @@ proxy = [ "mcp>=2.2.0,<3", "httpx2>=2.5.0,<3", "pydantic>=2.12.0,<3", - "litellm-proxy-extras==0.4.102", - "litellm-enterprise==0.1.71", + "litellm-proxy-extras==0.4.103", + "litellm-enterprise==0.1.72", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", @@ -357,7 +357,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras"] [tool.commitizen] -version = "1.104.0" +version = "1.105.0" version_files = [ "pyproject.toml:^version", ] diff --git a/schema.prisma b/schema.prisma index 03e59257f76..f29caa9ceb7 100644 --- a/schema.prisma +++ b/schema.prisma @@ -78,6 +78,11 @@ model LiteLLM_AgentsTable { object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) + identity_managed Boolean @default(false) + enabled Boolean @default(true) + execution_mode String @default("autonomous") + identity LiteLLM_AgentIdentity? + retired_identities LiteLLM_RetiredAgentIdentity[] tpm_limit Int? rpm_limit Int? session_tpm_limit Int? @@ -88,6 +93,56 @@ model LiteLLM_AgentsTable { updated_by String } +model LiteLLM_AgentIdentity { + agent_id String @id + active Boolean @default(true) + agent LiteLLM_AgentsTable @relation(fields: [agent_id], references: [agent_id], onDelete: Cascade) + provider String + issuer String + tenant_id String + client_id String + service_principal_id String? + required_roles String[] @default([]) + required_scopes String[] @default(["user_impersonation"]) + revision String @default(uuid()) + last_authenticated_at DateTime? + @@unique([provider, tenant_id, client_id]) + @@unique([issuer, service_principal_id]) +} + +model LiteLLM_RetiredAgentIdentity { + binding_id String @id @default(uuid()) + agent_id String? + agent LiteLLM_AgentsTable? @relation(fields: [agent_id], references: [agent_id], onDelete: SetNull) + provider String + issuer String + tenant_id String + client_id String + @@unique([provider, tenant_id, client_id]) +} + +model LiteLLM_RetiredAgent { + original_agent_id String @id + retired_at DateTime @default(now()) +} + +model LiteLLM_VerifiedSubject { + subject_id String @id @default(uuid()) + issuer String + tenant_id String + oid String + kind String @default("human") + user_id String? + user LiteLLM_UserTable? @relation(fields: [user_id], references: [user_id], onDelete: Cascade) + verified_via String @default("sso_interactive") + verified_at DateTime @default(now()) + @@unique([issuer, tenant_id, oid]) + @@index([user_id]) +} + + + + model LiteLLM_OrganizationTable { organization_id String @id @default(uuid()) organization_alias String @@ -241,6 +296,7 @@ model LiteLLM_DeletedTeamTable { // Track spend, rate limit, budget Users model LiteLLM_UserTable { + verified_subjects LiteLLM_VerifiedSubject[] user_id String @id user_alias String? team_id String? @@ -675,6 +731,7 @@ model LiteLLM_SpendLogs { session_id String? status String? mcp_namespaced_tool_name String? + billing_agent_id String? agent_id String? proxy_server_request Json? @default("{}") litellm_call_id String? diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 8e221b2da5e..fb2cf2dfa24 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -64,7 +64,7 @@ The suites run against a live proxy, so bring one up first by running the litell For the opt-in browser profile, start the existing IdP first, then run `.github/e2e-stack/oidc-profile.sh "$PROXY_BASE_URL" `. The wrapper creates a confidential client with an exact `/sso/callback` redirect and S256 PKCE, passes the client secret only through the child process environment, and removes the client on exit. It uses the existing generic OIDC handler with `GENERIC_USER_ID_ATTRIBUTE=sub`. Preserve the IdP's PostgreSQL data across restarts - `tests/e2e/ui/playwright.oidc.config.ts` uses an already running OIDC stack and separate storage/output files. Supply `E2E_OIDC_UI_URL`, `JWT_ISSUER`, `E2E_OIDC_USERNAME` and `E2E_OIDC_PASSWORD` for a seeded actor. Its setup follows the real login and callback path. The current Python canary qualifies browser-client configuration and token/userinfo identity mapping; browser journey specs under `ui/oidc/` are a separate coverage step + `tests/e2e/ui/playwright.oidc.config.ts` uses an already running OIDC stack and separate storage/output files. Supply `E2E_OIDC_UI_URL`, `JWT_ISSUER`, `E2E_OIDC_USERNAME` and `E2E_OIDC_PASSWORD` for a seeded actor. Its setup follows the real login and callback path. The current Python canary qualifies browser-client configuration and token/userinfo identity mapping. The specs under `ui/oidc/` drive a real dashboard SSO login and a real `lite login`, so start the proxy with `EXPERIMENTAL_UI_LOGIN=true` and at least one model it can actually serve. The CLI spec runs `lite` from `PATH` unless `E2E_LITE_CLI` names another executable, and it gives the CLI a temporary `HOME` with the keyring disabled so your own login is never touched. The main `playwright.config.ts` ignores `oidc/` Every successful IdP create immediately registers cleanup, including partial setup failures. Cleanup failures emit warnings. Tokens are minted on demand, and the expiration test waits relative to the token's actual `exp` with a bounded clock-drift check. To check first-attempt behavior locally, run both files with `--reruns 0`: diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index 3bd98ff5b0b..0b9249d7420 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -60,3 +60,6 @@ - {id: other.auth.jwt.wrong_issuer_denied, module: other, tier: P0, area: auth, assertions: [wrong_issuer_denied], source: "auth/handle_jwt.py", rationale: "A signed token with the correct audience and an unexpected issuer is rejected"} - {id: other.auth.jwt.wrong_audience_denied, module: other, tier: P0, area: auth, assertions: [wrong_audience_denied], source: "auth/handle_jwt.py", rationale: "A signed token from the trusted issuer intended for another app is rejected"} +- {id: other.auth.session_token.valid_allows, module: other, tier: P0, area: auth, assertions: [valid_allows], source: "auth/auth_checks.py ExperimentalUIJWTToken", rationale: "An unexpired LiteLLM-minted session token authenticates with the role it carries"} +- {id: other.auth.session_token.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "auth/user_api_key_auth.py expiry check", rationale: "An expired session token is rejected with the expired-key error"} +- {id: other.auth.session_token.encrypted_value_denied, module: other, tier: P0, area: auth, assertions: [encrypted_value_denied], source: "auth/auth_checks.py ExperimentalUIJWTToken", rationale: "An encrypted value read back from a management route is not accepted as a bearer token"} diff --git a/tests/e2e/other/test_session_token_e2e.py b/tests/e2e/other/test_session_token_e2e.py new file mode 100644 index 00000000000..51791278026 --- /dev/null +++ b/tests/e2e/other/test_session_token_e2e.py @@ -0,0 +1,91 @@ +"""Live e2e: UI/CLI session tokens are accepted only while valid and only when minted as session tokens. + +The runner mints its own session tokens under the proxy's salt key, so the valid and expired cases run in +seconds instead of waiting out a real login's expiry. +""" + +from __future__ import annotations + +import base64 +import hashlib +import json +import os +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest +from cryptography.hazmat.primitives.ciphers.aead import AESGCM +from e2e_config import MASTER_KEY, unique_marker +from e2e_http import UnauthorizedError, unwrap +from lifecycle import ResourceManager +from models import KeyGenerateBody, KeyLoggingCallback, KeyLoggingCallbackVars, KeyMetadata +from other_client import OtherClient + +pytestmark = pytest.mark.e2e + +SALT_KEY: Final = os.environ.get("LITELLM_SALT_KEY") or MASTER_KEY +SESSION_TOKEN_PREFIX: Final = "litellm_login_" +ENCRYPTED_PREFIX: Final = "litellm_enc::" + + +def _admin_session_token(expires_at: datetime) -> str: + claims: Final = json.dumps( + { + "token": f"ui-token-{unique_marker()}", + "user_id": f"e2e-session-{unique_marker()}", + "user_role": "proxy_admin", + "team_id": "litellm-dashboard", + "expires": expires_at.isoformat(), + } + ) + nonce: Final = os.urandom(12) + sealed: Final = AESGCM(hashlib.sha256(SALT_KEY.encode()).digest()).encrypt( + nonce, claims.encode(), SESSION_TOKEN_PREFIX.encode() + ) + return SESSION_TOKEN_PREFIX + base64.urlsafe_b64encode(nonce + sealed).decode().rstrip("=") + + +class TestSessionToken: + @pytest.mark.covers("other.auth.session_token.valid_allows") + def test_unexpired_session_token_reaches_admin_route(self, client: OtherClient) -> None: + token: Final = _admin_session_token(datetime.now(timezone.utc) + timedelta(minutes=10)) + listing: Final = unwrap(client.list_users_as(token)) + assert listing.total >= 0, f"an unexpired admin session token did not reach /user/list: {listing}" + + @pytest.mark.covers("other.auth.session_token.expired_denied") + def test_expired_session_token_is_denied(self, client: OtherClient) -> None: + token: Final = _admin_session_token(datetime.now(timezone.utc) - timedelta(minutes=1)) + result: Final = client.list_users_as(token) + assert isinstance(result, UnauthorizedError), f"an expired session token must get 401, got {result}" + assert "expired" in result.body.lower(), f"expected the expired-key error, got {result.body[:300]}" + + @pytest.mark.covers("other.auth.session_token.encrypted_value_denied") + def test_encrypted_stored_value_is_not_a_bearer_token( + self, client: OtherClient, resources: ResourceManager + ) -> None: + stored_value: Final = f'{{"token": "{unique_marker()}", "user_role": "proxy_admin"}}' + key: Final = client.proxy.generate_key( + KeyGenerateBody( + key_alias=f"e2e-session-{unique_marker()}", + metadata=KeyMetadata( + logging=[ + KeyLoggingCallback( + callback_name="langfuse", + callback_vars=KeyLoggingCallbackVars(langfuse_secret_key=stored_value), + ) + ] + ), + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + metadata: Final = client.proxy.key_info(key).metadata + assert metadata is not None and metadata.logging, f"/key/info dropped the logging metadata: {metadata}" + encrypted: Final = metadata.logging[0].callback_vars.langfuse_secret_key + assert encrypted is not None and encrypted.startswith(ENCRYPTED_PREFIX), ( + f"expected /key/info to return the stored secret encrypted, got {encrypted!r}" + ) + + for bearer in (encrypted.removeprefix(ENCRYPTED_PREFIX), encrypted): + result = client.list_users_as(bearer) + assert isinstance(result, UnauthorizedError), f"an encrypted stored value must get 401, got {result}" diff --git a/tests/e2e/ui/helpers/traffic.ts b/tests/e2e/ui/helpers/traffic.ts index cb68747b364..b534c475221 100644 --- a/tests/e2e/ui/helpers/traffic.ts +++ b/tests/e2e/ui/helpers/traffic.ts @@ -51,6 +51,21 @@ export async function sendChatCompletion(request: APIRequestContext, opts: ChatO return body.id as string; } +export interface ServedChat { + requestId: string; + callId: string; +} + +export async function sendChatCompletionWithCallId(request: APIRequestContext, opts: ChatOptions): Promise { + const res = await postChatCompletion(request, opts); + expect(res.ok(), `chat completion for ${opts.model} failed (${res.status()}): ${await res.text()}`).toBe(true); + const callId = res.headers()["x-litellm-call-id"]; + expect(callId, "proxy did not return an x-litellm-call-id header").toBeTruthy(); + const body = await res.json(); + expect(body.choices?.[0]?.message?.content).toContain(MOCK_RESPONSE_TEXT); + return { requestId: body.id as string, callId }; +} + export interface ChatAttempt { status: number; body: string; @@ -124,7 +139,7 @@ export async function waitForSpendLog( lastStatus = res.status(); if (res.ok()) { const body = await res.json(); - const rows = Array.isArray(body) ? body : (body?.data ?? []); + const rows = Array.isArray(body) ? body : body?.data ?? []; if (rows.length > 0) { return; } diff --git a/tests/e2e/ui/oidc/cliLogin.spec.ts b/tests/e2e/ui/oidc/cliLogin.spec.ts new file mode 100644 index 00000000000..89b9a7c7439 --- /dev/null +++ b/tests/e2e/ui/oidc/cliLogin.spec.ts @@ -0,0 +1,87 @@ +import { expect, test } from "@playwright/test"; +import { execFile, spawn } from "node:child_process"; +import * as fs from "node:fs"; +import * as os from "node:os"; +import * as path from "node:path"; +import { promisify } from "node:util"; + +const LITE_CLI = process.env.E2E_LITE_CLI ?? "lite"; +const SKIP_TEAM_SELECTION = "skip\n"; +const execFileAsync = promisify(execFile); + +function requiredEnv(name: string): string { + const value = process.env[name]; + if (!value) throw new Error(`${name} must be set for the OIDC suite`); + return value; +} + +test("CLI SSO login stores a session that lists models and completes a chat request", async ({ browser, baseURL }) => { + test.setTimeout(180_000); + const issuer = requiredEnv("JWT_ISSUER"); + const home = fs.mkdtempSync(path.join(os.tmpdir(), "lite-cli-login-")); + const browserUrlFile = path.join(home, "browser-url"); + const browserCommand = path.join(home, "browser.sh"); + fs.writeFileSync(browserCommand, `#!/bin/sh\nprintf '%s' "$1" > '${browserUrlFile}'\n`, { mode: 0o700 }); + const env = { + ...process.env, + HOME: home, + LITELLM_CLI_DISABLE_KEYRING: "1", + BROWSER: browserCommand, + PYTHONUNBUFFERED: "1", + FORCE_COLOR: undefined, + NO_COLOR: "1", + LITELLM_PROXY_URL: baseURL, + LITELLM_PROXY_API_KEY: undefined, + }; + const login = spawn(LITE_CLI, ["login"], { env }); + let loginOutput = ""; + login.stdout.on("data", (chunk: Buffer) => (loginOutput += chunk.toString())); + login.stderr.on("data", (chunk: Buffer) => (loginOutput += chunk.toString())); + const loginExit = new Promise((resolve) => login.on("close", resolve)); + login.stdin.end(SKIP_TEAM_SELECTION); + try { + await expect.poll(() => fs.existsSync(browserUrlFile), { timeout: 30_000 }).toBe(true); + await expect.poll(() => loginOutput).toMatch(/Verification code: \S+/); + const userCode = /Verification code: (\S+)/.exec(loginOutput)?.[1] ?? ""; + + const context = await browser.newContext({ storageState: { cookies: [], origins: [] } }); + try { + const page = await context.newPage(); + await page.goto(fs.readFileSync(browserUrlFile, "utf8")); + await expect(page).toHaveURL((url) => url.href.startsWith(`${issuer}/`)); + await page.getByLabel("Username or email").fill(requiredEnv("E2E_OIDC_USERNAME")); + await page.getByLabel("Password", { exact: true }).fill(requiredEnv("E2E_OIDC_PASSWORD")); + await page.getByRole("button", { name: "Sign In", exact: true }).click(); + await page.getByLabel("Verification code").fill(userCode); + await page.getByRole("button", { name: "Continue", exact: true }).click(); + await expect(page.getByRole("heading", { name: "Authentication Successful!" })).toBeVisible(); + } finally { + await context.close(); + } + + expect(await loginExit, loginOutput).toBe(0); + expect(loginOutput).toContain("Login successful!"); + const stored: { key?: unknown } = JSON.parse(fs.readFileSync(path.join(home, ".litellm", "token.json"), "utf8")); + expect(typeof stored.key).toBe("string"); + expect(stored.key, "CLI login issues a session token, not a virtual key").not.toMatch(/^sk-/); + + const { stdout: modelsJson } = await execFileAsync(LITE_CLI, ["models", "list", "--format", "json"], { env }); + const models: { id: string }[] = JSON.parse(modelsJson); + expect(models.length, "the stack serves at least one model").toBeGreaterThan(0); + + const chatRequest = JSON.stringify({ + model: models[0].id, + messages: [{ role: "user", content: "Reply with the single word: ok" }], + }); + const { stdout: completionJson } = await execFileAsync( + LITE_CLI, + ["http", "request", "POST", "/chat/completions", "-j", chatRequest], + { env }, + ); + const completion: { choices: { message: { content: string | null } }[] } = JSON.parse(completionJson); + expect(completion.choices[0]?.message.content).toBeTruthy(); + } finally { + login.kill(); + fs.rmSync(home, { recursive: true, force: true }); + } +}); diff --git a/tests/e2e/ui/oidc/dashboardLogin.spec.ts b/tests/e2e/ui/oidc/dashboardLogin.spec.ts new file mode 100644 index 00000000000..106646949ed --- /dev/null +++ b/tests/e2e/ui/oidc/dashboardLogin.spec.ts @@ -0,0 +1,35 @@ +import { expect, test, type Page as PlaywrightPage, type Response } from "@playwright/test"; +import { Page } from "../fixtures/pages"; +import { navigateToPage } from "../helpers/navigation"; + +function sessionKey(tokenCookie: string): string { + const claims: unknown = JSON.parse(Buffer.from(tokenCookie.split(".")[1] ?? "", "base64url").toString("utf8")); + const key = claims !== null && typeof claims === "object" && "key" in claims ? claims.key : undefined; + if (typeof key !== "string") throw new Error("The dashboard token cookie carries no key claim"); + return key; +} + +async function openPageAndCapture(page: PlaywrightPage, target: Page, apiPath: string): Promise { + const response = page.waitForResponse((r) => new URL(r.url()).pathname === apiPath); + await navigateToPage(page, target); + return response; +} + +test("SSO login issues a session that authorizes dashboard data requests", async ({ page, context, baseURL }) => { + const tokenCookie = (await context.cookies(baseURL)).find((cookie) => cookie.name === "token"); + expect(tokenCookie, "SSO login sets the dashboard token cookie").toBeDefined(); + const key = sessionKey(tokenCookie?.value ?? ""); + expect(key, "SSO login issues a session token, not a virtual key").not.toMatch(/^sk-/); + + const keyList = await openPageAndCapture(page, Page.ApiKeys, "/key/list"); + expect(keyList.request().headers()["authorization"]).toBe(`Bearer ${key}`); + expect(keyList.status()).toBe(200); + expect(Array.isArray((await keyList.json()).keys)).toBe(true); + + const modelInfo = await openPageAndCapture(page, Page.Models, "/v2/model/info"); + expect(modelInfo.request().headers()["authorization"]).toBe(`Bearer ${key}`); + expect(modelInfo.status()).toBe(200); + const models: { model_name: string }[] = (await modelInfo.json()).data; + expect(models.length, "the stack serves at least one model").toBeGreaterThan(0); + await expect(page.getByText(models[0].model_name, { exact: true }).first()).toBeVisible(); +}); diff --git a/tests/e2e/ui/playwright.config.ts b/tests/e2e/ui/playwright.config.ts index 2fc3b5f2d81..aed70620280 100644 --- a/tests/e2e/ui/playwright.config.ts +++ b/tests/e2e/ui/playwright.config.ts @@ -8,7 +8,7 @@ import { ARTIFACT_DIR, UI_BASE_URL } from "./constants"; export default defineConfig({ testDir: ".", testMatch: ["**/*.spec.ts", "**/*.setup.ts"], - testIgnore: ["**/*.test.*", "**/integrationCritical/**"], + testIgnore: ["**/*.test.*", "**/integrationCritical/**", "oidc/**"], /* Run tests in files in parallel */ fullyParallel: true, /* Fail the build on CI if you accidentally left test.only in the source code. */ diff --git a/tests/e2e/ui/tests/logs/logs.spec.ts b/tests/e2e/ui/tests/logs/logs.spec.ts index 2748c91395f..3908d79b29a 100644 --- a/tests/e2e/ui/tests/logs/logs.spec.ts +++ b/tests/e2e/ui/tests/logs/logs.spec.ts @@ -6,6 +6,7 @@ import { CHAT_MODEL_A, MOCK_RESPONSE_TEXT, sendChatCompletion, + sendChatCompletionWithCallId, waitForSpendLog, waitForSpendLogByPrompt, } from "../../helpers/traffic"; @@ -95,6 +96,50 @@ test.describe("Logs page", () => { await expect(drawer.getByText(MOCK_RESPONSE_TEXT, { exact: false }).first()).toBeVisible({ timeout: 20_000 }); }); + test("a served request's Logs row and drawer show its x-litellm-call-id", async ({ page, request }) => { + const prompt = `logs-call-id-prompt-${uniqueSuffix()}`; + const { requestId, callId } = await sendChatCompletionWithCallId(request, { + model: CHAT_MODEL_A, + prompt, + }); + expect(callId, "call id must differ from the provider response id for this check to mean anything").not.toBe( + requestId, + ); + await waitForSpendLog(request, requestId); + + await navigateToPage(page, Page.Logs); + await dismissFeedbackPopup(page); + const search = visibleTestId(page, "datatable-search"); + await expect(search).toBeVisible({ timeout: 20_000 }); + await search.fill(callId); + + const row = requestLogsRows(page).filter({ hasText: requestId }); + await expect(row, `no logs row for call id ${callId}`).toHaveCount(1, { timeout: 30_000 }); + await expect(row, "the row itself shows only the request id").not.toContainText(callId); + + await row.getByText(requestId).hover(); + const tooltip = page.locator("[data-slot='tooltip-content']"); + await expect(tooltip, "hovering the Request ID cell does not list the x-litellm-call-id").toContainText( + `x-litellm-call-id: ${callId}`, + { timeout: 10_000 }, + ); + await tooltip.getByRole("button", { name: "Copy x-litellm-call-id" }).click(); + if (await page.evaluate(() => window.isSecureContext)) { + await expect.poll(() => page.evaluate(() => navigator.clipboard.readText())).toBe(callId); + } + + await row.click(); + const drawer = page.getByRole("dialog").first(); + await expect(drawer.getByText("Request & Response")).toBeVisible({ timeout: 20_000 }); + await expect(drawer.getByText("x-litellm-call-id:"), "drawer header lacks the x-litellm-call-id line").toBeVisible({ + timeout: 10_000, + }); + await expect( + drawer.getByText(callId, { exact: false }).first(), + `drawer does not show x-litellm-call-id ${callId}`, + ).toBeVisible({ timeout: 10_000 }); + }); + // Split out because only the copy path needs a secure context; folding it in would // take the drawer-rendering coverage down with it. test("the drawer copies the request and the response to the clipboard", async ({ page, request }) => { diff --git a/tests/integration/_support/daily_activity.py b/tests/integration/_support/daily_activity.py index debb8c4cdb4..346fc156ea9 100644 --- a/tests/integration/_support/daily_activity.py +++ b/tests/integration/_support/daily_activity.py @@ -3,6 +3,8 @@ import uuid from collections.abc import Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass +from datetime import datetime, timedelta +from hashlib import sha256 from itertools import chain from typing import Final @@ -34,7 +36,16 @@ INSERT_SPEND_LOG: Final = ( " VALUES (%s, 'acompletion', %s, %s::timestamp, %s::timestamp, %s)" ) DELETE_SPEND_LOG: Final = 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = %s' +INSERT_SPEND_LOG_ROW: Final = ( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, api_key, "startTime", "endTime", metadata, team_id, "user")' + " VALUES (%s, 'acompletion', %s, %s::timestamp, %s::timestamp, %s, %s, %s)" +) +DELETE_SPEND_LOG_ROWS: Final = 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)' +DELETE_KEY_ROW: Final = 'DELETE FROM "LiteLLM_VerificationToken" WHERE token = %s' +DELETE_ARCHIVED_KEY_ROW: Final = 'DELETE FROM "LiteLLM_DeletedVerificationToken" WHERE token = %s' LOCK_TABLE: Final = sql.SQL("LOCK TABLE {table} IN ACCESS EXCLUSIVE MODE") +SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs" +FIRST_SPEND_LOG_AT: Final = datetime(2026, 2, 3, 12, 0, 0) @dataclass(frozen=True, slots=True) @@ -67,6 +78,10 @@ def key_no_key_table_holds() -> str: return f"integration-ownerless-{uuid.uuid4().hex}" +def digest_no_key_table_holds() -> str: + return sha256(uuid.uuid4().bytes).hexdigest() + + def activity_of_key( gateway: Gateway, path: str, api_key: str, *, reader: str | None = None, **filters: str ) -> httpx.Response: @@ -148,6 +163,56 @@ def spend_log_naming_only_an_alias(request_id: str, api_key: str, started: str, connection.execute(DELETE_SPEND_LOG, (request_id,)) +@dataclass(frozen=True, slots=True) +class SpendLogRow: + started: str + metadata: JsonValue = None + team_id: str | None = None + user: str | None = None + + +def started_at(index: int) -> str: + return (FIRST_SPEND_LOG_AT + timedelta(seconds=index)).strftime("%Y-%m-%d %H:%M:%S") + + +def nameless_rows(count: int, first_index: int = 0) -> tuple[SpendLogRow, ...]: + return tuple(SpendLogRow(started_at(first_index + offset), {}) for offset in range(count)) + + +def named_row(index: int, alias: str) -> SpendLogRow: + return SpendLogRow(started_at(index), {"user_api_key_alias": alias}) + + +@contextmanager +def spend_logs_of_key( + api_key: str, rows: Sequence[SpendLogRow], *, database_url: str | None = None +) -> Iterator[tuple[str, ...]]: + request_ids: Final = tuple(f"integration-{uuid.uuid4().hex}" for _ in rows) + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + connection.cursor().executemany( + INSERT_SPEND_LOG_ROW, + tuple( + (request_id, api_key, row.started, row.started, Jsonb(row.metadata), row.team_id, row.user) + for request_id, row in zip(request_ids, rows, strict=True) + ), + ) + try: + yield request_ids + finally: + delete_spend_logs(request_ids, database_url=database_url) + + +def delete_spend_logs(request_ids: Sequence[str], *, database_url: str | None = None) -> None: + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + connection.execute(DELETE_SPEND_LOG_ROWS, (list(request_ids),)) + + +def purge_key_from_the_key_tables(digest: str, *, database_url: str | None = None) -> None: + with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: + connection.execute(DELETE_KEY_ROW, (digest,)) + connection.execute(DELETE_ARCHIVED_KEY_ROW, (digest,)) + + @contextmanager def locked_table(table: str, *, database_url: str | None = None) -> Iterator[None]: with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection: diff --git a/tests/integration/authorization/test_team_admin_gate.py b/tests/integration/authorization/test_team_admin_gate.py index 2b02e90fcc5..d62a2a1a9a6 100644 --- a/tests/integration/authorization/test_team_admin_gate.py +++ b/tests/integration/authorization/test_team_admin_gate.py @@ -324,7 +324,7 @@ ROUTES: Final[tuple[Route, ...]] = ( lambda s: Call("POST", "/team/update", {"team_id": s.team_id, "max_budget": 5}), team_admin=403, others=403, org_admin=200), Route("team_update_budget_permitted", - lambda s: Call("POST", "/team/update", {"team_id": s.team_id, "max_budget": 7}), + lambda s: Call("POST", "/team/update", {"team_id": s.team_id, "max_budget": 4}), team_admin=200, others=403, org_admin=200, permission="max_budget"), Route("project_new", lambda s: Call("POST", "/project/new", {"team_id": s.team_id, "project_alias": f"matrix-{uuid.uuid4().hex}"}), @@ -414,6 +414,8 @@ def test_status_code(shared: TeamScenario, org_team: TeamScenario, route: Route, team: Final = org_team if caller in ORG_CALLERS else shared with team.gateway.scenario() as scenario: s: Final = replace(team, scenario=scenario) + if route.name == "team_update_budget_permitted": + s.gateway.post("/team/update", {"team_id": s.team_id, "max_budget": 5}) if route.permission: scenario.cleanups.enter_context(team_admin_permissions(s.gateway, (route.permission,))) call: Final = route.call(s) @@ -421,5 +423,9 @@ def test_status_code(shared: TeamScenario, org_team: TeamScenario, route: Route, assert response.status_code == route.expected(caller), ( f"{caller} {call.method} {call.path}: {response.status_code} {response.text}" ) + if route.name == "team_update_budget_permitted": + assert read_rows( + 'SELECT max_budget FROM "LiteLLM_TeamTable" WHERE team_id = %s', (s.team_id,) + ) == [{"max_budget": 4.0 if response.status_code == 200 else 5.0}] if response.status_code == 200 and route.cleanup is not None: route.cleanup(s, object_value(response.json())) diff --git a/tests/integration/pricing/test_service_tier_pricing.py b/tests/integration/pricing/test_service_tier_pricing.py index e0d26392f7f..0917c744bbf 100644 --- a/tests/integration/pricing/test_service_tier_pricing.py +++ b/tests/integration/pricing/test_service_tier_pricing.py @@ -1,11 +1,15 @@ import json -from typing import Final +import uuid +from typing import Final, Literal import httpx import pytest +from pydantic import JsonValue from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value from tests.integration._support.database import read_rows +from tests.integration._support.upstream import delete_scenario, register_scenario +from tests.integration.cost_calculation.cost_tracking_case import JsonResponse STANDARD_INPUT_RATE: Final = 0.001 STANDARD_OUTPUT_RATE: Final = 0.002 @@ -69,3 +73,190 @@ def test_ultrafast_service_tier_bills_ultrafast_rates_and_keeps_pricing_off_the_ ) assert_chat_bills_rates(gateway, model, "ultrafast", ULTRAFAST_INPUT_RATE, ULTRAFAST_OUTPUT_RATE) assert_chat_bills_rates(gateway, model, None, STANDARD_INPUT_RATE, STANDARD_OUTPUT_RATE) + + +LONG_CONTEXT_PRICING: Final[dict[str, JsonValue]] = { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token_above_272k_tokens": 3e-06, + "output_cost_per_token_above_272k_tokens": 4e-06, + "cache_read_input_token_cost_above_272k_tokens": 3e-07, + "input_cost_per_token_ultrafast": 1e-05, + "output_cost_per_token_ultrafast": 2e-05, + "cache_read_input_token_cost_ultrafast": 1e-06, + "input_cost_per_token_above_272k_tokens_ultrafast": 5e-05, + "output_cost_per_token_above_272k_tokens_ultrafast": 6e-05, + "cache_read_input_token_cost_above_272k_tokens_ultrafast": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_ultrafast": 6e-06, +} +LONG_PROMPT_TOKENS: Final = 300_000 +SHORT_PROMPT_TOKENS: Final = 1_000 +CACHED_TOKENS: Final = 400 +COMPLETION_TOKENS: Final = 1_000 + + +def _chat_response(service_tier: str | None, prompt_tokens: int) -> JsonResponse: + return JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-$UNIQUE_ID", + "object": "chat.completion", + "created": 1, + "model": "integration-ultrafast-long-context", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "long context answer"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": prompt_tokens + COMPLETION_TOKENS, + "prompt_tokens_details": {"cached_tokens": CACHED_TOKENS}, + }, + **({} if service_tier is None else {"service_tier": service_tier}), + }, + ) + + +def _responses_response(service_tier: str | None, prompt_tokens: int) -> JsonResponse: + return JsonResponse( + content_type="application/json", + body={ + "id": "resp_$UNIQUE_ID", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "integration-ultrafast-long-context", + "output": [ + { + "type": "message", + "id": "msg_$UNIQUE_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "long context answer", "annotations": []}], + } + ], + "usage": { + "input_tokens": prompt_tokens, + "output_tokens": COMPLETION_TOKENS, + "total_tokens": prompt_tokens + COMPLETION_TOKENS, + "input_tokens_details": {"cached_tokens": CACHED_TOKENS}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + **({} if service_tier is None else {"service_tier": service_tier}), + }, + ) + + +def _surface_response( + surface: Literal["chat", "responses"], service_tier: str | None, prompt_tokens: int +) -> JsonResponse: + match surface: + case "chat": + return _chat_response(service_tier, prompt_tokens) + case "responses": + return _responses_response(service_tier, prompt_tokens) + + +def _surface_request( + surface: Literal["chat", "responses"], scenario_id: str, model: str, service_tier: str | None +) -> tuple[str, dict[str, JsonValue], str]: + match surface: + case "chat": + return ( + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "long context ultrafast control"}], + **({} if service_tier is None else {"service_tier": service_tier}), + }, + f"/{scenario_id}/chat/completions", + ) + case "responses": + return ( + "/v1/responses", + { + "model": model, + "input": "long context ultrafast control", + **({} if service_tier is None else {"service_tier": service_tier}), + }, + f"/{scenario_id}/responses", + ) + + +@pytest.mark.parametrize( + ("service_tier", "prompt_tokens", "input_rate", "cache_read_rate", "output_rate"), + ( + ("ultrafast", LONG_PROMPT_TOKENS, 5e-05, 5e-06, 6e-05), + ("ultrafast", SHORT_PROMPT_TOKENS, 1e-05, 1e-06, 2e-05), + (None, LONG_PROMPT_TOKENS, 3e-06, 3e-07, 4e-06), + ), + ids=("ultrafast_above_272k", "ultrafast_below_272k", "standard_above_272k"), +) +@pytest.mark.parametrize("surface", ("chat", "responses"), ids=("chat", "responses")) +def test_ultrafast_long_context_prompt_bills_ultrafast_long_context_rates( + gateway: Gateway, + surface: Literal["chat", "responses"], + service_tier: str | None, + prompt_tokens: int, + input_rate: float, + cache_read_rate: float, + output_rate: float, +) -> None: + with gateway.scenario() as scenario: + scenario_id: Final = f"ultrafast-long-context-{uuid.uuid4().hex}" + handle: Final = register_scenario( + scenario_id, _surface_response(surface, service_tier, prompt_tokens) + ) + scenario.cleanups.callback(delete_scenario, handle) + key: Final = scenario.key() + model: Final = scenario.model( + model=f"openai/integration-ultrafast-long-context-{uuid.uuid4().hex}", + api_key=scenario_id, + api_base=handle.api_base(), + **LONG_CONTEXT_PRICING, + ) + request_path, request_body, expected_upstream_path = _surface_request(surface, scenario_id, model, service_tier) + with httpx.Client(base_url=gateway.upstream_url, trust_env=False) as upstream: + upstream.get("/__observations").raise_for_status() + response: Final = gateway.request( + "POST", + request_path, + request_body, + key=key, + ) + observations: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"] + assert response.status_code == 200, response.text + expected_input: Final = (prompt_tokens - CACHED_TOKENS) * input_rate + CACHED_TOKENS * cache_read_rate + expected_output: Final = COMPLETION_TOKENS * output_rate + expected: Final = expected_input + expected_output + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6), response.text + request_id: Final = string_value(object_value(response.json())["id"]) + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id = %s', + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["prompt_tokens"] == prompt_tokens + assert rows[0]["completion_tokens"] == COMPLETION_TOKENS + assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6) + metadata: Final = rows[0]["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + breakdown: Final = object_value(parsed["cost_breakdown"]) + assert float(breakdown["input_cost"]) == pytest.approx(expected_input, rel=1e-6) + assert float(breakdown["output_cost"]) == pytest.approx(expected_output, rel=1e-6) + assert isinstance(observations, list) + assert len(observations) == 1 + observation: Final = object_value(observations[0]) + upstream_path: Final = string_value(observation["path"]) + assert upstream_path == expected_upstream_path, upstream_path + body: Final = object_value(observation["body"]) + assert body.get("service_tier") == service_tier, body + assert not set(LONG_CONTEXT_PRICING).intersection(body), body diff --git a/tests/integration/providers/test_responses_bridge_incomplete.py b/tests/integration/providers/test_responses_bridge_incomplete.py index e700d17ea88..2252f1634e0 100644 --- a/tests/integration/providers/test_responses_bridge_incomplete.py +++ b/tests/integration/providers/test_responses_bridge_incomplete.py @@ -12,6 +12,8 @@ def test_chat_over_responses_deployment_returns_length_when_output_tokens_run_ou identity: Final = "responses-incomplete-" + uuid.uuid4().hex def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') assert request.method == "POST" and request.target == "/responses", request.target assert request.headers["authorization"] == "Bearer synthetic-openai-key" body: Final = json.loads(request.body) @@ -56,7 +58,7 @@ def test_chat_over_responses_deployment_returns_length_when_output_tokens_run_ou ) assert response.status_code == 200, response.text body: Final = response.json() - assert len(wire.drain()) == 1 + assert len(tuple(request for request in wire.drain() if request.method == "POST")) == 1 assert [choice["finish_reason"] for choice in body["choices"]] == ["length"], response.text assert body["choices"][0]["message"]["content"] == "", response.text assert body["choices"][0]["message"]["role"] == "assistant", response.text @@ -69,6 +71,8 @@ def test_messages_over_responses_deployment_with_max_tokens_1_is_clamped_to_16_i identity: Final = "responses-clamp-" + uuid.uuid4().hex def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') assert request.method == "POST" and request.target == "/responses", request.target assert request.headers["authorization"] == "Bearer synthetic-openai-key" body: Final = json.loads(request.body) @@ -132,7 +136,7 @@ def test_messages_over_responses_deployment_with_max_tokens_1_is_clamped_to_16_i ) assert response.status_code == 200, response.text body: Final = response.json() - assert len(wire.drain()) == 1 + assert len(tuple(request for request in wire.drain() if request.method == "POST")) == 1 assert body["role"] == "assistant", response.text assert body["content"] == [{"type": "text", "text": "ok"}], response.text assert body["stop_reason"] == "end_turn", response.text @@ -143,6 +147,8 @@ def test_messages_over_responses_deployment_with_max_tokens_one_reaches_openai_a identity: Final = "responses-min-tokens-" + uuid.uuid4().hex def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=b'{"object":"list","data":[]}') assert request.method == "POST" and request.target == "/responses", request.target assert request.headers["authorization"] == "Bearer synthetic-openai-key" body: Final = json.loads(request.body) @@ -185,6 +191,6 @@ def test_messages_over_responses_deployment_with_max_tokens_one_reaches_openai_a ) assert response.status_code == 200, response.text body: Final = response.json() - assert len(wire.drain()) == 1 + assert len(tuple(request for request in wire.drain() if request.method == "POST")) == 1 assert body["content"] == [{"type": "text", "text": "ok"}], response.text assert body["usage"]["input_tokens"] == 9 and body["usage"]["output_tokens"] == 1, response.text diff --git a/tests/integration/routing/test_usage_based_routing_redis_reads.py b/tests/integration/routing/test_usage_based_routing_redis_reads.py new file mode 100644 index 00000000000..f4801eb3318 --- /dev/null +++ b/tests/integration/routing/test_usage_based_routing_redis_reads.py @@ -0,0 +1,262 @@ +from __future__ import annotations + +import json +import shlex +import threading +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from datetime import UTC, datetime +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.process import owned_proxy +from integration._support.redis_process import owned_redis +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter +from redis import Redis +from redis.exceptions import TimeoutError as RedisTimeoutError + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +MONITOR_COMMAND: Final = TypeAdapter(dict[str, JsonValue]) +OPENAI_MODEL: Final = "gpt-4o-mini" +MASTER_KEY: Final = "sk-integration-usage-routing-redis-reads" +API_KEY: Final = "synthetic-usage-routing-key" +ENDPOINT_PATHS: Final = MappingProxyType( + { + "/v1/chat/completions": ("/v1/chat/completions", "/v1/chat/completions"), + "/v1/messages": ("/v1/responses", "/v1/responses"), + "/v1/responses": ("/v1/responses", "/v1/responses"), + } +) +CHAT_RESPONSE: Final = json.dumps( + { + "id": "chatcmpl_usage_routing_redis_reads", + "object": "chat.completion", + "created": 1700000000, + "model": OPENAI_MODEL, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "redis read contract"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11}, + } +).encode() +RESPONSES_RESPONSE: Final = json.dumps( + { + "id": "resp_usage_routing_redis_reads", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": OPENAI_MODEL, + "output": [ + { + "id": "msg_usage_routing_redis_reads", + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "redis read contract", "annotations": []}], + } + ], + "usage": {"input_tokens": 7, "output_tokens": 4, "total_tokens": 11}, + } +).encode() + + +def _request_object(body: bytes) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_json(body) + + +def _deployment_list( + model_name: str, api_base: str, deployment_ids: tuple[str, str] +) -> list[dict[str, JsonValue]]: + return [ + { + "model_name": model_name, + "litellm_params": { + "model": f"openai/{OPENAI_MODEL}", + "api_base": api_base, + "api_key": API_KEY, + "rpm": 1, + }, + "model_info": {"id": deployment_id}, + } + for deployment_id in deployment_ids + ] + + +def _request_payload(endpoint: str, model_name: str, marker: str) -> dict[str, JsonValue]: + if endpoint == "/v1/responses": + return {"model": model_name, "input": marker, "max_output_tokens": 16, "store": False} + return {"model": model_name, "messages": [{"role": "user", "content": marker}], "max_tokens": 16} + + +def _expected_wire_body(endpoint: str, marker: str) -> dict[str, JsonValue]: + if endpoint == "/v1/messages": + return { + "model": OPENAI_MODEL, + "input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": marker}]}], + "include": ["reasoning.encrypted_content"], + "max_output_tokens": 16, + } + if endpoint == "/v1/responses": + return {"model": OPENAI_MODEL, "input": marker, "max_output_tokens": 16, "store": False} + return {"model": OPENAI_MODEL, "messages": [{"role": "user", "content": marker}], "max_tokens": 16} + + +def _reply(request: Request) -> Reply: + if request.target == "/v1/models": + return Reply(body=json.dumps({"object": "list", "data": [{"id": OPENAI_MODEL, "object": "model"}]}).encode()) + return Reply(body=RESPONSES_RESPONSE if request.target == "/v1/responses" else CHAT_RESPONSE) + + +@contextmanager +def _capture_redis_commands(host: str, port: int) -> Iterator[SimpleQueue[str]]: + commands: Final = SimpleQueue[str]() + started: Final = threading.Event() + armed: Final = threading.Event() + stopped: Final = threading.Event() + ready_marker: Final = f"monitor-ready-{uuid.uuid4().hex}" + stop_marker: Final = f"monitor-stop-{uuid.uuid4().hex}" + + def capture() -> None: + with Redis(host=host, port=port, socket_timeout=1, decode_responses=True) as client: + with client.monitor() as monitor: + started.set() + stream: Final = iter(monitor.listen()) + while not stopped.is_set(): + try: + record: Final = MONITOR_COMMAND.validate_python(next(stream)) + except RedisTimeoutError: + continue + command: Final = record.get("command") + if not isinstance(command, str): + continue + commands.put(command) + if ready_marker in command: + armed.set() + + thread: Final = threading.Thread(target=capture, daemon=True) + thread.start() + try: + assert started.wait(timeout=5), "Redis MONITOR did not start" + with Redis(host=host, port=port, socket_timeout=1, decode_responses=True) as client: + client.set(ready_marker, "ready", ex=1) + assert armed.wait(timeout=5), "Redis MONITOR did not capture its readiness command" + yield commands + finally: + stopped.set() + with Redis(host=host, port=port, socket_timeout=1, decode_responses=True) as client: + client.set(stop_marker, "stop", ex=1) + thread.join(timeout=5) + assert not thread.is_alive(), "Redis MONITOR thread survived cleanup" + + +def _drain_mgets(commands: SimpleQueue[str]) -> tuple[tuple[str, tuple[str, ...]], ...]: + captured: Final = tuple(commands.get_nowait() for _ in range(commands.qsize())) + parsed: Final = tuple((line, tuple(shlex.split(line))) for line in captured) + return tuple((line, arguments) for line, arguments in parsed if arguments and arguments[0] == "MGET") + + +@pytest.mark.parametrize( + "endpoint", + ("/v1/chat/completions", "/v1/messages", "/v1/responses"), + ids=("chat-completions", "messages", "responses"), +) +def test_proxy_usage_routing_reads_cooldown_tpm_then_rpm_from_redis( + endpoint: str, tmp_path: Path +) -> None: + with owned_redis(tmp_path) as cache, wire_server(_reply) as wire: + run_id: Final = uuid.uuid4().hex + model_name: Final = f"usage-redis-{run_id}" + deployment_ids: Final = (f"dep-a-{run_id[:8]}", f"dep-b-{run_id[:8]}") + configuration: Final = JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + config: Final = { + **configuration, + "model_list": _deployment_list(model_name, f"{wire.url}/v1", deployment_ids), + "router_settings": { + "routing_strategy": "usage-based-routing-v2", + "redis_host": cache.host, + "redis_port": cache.port, + }, + } + config_path: Final = tmp_path / "usage-routing.yaml" + config_path.write_text(yaml.safe_dump(config)) + with httpx.Client(base_url=wire.url, timeout=15, trust_env=False) as bootstrap_client: + bootstrap: Final = Gateway(bootstrap_client, MASTER_KEY, wire.url) + with owned_proxy( + bootstrap, + tmp_path, + {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)}, + config=config_path, + ) as candidate: + eventually( + lambda: wire.received.qsize(), + lambda received: received >= len(deployment_ids), + seconds=15, + ) + wire.drain() + eventually( + lambda: datetime.now(UTC), + lambda current: current.second < 40, + seconds=65, + ) + minute: Final = datetime.now(UTC).strftime("%H-%M") + markers: Final = tuple(f"{run_id}-{index}" for index in range(3)) + payloads: Final = tuple(_request_payload(endpoint, model_name, marker) for marker in markers) + request_headers: Final = ( + {"anthropic-version": "2023-06-01"} if endpoint == "/v1/messages" else {} + ) + with _capture_redis_commands(cache.host, cache.port) as commands: + responses: Final = tuple( + candidate.request("POST", endpoint, payload, headers=request_headers) for payload in payloads + ) + assert tuple(response.status_code for response in responses) == (200, 200, 429), [ + response.text for response in responses + ] + assert "No deployments available" in responses[2].text + served_ids: Final = tuple(response.headers["x-litellm-model-id"] for response in responses[:2]) + assert set(served_ids) == set(deployment_ids), served_ids + rpm_keys: Final = tuple( + f"{deployment_id}:openai/{OPENAI_MODEL}:rpm:{minute}" for deployment_id in deployment_ids + ) + with Redis(host=cache.host, port=cache.port, decode_responses=True) as redis_client: + rpm_values: Final = eventually( + lambda: tuple(redis_client.get(key) for key in rpm_keys), + lambda values: values == ("1", "1"), + seconds=15, + ) + assert rpm_values == ("1", "1") + received: Final = wire.drain() + assert len(received) == 2 + assert tuple(request.method for request in received) == ("POST", "POST") + assert tuple(request.target for request in received) == ENDPOINT_PATHS[endpoint] + observed_bodies: Final = tuple(_request_object(request.body) for request in received) + expected_bodies: Final = tuple(_expected_wire_body(endpoint, marker) for marker in markers[:2]) + assert observed_bodies == expected_bodies, observed_bodies + expected_mget: Final = ( + "MGET", + f"deployment:{deployment_ids[0]}:cooldown", + f"deployment:{deployment_ids[1]}:cooldown", + f"{deployment_ids[0]}:openai/{OPENAI_MODEL}:tpm:{minute}", + f"{deployment_ids[1]}:openai/{OPENAI_MODEL}:tpm:{minute}", + *( + f"{deployment_id}:openai/{OPENAI_MODEL}:rpm:{minute}" + for deployment_id in deployment_ids + ), + ) + mgets: Final = _drain_mgets(commands) + assert any(arguments == expected_mget for _, arguments in mgets), mgets + raw_mgets: Final = tuple(line for line, _ in mgets) + print(f"proxy {endpoint} MGETs: {raw_mgets}") diff --git a/tests/integration/sdk/test_usage_based_routing_sdk_redis_reads.py b/tests/integration/sdk/test_usage_based_routing_sdk_redis_reads.py new file mode 100644 index 00000000000..da9a8309c8d --- /dev/null +++ b/tests/integration/sdk/test_usage_based_routing_sdk_redis_reads.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +import asyncio +import json +import shlex +import threading +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from datetime import UTC, datetime +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import litellm +import pytest +from integration._support.client import eventually +from integration._support.redis_process import owned_redis +from integration._support.wire import Reply, Request, wire_server +from litellm import Router +from pydantic import JsonValue, TypeAdapter +from redis import Redis +from redis.exceptions import TimeoutError as RedisTimeoutError + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +MONITOR_COMMAND: Final = TypeAdapter(dict[str, JsonValue]) +OPENAI_MODEL: Final = "gpt-4o-mini" +API_KEY: Final = "synthetic-usage-routing-key" +CHAT_RESPONSE: Final = json.dumps( + { + "id": "chatcmpl_usage_routing_sdk_redis_reads", + "object": "chat.completion", + "created": 1700000000, + "model": OPENAI_MODEL, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "redis read contract"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 7, "completion_tokens": 4, "total_tokens": 11}, + } +).encode() + + +def _request_object(body: bytes) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_json(body) + + +def _deployment_list( + model_name: str, api_base: str, deployment_ids: tuple[str, str] +) -> list[dict[str, JsonValue]]: + return [ + { + "model_name": model_name, + "litellm_params": { + "model": f"openai/{OPENAI_MODEL}", + "api_base": api_base, + "api_key": API_KEY, + "rpm": 1, + }, + "model_info": {"id": deployment_id}, + } + for deployment_id in deployment_ids + ] + + +def _reply(request: Request) -> Reply: + return Reply(body=CHAT_RESPONSE) + + +@contextmanager +def _capture_redis_commands(host: str, port: int) -> Iterator[SimpleQueue[str]]: + commands: Final = SimpleQueue[str]() + started: Final = threading.Event() + armed: Final = threading.Event() + stopped: Final = threading.Event() + ready_marker: Final = f"monitor-ready-{uuid.uuid4().hex}" + stop_marker: Final = f"monitor-stop-{uuid.uuid4().hex}" + + def capture() -> None: + with Redis(host=host, port=port, socket_timeout=1, decode_responses=True) as client: + with client.monitor() as monitor: + started.set() + stream: Final = iter(monitor.listen()) + while not stopped.is_set(): + try: + record: Final = MONITOR_COMMAND.validate_python(next(stream)) + except RedisTimeoutError: + continue + command: Final = record.get("command") + if not isinstance(command, str): + continue + commands.put(command) + if ready_marker in command: + armed.set() + + thread: Final = threading.Thread(target=capture, daemon=True) + thread.start() + try: + assert started.wait(timeout=5), "Redis MONITOR did not start" + with Redis(host=host, port=port, socket_timeout=1, decode_responses=True) as client: + client.set(ready_marker, "ready", ex=1) + assert armed.wait(timeout=5), "Redis MONITOR did not capture its readiness command" + yield commands + finally: + stopped.set() + with Redis(host=host, port=port, socket_timeout=1, decode_responses=True) as client: + client.set(stop_marker, "stop", ex=1) + thread.join(timeout=5) + assert not thread.is_alive(), "Redis MONITOR thread survived cleanup" + + +def _drain_mgets(commands: SimpleQueue[str]) -> tuple[tuple[str, tuple[str, ...]], ...]: + captured: Final = tuple(commands.get_nowait() for _ in range(commands.qsize())) + parsed: Final = tuple((line, tuple(shlex.split(line))) for line in captured) + return tuple((line, arguments) for line, arguments in parsed if arguments and arguments[0] == "MGET") + + +def _model_id(response: object) -> str: + response_params: Final = getattr(response, "_hidden_params") + hidden_params: Final = JSON_OBJECT.validate_python(response_params) + model_id: Final = hidden_params.get("model_id") + assert isinstance(model_id, str), hidden_params + return model_id + + +async def _exercise_router(router: Router, model_name: str, markers: tuple[str, str, str]) -> tuple[str, str]: + first: Final = await router.acompletion( + model=model_name, messages=[{"role": "user", "content": markers[0]}], max_tokens=8 + ) + second: Final = await router.acompletion( + model=model_name, messages=[{"role": "user", "content": markers[1]}], max_tokens=8 + ) + with pytest.raises(litellm.RateLimitError, match="No deployments available"): + await router.acompletion( + model=model_name, messages=[{"role": "user", "content": markers[2]}], max_tokens=8 + ) + return _model_id(first), _model_id(second) + + +def test_sdk_usage_routing_reads_tpm_then_rpm_from_redis(tmp_path: Path) -> None: + with owned_redis(tmp_path) as cache, wire_server(_reply) as wire: + run_id: Final = uuid.uuid4().hex + model_name: Final = f"usage-redis-{run_id}" + deployment_ids: Final = (f"dep-a-{run_id[:8]}", f"dep-b-{run_id[:8]}") + router: Final = Router( + model_list=_deployment_list(model_name, f"{wire.url}/v1", deployment_ids), + routing_strategy="usage-based-routing-v2", + redis_host=cache.host, + redis_port=cache.port, + ) + try: + eventually( + lambda: datetime.now(UTC), + lambda current: current.second < 40, + seconds=65, + ) + minute: Final = datetime.now(UTC).strftime("%H-%M") + markers: Final = tuple(f"{run_id}-{index}" for index in range(3)) + with _capture_redis_commands(cache.host, cache.port) as commands: + served_ids: Final = asyncio.run(_exercise_router(router, model_name, markers)) + assert set(served_ids) == set(deployment_ids), served_ids + received: Final = wire.drain() + assert len(received) == 2 + assert tuple(request.method for request in received) == ("POST", "POST") + assert tuple(request.target for request in received) == ("/v1/chat/completions",) * 2 + observed_bodies: Final = tuple(_request_object(request.body) for request in received) + expected_bodies: Final = tuple( + { + "model": OPENAI_MODEL, + "messages": [{"role": "user", "content": marker}], + "max_tokens": 8, + } + for marker in markers[:2] + ) + assert observed_bodies == expected_bodies, observed_bodies + expected_mget: Final = ( + "MGET", + f"deployment:{deployment_ids[0]}:cooldown", + f"deployment:{deployment_ids[1]}:cooldown", + *(f"{deployment_id}:openai/{OPENAI_MODEL}:tpm:{minute}" for deployment_id in deployment_ids), + *(f"{deployment_id}:openai/{OPENAI_MODEL}:rpm:{minute}" for deployment_id in deployment_ids), + ) + mgets: Final = _drain_mgets(commands) + assert any(arguments == expected_mget for _, arguments in mgets), mgets + raw_mgets: Final = tuple(line for line, _ in mgets) + print(f"sdk MGETs: {raw_mgets}") + finally: + router.reset() diff --git a/tests/integration/spend/test_batch_completion_accounting.py b/tests/integration/spend/test_batch_completion_accounting.py index 0cbeda934f6..4ecab10f942 100644 --- a/tests/integration/spend/test_batch_completion_accounting.py +++ b/tests/integration/spend/test_batch_completion_accounting.py @@ -2,11 +2,12 @@ from __future__ import annotations import json import uuid +from datetime import datetime, timedelta, timezone from hashlib import sha256 from typing import Final import pytest -from integration._support.client import JSON_OBJECT, Gateway, eventually, string_value +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value from integration._support.database import read_rows from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse @@ -116,6 +117,28 @@ def _batch_routes(model: str) -> RoutedResponse: ) +def _team_day_endpoints(gateway: Gateway, team: str, start_date: str, end_date: str) -> dict[str, object] | None: + response: Final = gateway.request( + "GET", + "/team/daily/activity", + params={"team_ids": team, "start_date": start_date, "end_date": end_date}, + ) + if response.status_code != 200: + return None + days: Final = response.json()["results"] + if not days: + return None + return object_value(object_value(object_value(days[0])["breakdown"])["endpoints"]) + + +def _batches_total_tokens(endpoints: dict[str, object] | None) -> int | None: + if endpoints is None or "/batches" not in endpoints: + return None + metrics: Final = object_value(object_value(endpoints["/batches"])["metrics"]) + total_tokens: Final = metrics["total_tokens"] + return int(total_tokens) if isinstance(total_tokens, (int, float, str)) else None + + def _input_file(model: str) -> bytes: return ( "\n".join( @@ -200,3 +223,79 @@ def test_completed_batch_spend_row_records_reasoning_tokens_and_error_file_failu "reasoning_tokens": reasoning_tokens, "text_tokens": completion_tokens - reasoning_tokens, }, json.dumps(metadata) + + +INPUT_COST_PER_TOKEN: Final = 0.001 +OUTPUT_COST_PER_TOKEN: Final = 0.002 +BATCH_PROMPT_TOKENS: Final = FIRST_LINE["prompt_tokens"] + SECOND_LINE["prompt_tokens"] +BATCH_COMPLETION_TOKENS: Final = FIRST_LINE["completion_tokens"] + SECOND_LINE["completion_tokens"] +BATCH_SPEND: Final = (BATCH_PROMPT_TOKENS * INPUT_COST_PER_TOKEN + BATCH_COMPLETION_TOKENS * OUTPUT_COST_PER_TOKEN) / 2 + + +def test_completed_batch_spend_lands_under_batches_in_team_endpoint_activity(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + scenario_id: Final = f"batch-endpoint-{uuid.uuid4().hex[:12]}" + handle: Final = register_scenario(scenario_id, _batch_routes("gpt-4o-mini")) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model( + api_base=handle.api_base(), + input_cost_per_token=INPUT_COST_PER_TOKEN, + output_cost_per_token=OUTPUT_COST_PER_TOKEN, + ) + team: Final = scenario.team(models=[model]) + key: Final = scenario.key(team_id=team, models=[model]) + file_response: Final = gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "model": model}, + {"file": ("in.jsonl", _input_file(model), "application/jsonl")}, + key=key, + ) + assert file_response.status_code == 200, file_response.text + batch_response: Final = gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": string_value(JSON_OBJECT.validate_json(file_response.content)["id"]), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": model, + }, + key=key, + ) + assert batch_response.status_code == 200, batch_response.text + batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) + retrieval: Final = gateway.request("GET", f"/v1/batches/{batch_id}", key=key) + assert retrieval.status_code == 200, retrieval.text + assert retrieval.json()["status"] == "completed", retrieval.text + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens, total_tokens FROM "LiteLLM_SpendLogs" ' + "WHERE api_key=%s AND call_type='aretrieve_batch'", + (sha256(key.encode()).hexdigest(),), + ), + lambda values: len(values) == 1, + seconds=70, + ) + row: Final = rows[0] + assert float(row["spend"]) == pytest.approx(BATCH_SPEND), dict(row) + assert (row["prompt_tokens"], row["completion_tokens"]) == ( + BATCH_PROMPT_TOKENS, + BATCH_COMPLETION_TOKENS, + ), dict(row) + today: Final = datetime.now(timezone.utc) + endpoints: Final = eventually( + lambda: _team_day_endpoints( + gateway, + team, + (today - timedelta(days=1)).strftime("%Y-%m-%d"), + (today + timedelta(days=1)).strftime("%Y-%m-%d"), + ), + lambda value: _batches_total_tokens(value) == BATCH_PROMPT_TOKENS + BATCH_COMPLETION_TOKENS, + seconds=70, + return_last_on_timeout=True, + ) + assert endpoints is not None, "team daily activity returned no endpoint breakdown for the day" + assert set(endpoints) == {"/batches"}, endpoints + endpoint_metrics: Final = object_value(object_value(endpoints["/batches"])["metrics"]) + assert float(endpoint_metrics["spend"]) == pytest.approx(BATCH_SPEND), endpoints + assert endpoint_metrics["total_tokens"] == BATCH_PROMPT_TOKENS + BATCH_COMPLETION_TOKENS, endpoints diff --git a/tests/integration/spend/test_daily_activity_key_alias_probes.py b/tests/integration/spend/test_daily_activity_key_alias_probes.py new file mode 100644 index 00000000000..8d9b8616435 --- /dev/null +++ b/tests/integration/spend/test_daily_activity_key_alias_probes.py @@ -0,0 +1,490 @@ +import time +import uuid +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.daily_activity import ( + AGGREGATED_USER_ACTIVITY, + DAY, + ROUTES, + SPEND_LOGS_TABLE, + USER_SPEND, + Route, + SpendLogRow, + activity_of_key, + assert_key_reported, + daily_rows, + digest_no_key_table_holds, + key_metadata, + locked_table, + named_row, + nameless_rows, + records_of_key, + seeded_metrics, + seeded_row, + spend_logs_of_key, + started_at, + user_row, + user_with_an_email, +) +from integration._support.database import read_rows, scratch_database +from integration._support.process import OwnedProxy, owned_proxy_process +from pydantic import JsonValue + +DAY_OUTSIDE_THE_WINDOW: Final = "2026-02-10" +GIVES_UP_WITHIN_SECONDS: Final = 10 +CONCURRENT_READS: Final = 20 +CACHED_MISS_CLEARS_WITHIN_SECONDS: Final = 45 +ALIAS_OF_ONE_SPEND_LOG: Final = ( + "SELECT metadata->>'user_api_key_alias' AS alias FROM \"LiteLLM_SpendLogs\" WHERE request_id = %s" +) + + +def _alias() -> str: + return f"integration-alias-{uuid.uuid4().hex}" + + +def _named_between_fifty_and_fifty(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(50), named_row(50, alias), *nameless_rows(50, 51)) + + +def _oldest_named(alias: str) -> tuple[SpendLogRow, ...]: + return (named_row(0, alias), *nameless_rows(150, 1)) + + +def _newest_named(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(150), named_row(150, alias)) + + +def _both_edges_named(alias: str) -> tuple[SpendLogRow, ...]: + return (named_row(0, alias), *nameless_rows(150, 1), named_row(151, alias)) + + +def _named_after_one_hundred(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(100), named_row(100, alias), *nameless_rows(99, 101)) + + +def _named_after_ninety_nine(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(99), named_row(99, alias), *nameless_rows(100, 100)) + + +def _named_only_in_the_middle(alias: str) -> tuple[SpendLogRow, ...]: + return (*nameless_rows(100), named_row(100, alias), *nameless_rows(100, 101)) + + +def _renamed_and_renamed_back(alias: str, other: str) -> tuple[SpendLogRow, ...]: + return ( + named_row(0, alias), + *nameless_rows(100, 1), + named_row(101, other), + *nameless_rows(100, 102), + named_row(202, alias), + ) + + +def _team_in_the_column(team: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {}, team_id=team) + + +def _team_in_the_metadata(team: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {"user_api_key_team_id": team}) + + +def _user_in_the_column(user: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {}, user=user) + + +def _user_in_the_metadata(user: str) -> SpendLogRow: + return SpendLogRow(started_at(0), {"user_api_key_user_id": user}) + + +def _activity_on_route(gateway: Gateway, route: Route, api_key: str, entity: str) -> httpx.Response: + filters: Final = {} if route.entity_filter is None else {route.entity_filter: entity} + return activity_of_key(gateway, route.path, api_key, **filters) + + +def _reported_aliases(response: httpx.Response, api_key: str) -> tuple[JsonValue, ...]: + if response.status_code != 200: + return () + return tuple( + object_value(object_value(record)["metadata"])["key_alias"] + for record in records_of_key(object_value(response.json()), api_key) + ) + + +def _names_the_key(api_key: str, alias: str) -> Callable[[httpx.Response], bool]: + def names(response: httpx.Response) -> bool: + reported: Final = _reported_aliases(response, api_key) + return bool(reported) and frozenset(reported) == frozenset((alias,)) + + return names + + +@contextmanager +def _proxy_on(gateway: Gateway, directory: Path, database_url: str, *, workers: int = 1) -> Iterator[OwnedProxy]: + with owned_proxy_process( + gateway, + directory, + {"DATABASE_URL": database_url}, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=workers, + ) as owned: + yield owned + + +def _owner_on(candidate: Gateway) -> tuple[str, str]: + owner: Final = f"integration-{uuid.uuid4().hex}" + email: Final = f"{owner}@example.com" + candidate.post("/user/new", {"user_id": owner, "user_email": email, "auto_create_key": False}) + return owner, email + + +@pytest.mark.parametrize("route", ROUTES, ids=lambda route: route.path.strip("/").replace("/", "_")) +def test_alias_named_only_by_a_spend_log_is_reported_on_every_daily_activity_route( + gateway: Gateway, route: Route +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + entity: Final = f"integration-entity-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + entity_rows: Final = ( + () if route.table == USER_SPEND else (seeded_row(route.table, route.entity_column, entity, api_key, DAY),) + ) + filters: Final = {} if route.entity_filter is None else {route.entity_filter: entity} + with ( + daily_rows((user_row(owner, api_key, DAY), *entity_rows)), + spend_logs_of_key(api_key, (named_row(0, alias),)), + ): + assert_key_reported( + activity_of_key(gateway, route.path, api_key, **filters), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "layout", + ( + pytest.param(_named_between_fifty_and_fifty, id="named_between_50_and_50_nameless"), + pytest.param(_oldest_named, id="oldest_named_150_nameless_newer"), + pytest.param(_newest_named, id="newest_named_150_nameless_older"), + pytest.param(_both_edges_named, id="both_edges_named_150_nameless_between"), + pytest.param(_named_after_one_hundred, id="100_nameless_named_99_nameless"), + pytest.param(_named_after_ninety_nine, id="99_nameless_named_100_nameless"), + ), +) +def test_alias_on_an_edge_of_the_window_is_reported_whatever_surrounds_it( + gateway: Gateway, layout: Callable[[str], tuple[SpendLogRow, ...]] +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, layout(alias)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +def test_alias_named_only_in_the_middle_of_two_hundred_nameless_rows_is_not_picked_up(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with ( + daily_rows((user_row(owner, api_key, DAY),)), + spend_logs_of_key(api_key, _named_only_in_the_middle(_alias())), + ): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=owner, email=email), + seeded_metrics(1), + ) + + +def test_key_renamed_and_renamed_back_is_reported_with_the_alias_on_both_edges(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + rows: Final = _renamed_and_renamed_back(alias, _alias()) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "spend_log_of_team", + ( + pytest.param(_team_in_the_column, id="team_id_column"), + pytest.param(_team_in_the_metadata, id="team_id_in_metadata"), + ), +) +def test_team_named_only_by_a_spend_log_is_reported_next_to_the_daily_owner( + gateway: Gateway, spend_log_of_team: Callable[[str], SpendLogRow] +) -> None: + api_key: Final = digest_no_key_table_holds() + team: Final = f"integration-team-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (spend_log_of_team(team),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(team=team, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "spend_log_of_user", + ( + pytest.param(_user_in_the_column, id="user_column"), + pytest.param(_user_in_the_metadata, id="user_id_in_metadata"), + ), +) +def test_user_named_by_a_spend_log_beats_the_owner_the_daily_rows_name( + gateway: Gateway, spend_log_of_user: Callable[[str], SpendLogRow] +) -> None: + api_key: Final = digest_no_key_table_holds() + with gateway.scenario() as scenario: + daily_owner, _ = user_with_an_email(scenario) + log_user, log_email = user_with_an_email(scenario) + with ( + daily_rows((user_row(daily_owner, api_key, DAY),)), + spend_logs_of_key(api_key, (spend_log_of_user(log_user),)), + ): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=log_user, email=log_email), + seeded_metrics(1), + ) + + +def test_hashed_jwt_digest_is_named_by_its_spend_log(gateway: Gateway) -> None: + api_key: Final = f"hashed-jwt-{sha256(uuid.uuid4().bytes).hexdigest()}" + alias: Final = _alias() + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (named_row(0, alias),)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + ("started", "inside_the_window"), + ( + pytest.param("2026-02-01 23:59:59", False, id="second_before_the_window"), + pytest.param("2026-02-02 00:00:00", True, id="first_second_of_the_window"), + pytest.param("2026-02-04 23:59:59", True, id="last_second_of_the_window"), + pytest.param("2026-02-05 00:00:00", False, id="first_second_after_the_window"), + ), +) +def test_spend_log_names_the_key_only_from_one_day_before_to_two_days_after_the_read( + gateway: Gateway, started: str, inside_the_window: bool +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + row: Final = SpendLogRow(started, {"user_api_key_alias": alias}) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (row,)): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias if inside_the_window else None, user=owner, email=email), + seeded_metrics(1), + ) + + +def test_two_aliases_on_the_two_edges_leave_the_key_unnamed(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + rows: Final = (named_row(0, _alias()), *nameless_rows(150, 1), named_row(151, _alias())) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "unnamed_rows", + ( + pytest.param((SpendLogRow(started_at(0), {"user_api_key_alias": ""}),), id="empty_string_alias"), + pytest.param( + (SpendLogRow(started_at(0), ["x"]), SpendLogRow(started_at(1), "x")), id="array_then_string_metadata" + ), + ), +) +def test_rows_without_a_usable_alias_do_not_hide_the_named_row_after_them( + gateway: Gateway, unnamed_rows: tuple[SpendLogRow, ...] +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + rows: Final = (*unnamed_rows, named_row(len(unnamed_rows), alias)) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows): + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=alias, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.parametrize( + "stored_alias", + ( + pytest.param(123, id="json_int"), + pytest.param(["a"], id="json_list"), + pytest.param("a" * 5000, id="five_kb_string"), + ), +) +def test_alias_of_an_unexpected_shape_is_reported_as_postgres_renders_it( + gateway: Gateway, stored_alias: JsonValue +) -> None: + api_key: Final = digest_no_key_table_holds() + row: Final = SpendLogRow(started_at(0), {"user_api_key_alias": stored_alias}) + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (row,)) as request_ids: + rendered: Final = read_rows(ALIAS_OF_ONE_SPEND_LOG, (request_ids[0],))[0]["alias"] + assert isinstance(rendered, str) and rendered, rendered + assert_key_reported( + activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key), + api_key, + DAY, + key_metadata(alias=rendered, user=owner, email=email), + seeded_metrics(1), + ) + + +@pytest.mark.timeout(300) +def test_alias_found_once_is_served_from_the_cache_for_the_same_window_only(gateway: Gateway, tmp_path: Path) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url) as owned: + owner, email = _owner_on(owned.gateway) + rows: Final = (user_row(owner, api_key, DAY), user_row(owner, api_key, DAY_OUTSIDE_THE_WINDOW)) + with daily_rows(rows, database_url=database_url): + with spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url): + first: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + cached: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + other_window: Final = owned.gateway.request( + "GET", + AGGREGATED_USER_ACTIVITY, + params={"start_date": DAY_OUTSIDE_THE_WINDOW, "end_date": DAY_OUTSIDE_THE_WINDOW, "api_key": api_key}, + ) + named: Final = key_metadata(alias=alias, user=owner, email=email) + assert_key_reported(first, api_key, DAY, named, seeded_metrics(1)) + assert_key_reported(cached, api_key, DAY, named, seeded_metrics(1)) + assert_key_reported( + other_window, api_key, DAY_OUTSIDE_THE_WINDOW, key_metadata(user=owner, email=email), seeded_metrics(1) + ) + + +@pytest.mark.timeout(300) +def test_alias_logged_after_a_cached_miss_shows_once_the_miss_expires(gateway: Gateway, tmp_path: Path) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url) as owned: + owner, email = _owner_on(owned.gateway) + with daily_rows((user_row(owner, api_key, DAY),), database_url=database_url): + missed: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + with spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url): + named: Final = eventually( + lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key), + _names_the_key(api_key, alias), + seconds=CACHED_MISS_CLEARS_WITHIN_SECONDS, + ) + assert_key_reported(missed, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) + assert_key_reported(named, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1)) + + +@pytest.mark.timeout(300) +def test_alias_lookup_gives_up_while_spend_logs_are_locked_and_answers_once_they_are_not( + gateway: Gateway, tmp_path: Path +) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url, workers=2) as owned: + owner, email = _owner_on(owned.gateway) + with ( + daily_rows((user_row(owner, api_key, DAY),), database_url=database_url), + spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url), + ): + with locked_table(SPEND_LOGS_TABLE, database_url=database_url): + started: Final = time.monotonic() + locked: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key) + waited: Final = time.monotonic() - started + unlocked: Final = eventually( + lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key), + _names_the_key(api_key, alias), + seconds=CACHED_MISS_CLEARS_WITHIN_SECONDS, + ) + assert waited < GIVES_UP_WITHIN_SECONDS, waited + assert_key_reported(locked, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1)) + assert_key_reported(unlocked, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1)) + + +def test_concurrent_reads_over_every_route_all_name_a_fresh_key(gateway: Gateway) -> None: + api_key: Final = digest_no_key_table_holds() + alias: Final = _alias() + entity: Final = f"integration-entity-{uuid.uuid4().hex}" + entity_columns: Final = {route.table: route.entity_column for route in ROUTES if route.table != USER_SPEND} + with gateway.scenario() as scenario: + owner, email = user_with_an_email(scenario) + rows: Final = ( + user_row(owner, api_key, DAY), + *(seeded_row(table, column, entity, api_key, DAY) for table, column in entity_columns.items()), + ) + with ( + daily_rows(rows), + spend_logs_of_key(api_key, (named_row(0, alias),)), + ThreadPoolExecutor(CONCURRENT_READS) as pool, + ): + reads: Final = tuple( + pool.submit(_activity_on_route, gateway, ROUTES[index % len(ROUTES)], api_key, entity) + for index in range(CONCURRENT_READS) + ) + responses: Final = tuple(read.result() for read in reads) + for response in responses: + assert_key_reported( + response, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1) + ) diff --git a/tests/integration/spend/test_daily_activity_key_owner_traffic.py b/tests/integration/spend/test_daily_activity_key_owner_traffic.py index e0ec1310485..b8f113aca49 100644 --- a/tests/integration/spend/test_daily_activity_key_owner_traffic.py +++ b/tests/integration/spend/test_daily_activity_key_owner_traffic.py @@ -12,7 +12,7 @@ from typing import Final import httpx import pytest -from integration._support.client import Gateway, Scenario, eventually +from integration._support.client import Gateway, Scenario, eventually, string_value from integration._support.daily_activity import ( AGGREGATED_USER_ACTIVITY, DAY, @@ -25,6 +25,7 @@ from integration._support.daily_activity import ( daily_rows, key_metadata, key_no_key_table_holds, + purge_key_from_the_key_tables, seeded_metrics, seeded_row, user_row, @@ -42,6 +43,10 @@ REQUESTS_OF_KEY: Final = ( 'SELECT COALESCE(SUM(api_requests), 0)::int AS requests FROM "LiteLLM_DailyUserSpend" ' "WHERE api_key=%s AND user_id=%s" ) +NAMED_SPEND_LOGS_OF_KEY: Final = ( + 'SELECT COUNT(*)::int AS named FROM "LiteLLM_SpendLogs" ' + "WHERE api_key=%s AND NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL" +) UNIFIED_ENDPOINTS: Final = ("/v1/chat/completions", "/v1/messages", "/v1/responses") REQUESTS_OF_A_BURST: Final = 21 READS_DURING_A_BURST: Final = 30 @@ -210,6 +215,14 @@ def _wait_for_requests(api_key: str, user: str, requests: int) -> None: ) +def _wait_for_named_spend_logs(api_key: str, requests: int) -> None: + eventually( + lambda: read_rows(NAMED_SPEND_LOGS_OF_KEY, (api_key,)), + lambda rows: rows[0]["named"] == requests, + seconds=70, + ) + + def _cli_session_token(user: str, team: str) -> str: cli_user: Final = LiteLLM_UserTable(user_id=user, user_role="internal_user", teams=[team], models=[]) return ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info=cli_user, team_id=team, team_alias="cli-team") @@ -244,6 +257,39 @@ def test_key_used_on_every_unified_endpoint_is_reported_with_its_own_alias_and_u ) +def test_key_purged_from_the_key_tables_is_reported_with_the_alias_its_spend_logs_name(gateway: Gateway) -> None: + prompts: Final = (_prompt(), _prompt(), _prompt()) + with wire_server(_provider) as wire, gateway.scenario() as scenario: + model: Final = _priced_model(scenario, wire.url) + owner, email = user_with_an_email(scenario) + alias: Final = f"integration-alias-{uuid.uuid4().hex}" + generated: Final = gateway.post("/key/generate", {"user_id": owner, "key_alias": alias, "models": [model]}) + key: Final = string_value(generated["key"]) + stored: Final = sha256(key.encode()).hexdigest() + try: + answers: Final = tuple( + gateway.request("POST", endpoint, _request_body(endpoint, model, prompt), key=key) + for endpoint, prompt in zip(UNIFIED_ENDPOINTS, prompts, strict=True) + ) + assert [answer.status_code for answer in answers] == [200, 200, 200], [answer.text for answer in answers] + received: Final = _sent_for_callers(wire.drain()) + assert [request.target for request in received] == [ + "/v1/chat/completions", + "/v1/responses", + "/v1/responses", + ] + _wait_for_requests(stored, owner, 3) + _wait_for_named_spend_logs(stored, 3) + finally: + purge_key_from_the_key_tables(stored) + assert_key_owner_and_totals( + _activity_around_today(gateway, stored), + stored, + key_metadata(alias=alias, user=owner, email=email, exists=False), + _totals_of_requests(3), + ) + + def test_cli_session_spend_is_reported_with_the_user_and_team_of_the_session( gateway: Gateway, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/local_testing/test_alangfuse.py b/tests/local_testing/test_alangfuse.py index ec80724d3ba..bc388aebecf 100644 --- a/tests/local_testing/test_alangfuse.py +++ b/tests/local_testing/test_alangfuse.py @@ -690,7 +690,7 @@ def test_langfuse_logging_tool_calling(): ] response = litellm.completion( - model="gpt-3.5-turbo-1106", + model="gpt-6-luna", messages=messages, tools=tools, tool_choice="auto", # auto is default, but we'll be explicit @@ -698,6 +698,8 @@ def test_langfuse_logging_tool_calling(): print("\nLLM Response1:\n", response) response_message = response.choices[0].message tool_calls = response.choices[0].message.tool_calls + assert response.choices[0].message.tool_calls + assert all(call.function.name == "get_current_weather" for call in response.choices[0].message.tool_calls) # test_langfuse_logging_tool_calling() diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index 2d79f8a6af6..4c216cc75fb 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -39,7 +39,7 @@ def get_current_weather(location, unit="fahrenheit"): @pytest.mark.parametrize( "model", [ - "gpt-3.5-turbo-1106", + "gpt-6-luna", "mistral/mistral-large-latest", "claude-haiku-4-5-20251001", "gemini/gemini-2.5-flash-lite", @@ -386,7 +386,7 @@ def test_parallel_function_call_stream(): } ] response = litellm.completion( - model="gpt-3.5-turbo-1106", + model="gpt-6-luna", messages=messages, tools=tools, stream=True, @@ -435,7 +435,7 @@ def test_parallel_function_call_stream(): ) # extend conversation with function response print(f"messages: {messages}") second_response = litellm.completion( - model="gpt-3.5-turbo-1106", messages=messages, temperature=0.2, seed=22 + model="gpt-6-luna", messages=messages, temperature=0.2, seed=22, reasoning_effort="none" ) # get a new response from the model where it can see the function response print("second response\n", second_response) return second_response diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index 79f6739a423..8e24dc23398 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -1,14 +1,18 @@ # What is this? ## Unit testing for the 'get_model_info()' function import os +import re +from collections.abc import Collection, Mapping -from typing import List, Dict, Any +from typing import List, Dict, Any, Final, Literal import pytest import litellm from litellm import get_model_info +from litellm.llms.bedrock.common_utils import BedrockModelInfo +from litellm.types.utils import ModelInfoBase from litellm.utils import _invalidate_model_cost_lowercase_map from unittest.mock import MagicMock, patch @@ -116,26 +120,31 @@ def test_get_model_info_ft_model_with_provider_prefix(): def _enforce_bedrock_converse_models( - model_cost: List[Dict[str, Any]], whitelist_models: List[str] -): + model_cost: Mapping[str, ModelInfoBase], whitelist_models: Collection[str] +) -> None: """ - Assert all new bedrock chat models are added as `bedrock_converse` unless explicitly whitelisted. + Assert unlisted Bedrock chat models declare or inherit Converse routing. """ # Check for unwhitelisted models - for model, info in litellm.model_cost.items(): + for model, info in model_cost.items(): if ( info["litellm_provider"] == "bedrock" and info["mode"] == "chat" and model not in whitelist_models + and not ( + (base_model := BedrockModelInfo.get_base_model(model)) != model + and model_cost.get(base_model, {}).get("litellm_provider") == "bedrock_converse" + and BedrockModelInfo.get_bedrock_route(model) == "converse" + ) ): raise AssertionError( - f"New bedrock chat model detected: {model}. Please set `litellm_provider='bedrock_converse'` for this model." + f"Unlisted Bedrock chat model does not route to Converse: {model}" ) def test_model_info_bedrock_converse(monkeypatch): """ - Assert all new bedrock chat models are added as `bedrock_converse` unless explicitly whitelisted. + Assert unlisted Bedrock chat models declare or inherit Converse routing. This ensures they are automatically routed to the converse endpoint. """ @@ -173,7 +182,7 @@ def test_model_info_bedrock_converse_enforcement(monkeypatch): whitelist_models = [line.strip() for line in file.readlines()] # Check for unwhitelisted models - with pytest.raises(AssertionError): + with pytest.raises(AssertionError, match=r"fake\.bedrock-chat-model"): _enforce_bedrock_converse_models( model_cost=litellm.model_cost, whitelist_models=whitelist_models ) @@ -181,6 +190,27 @@ def test_model_info_bedrock_converse_enforcement(monkeypatch): pytest.skip("whitelisted_bedrock_models.txt not found") +@pytest.mark.parametrize("region", ("us-gov-east-1", "us-gov-west-1")) +@pytest.mark.parametrize("base_provider", ("bedrock_converse", "bedrock")) +def test_regional_bedrock_alias_requires_canonical_converse_metadata( + region: str, base_provider: Literal["bedrock_converse", "bedrock"] +) -> None: + base_model: Final = next( + model for model in sorted(litellm.bedrock_converse_models) if BedrockModelInfo.get_base_model(model) == model + ) + model: Final = f"bedrock/{region}/{base_model}" + model_cost: Final[Mapping[str, ModelInfoBase]] = { + model: {"litellm_provider": "bedrock", "mode": "chat"}, + base_model: {"litellm_provider": base_provider, "mode": "chat"}, + } + assert BedrockModelInfo.get_bedrock_route(model) == "converse" + if base_provider == "bedrock": + with pytest.raises(AssertionError, match=re.escape(model)): + _enforce_bedrock_converse_models(model_cost, ()) + return + _enforce_bedrock_converse_models(model_cost, ()) + + def test_get_model_info_custom_provider(): # Custom provider example copied from https://docs.litellm.ai/docs/providers/custom_llm_server: import litellm diff --git a/tests/local_testing/test_lunary.py b/tests/local_testing/test_lunary.py index a2e137ed355..f561ce00f3e 100644 --- a/tests/local_testing/test_lunary.py +++ b/tests/local_testing/test_lunary.py @@ -83,13 +83,15 @@ def test_lunary_with_tools(): ] response = litellm.completion( - model="gpt-3.5-turbo-1106", + model="gpt-6-luna", messages=messages, tools=tools, tool_choice="auto", # auto is default, but we'll be explicit ) response_message = response.choices[0].message + assert response.choices[0].message.tool_calls + assert all(call.function.name == "get_current_weather" for call in response.choices[0].message.tool_calls) print("\nLLM Response:\n", response.choices[0].message) diff --git a/tests/store_model_in_db_tests/test_mcp_servers.py b/tests/store_model_in_db_tests/test_mcp_servers.py index 0e20880ede9..94e14798c54 100644 --- a/tests/store_model_in_db_tests/test_mcp_servers.py +++ b/tests/store_model_in_db_tests/test_mcp_servers.py @@ -157,6 +157,7 @@ async def test_create_mcp_server_direct(): # Mock server manager mock_manager.add_server = mock.AsyncMock() mock_manager.reload_servers_from_database = mock.AsyncMock() + mock_manager.get_mcp_server_by_id.return_value = None # Set up test data server_id = str(uuid.uuid4()) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_managed_agent_access.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_managed_agent_access.py new file mode 100644 index 00000000000..90403e5553f --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_managed_agent_access.py @@ -0,0 +1,559 @@ +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from litellm.proxy import proxy_server +from litellm.proxy._experimental.mcp_server import mcp_server_manager +from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth import auth_checks +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import ManagedAgentContext + + +def actor(tools: tuple[str, ...] | None, *, delegated: bool = False) -> UserAPIKeyAuth: + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="agent-permissions", + mcp_servers=["slack", "linear"], + mcp_tool_permissions={"slack": list(tools)} if tools is not None else None, + ) + agent: Final = AgentResponse( + agent_id="publisher", + agent_name="Publisher", + agent_card_params={}, + object_permission=permission.model_dump(), + identity_managed=True, + ) + auth: Final = UserAPIKeyAuth(agent_id=agent.agent_id) + auth.managed_agent_policy = agent + auth.managed_agent_context = ManagedAgentContext( + agent_id=agent.agent_id, + mode="delegated" if delegated else "autonomous", + user_id="human" if delegated else None, + ) + return auth + + +@pytest.fixture(autouse=True) +def isolated_manager(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", mcp_server_manager.MCPServerManager()) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tools", (None, (), ("read",), ("read", "write"))) +async def test_autonomous_agent_uses_only_its_own_tool_grants(tools: tuple[str, ...] | None) -> None: + auth: Final = actor(tools) + assert set(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == {"slack", "linear"} + actual: Final = await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) + assert (frozenset(actual) if actual is not None else None) == (frozenset(tools) if tools is not None else None) + assert await MCPRequestHandler.get_allowed_tools_for_server("ungranted-server", auth) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "agent_tools,user_tools,expected", + ( + (None, ("read",), ("read",)), + (("read",), None, ("read",)), + (("read", "write"), ("read",), ("read",)), + (("read",), ("write",), ()), + ((), None, ()), + ), +) +async def test_delegated_server_and_tool_intersections( + monkeypatch: pytest.MonkeyPatch, + agent_tools: tuple[str, ...] | None, + user_tools: tuple[str, ...] | None, + expected: tuple[str, ...], +) -> None: + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="user-permissions", + mcp_servers=["slack", "user-only"], + mcp_tool_permissions={"slack": list(user_tools)} if user_tools is not None else None, + ) + user: Final = LiteLLM_UserTable(user_id="human", teams=[], object_permission=permission) + monkeypatch.setattr(auth_checks, "get_user_object", AsyncMock(return_value=user)) + auth: Final = actor(agent_tools, delegated=True) + assert await MCPRequestHandler.get_allowed_mcp_servers(auth) == ["slack"] + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == list(expected) + assert await MCPRequestHandler.get_allowed_tools_for_server("linear", auth) == [] + assert await MCPRequestHandler.get_allowed_tools_for_server("user-only", auth) == [] + + +@pytest.mark.asyncio +async def test_unavailable_delegated_user_never_leaves_agent_permissions_unrestricted( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(auth_checks, "get_user_object", AsyncMock(side_effect=RuntimeError("DB unavailable"))) + with pytest.raises(HTTPException) as failure: + await MCPRequestHandler.get_allowed_tools_for_server("slack", actor(None, delegated=True)) + assert failure.value.status_code == 503 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("servers,expected", (((), ()), (("slack",), ("slack",)), (("user-only",), ()))) +async def test_access_groups_cap_agent_servers_without_granting_new_ones( + monkeypatch: pytest.MonkeyPatch, + servers: tuple[str, ...], + expected: tuple[str, ...], +) -> None: + from litellm.proxy._types import LiteLLM_AccessGroupTable + + group: Final = LiteLLM_AccessGroupTable( + access_group_id="group", access_group_name="Restricted", access_mcp_server_ids=list(servers) + ) + monkeypatch.setattr(auth_checks, "get_access_object", AsyncMock(return_value=group)) + auth: Final = actor(None) + assert auth.managed_agent_policy is not None + auth.managed_agent_policy = auth.managed_agent_policy.model_copy(update={"access_group_ids": ["group"]}) + assert tuple(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == expected + if "slack" not in expected: + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["tools", "servers", "disabled", "outage"]) +async def test_delegated_mcp_revokes_warm_human_policy_before_tool_execution( + monkeypatch: pytest.MonkeyPatch, change: str +) -> None: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, object_permission_cache_key + + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="user-grant", mcp_servers=["slack"], mcp_tool_permissions={"slack": ["read", "write"]} + ) + user: Final = LiteLLM_UserTable( + user_id="human", teams=[], organization_memberships=[], object_permission_id="user-grant" + ) + cache: Final = UserApiKeyCache() + cache.set_cache("human", user) + cache.set_cache(object_permission_cache_key("user-grant"), permission) + client: Final = MagicMock() + client.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=user) + client.writer_db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=permission) + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + auth: Final = actor(("read", "write"), delegated=True) + assert set(await MCPRequestHandler.get_allowed_tools_for_server("slack", auth)) == {"read", "write"} + if change == "disabled": + client.writer_db.litellm_usertable.find_unique.return_value = user.model_copy( + update={"metadata": {"scim_active": False}} + ) + elif change == "outage": + client.writer_db.litellm_usertable.find_unique.side_effect = RuntimeError("writer unavailable") + elif change == "servers": + client.writer_db.litellm_objectpermissiontable.find_unique.return_value = permission.model_copy( + update={"mcp_servers": [], "mcp_tool_permissions": {}} + ) + else: + client.writer_db.litellm_objectpermissiontable.find_unique.return_value = permission.model_copy( + update={"mcp_tool_permissions": {"slack": ["read"]}} + ) + if change in ("disabled", "outage"): + with pytest.raises(HTTPException): + await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) + else: + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == ( + ["read"] if change == "tools" else [] + ) + client.db.litellm_usertable.find_unique.assert_not_called() + client.db.litellm_objectpermissiontable.find_unique.assert_not_called() + + +def _server_row(server_id: str, access_groups: tuple[str, ...]) -> MagicMock: + row: Final = MagicMock() + row.server_id = server_id + row.mcp_access_groups = list(access_groups) + return row + + +def _toolset_row(server_id: str, tool_name: str) -> MagicMock: + row: Final = MagicMock() + row.tools = [{"server_id": server_id, "tool_name": tool_name}] + return row + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["tool", "server", "outage"]) +async def test_autonomous_agent_toolset_and_access_group_revocations_bind_on_the_next_request( + monkeypatch: pytest.MonkeyPatch, change: str +) -> None: + """The agent's entitlements are read through the shared toolset and access-group resolvers. Once the + writer revokes a tool or drops the server from the group, the next managed request must be denied + even though the legacy cache still holds the warm grant and the replica still shows the old rows""" + from litellm.caching.caching import DualCache + from litellm.proxy._experimental.mcp_server import toolset_db + + warm_toolset: Final = _toolset_row("slack", "read") + list_toolsets: Final = AsyncMock(return_value=[warm_toolset]) + monkeypatch.setattr(toolset_db, "list_mcp_toolsets", list_toolsets) + client: Final = MagicMock() + client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[_server_row("linear", ("grp",))]) + client.writer_db.litellm_mcpservertable.find_many = AsyncMock(return_value=[_server_row("linear", ("grp",))]) + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="agent-permissions", mcp_toolsets=["ts"], mcp_access_groups=["grp"] + ) + auth: Final = actor(None) + assert auth.managed_agent_policy is not None + auth.managed_agent_policy = auth.managed_agent_policy.model_copy( + update={"object_permission": permission.model_dump()} + ) + auth.requires_fresh_policy = True + + assert set(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == {"slack", "linear"} + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == ["read"] + + if change == "tool": + list_toolsets.return_value = [_toolset_row("slack", "other")] + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == ["other"] + elif change == "server": + client.writer_db.litellm_mcpservertable.find_many.return_value = [] + assert set(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == {"slack"} + assert await MCPRequestHandler.get_allowed_tools_for_server("linear", auth) == [] + else: + list_toolsets.side_effect = RuntimeError("writer unavailable") + with pytest.raises(HTTPException) as failure: + await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) + assert failure.value.status_code == 503 + for call in list_toolsets.await_args_list: + assert call.kwargs["use_writer"] is True, "managed agent toolsets must be read from the writer" + client.db.litellm_mcpservertable.find_many.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("role", ["proxy_admin", "proxy_admin_viewer", "internal_user"]) +@pytest.mark.parametrize("open_channel", ["none", "operator", "submitted"]) +@pytest.mark.parametrize("has_grant", [True, False]) +@pytest.mark.parametrize("agent_tools", [("read", "write"), None]) +async def test_delegated_mcp_uses_explicit_team_grants_even_for_dashboard_admins( + monkeypatch: pytest.MonkeyPatch, + role: str, + open_channel: str, + has_grant: bool, + agent_tools: tuple[str, ...] | None, +) -> None: + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + manager: Final = mcp_server_manager.global_mcp_server_manager + manager.registry = { + name: MCPServer( + server_id=name, + name=name, + transport="http", + url="https://example.com/mcp", + allow_all_keys=open_channel == "operator", + ) + for name in ("slack", "linear") + } + from litellm.proxy._experimental.mcp_server import db + + monkeypatch.setattr( + db, + "get_active_submitted_mcp_server_ids_for_user", + AsyncMock(return_value=["slack", "linear"] if open_channel == "submitted" else []), + ) + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="team-grant", mcp_servers=["slack"], mcp_tool_permissions={"slack": ["read"]} + ) + user: Final = LiteLLM_UserTable( + user_id="human", user_role=role, teams=["team"] if has_grant else [], organization_memberships=[] + ) + team: Final = LiteLLM_TeamTable( + team_id="team", + models=[], + members_with_roles=[{"user_id": "human", "role": "user"}], + object_permission_id="team-grant", + ) + client: Final = MagicMock() + client.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=user) + client.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + client.writer_db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=permission) + monkeypatch.setattr(proxy_server, "prisma_client", client) + auth: Final = actor(agent_tools, delegated=True) + auth.team_id = "team" + assert await MCPRequestHandler.get_allowed_mcp_servers(auth) == (["slack"] if has_grant else []) + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == (["read"] if has_grant else []) + assert await MCPRequestHandler.get_allowed_tools_for_server("linear", auth) == [] + admitted: Final = await MCPRequestHandler.reload_admitted_user("human", requires_fresh_policy=True) + assert admitted.user_role == role + + +@pytest.mark.asyncio +async def test_explicit_grants_never_fall_back_to_open_servers_on_resolution_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy._experimental.mcp_server import db + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + manager: Final = mcp_server_manager.global_mcp_server_manager + manager.registry = {"slack": MCPServer(server_id="slack", name="slack", transport="http", allow_all_keys=True)} + monkeypatch.setattr(db, "get_active_submitted_mcp_server_ids_for_user", AsyncMock(return_value=["slack"])) + auth: Final = UserAPIKeyAuth(user_id="human") + auth.mcp_explicit_grants_only = True + with pytest.MonkeyPatch.context() as patcher: + patcher.setattr(MCPRequestHandler, "get_mcp_server_access", AsyncMock(side_effect=RuntimeError("unavailable"))) + assert await manager.get_allowed_mcp_servers(auth) == [] + auth.mcp_explicit_grants_only = False + assert await manager.get_allowed_mcp_servers(auth) == ["slack"] + + +@pytest.mark.asyncio +async def test_absent_agent_policy_and_missing_delegated_subject_grant_no_servers() -> None: + from litellm.proxy._experimental.mcp_server.auth.managed_agent_access import managed_agent_servers + + assert await managed_agent_servers(UserAPIKeyAuth()) == () + auth: Final = actor(None, delegated=True) + assert auth.managed_agent_context is not None + auth.managed_agent_context = auth.managed_agent_context.model_copy(update={"user_id": None}) + assert await MCPRequestHandler.get_allowed_mcp_servers(auth) == [] + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == [] + + +@pytest.mark.asyncio +async def test_tool_policy_outage_after_server_admission_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="human-grant", mcp_servers=["slack"]) + user: Final = LiteLLM_UserTable(user_id="human", teams=[], object_permission=permission) + monkeypatch.setattr( + auth_checks, "get_user_object", AsyncMock(side_effect=[user, RuntimeError("tool lookup unavailable")]) + ) + with pytest.raises(HTTPException) as failure: + await MCPRequestHandler.get_allowed_tools_for_server("slack", actor(None, delegated=True)) + assert failure.value.status_code == 503 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("role", (None, "proxy_admin", "internal_user")) +@pytest.mark.parametrize("scoped", (False, True)) +async def test_manager_preserves_managed_server_grants_across_open_channels( + monkeypatch: pytest.MonkeyPatch, role: str | None, scoped: bool +) -> None: + from litellm.proxy._experimental.mcp_server import db + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPServerAccess + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + manager: Final = mcp_server_manager.global_mcp_server_manager + manager.registry = { + "open": MCPServer(server_id="open", name="open", transport="http", allow_all_keys=True), + "submitted": MCPServer(server_id="submitted", name="submitted", transport="http"), + "passthrough": MCPServer( + server_id="passthrough", name="passthrough", transport="http", auth_type="true_passthrough" + ), + } + monkeypatch.setattr(db, "get_active_submitted_mcp_server_ids_for_user", AsyncMock(return_value=["submitted"])) + auth: Final = actor(None) + auth.user_role = role + assert not auth.mcp_explicit_grants_only + access: Final = MCPServerAccess(server_ids=("slack", "open")) if scoped else None + assert set(await manager.get_allowed_mcp_servers(auth, access=access)) == ( + {"slack"} if scoped else {"slack", "linear"} + ) + + +@pytest.mark.asyncio +async def test_manager_does_not_replace_managed_policy_failure_with_open_servers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + manager: Final = mcp_server_manager.global_mcp_server_manager + manager.registry = {"open": MCPServer(server_id="open", name="open", transport="http", allow_all_keys=True)} + monkeypatch.setattr(auth_checks, "get_user_object", AsyncMock(side_effect=RuntimeError("writer unavailable"))) + with pytest.raises(HTTPException) as failure: + await manager.get_allowed_mcp_servers(actor(None, delegated=True)) + assert failure.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_inline_tool_grant_admits_its_server_without_widening_tools() -> None: + auth: Final = actor(("read",)) + assert auth.managed_agent_policy is not None + auth.managed_agent_policy = auth.managed_agent_policy.model_copy( + update={"object_permission": {"object_permission_id": "tools", "mcp_tool_permissions": {"slack": ["read"]}}} + ) + assert await MCPRequestHandler.get_allowed_mcp_servers(auth) == ["slack"] + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == ["read"] + assert await MCPRequestHandler.get_allowed_tools_for_server("linear", auth) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("selected_team", (None, "selected")) +@pytest.mark.parametrize("selected_grant", (False, True)) +async def test_delegation_never_borrows_another_teams_server_or_tools( + monkeypatch: pytest.MonkeyPatch, selected_team: str | None, selected_grant: bool +) -> None: + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + user: Final = LiteLLM_UserTable(user_id="human", teams=["selected", "other"], organization_memberships=[]) + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="selected-grant", + mcp_servers=["slack"] if selected_grant else [], + mcp_tool_permissions={"slack": ["read"]} if selected_grant else {}, + ) + teams: Final = { + name: LiteLLM_TeamTable( + team_id=name, + models=[], + members_with_roles=[{"user_id": "human", "role": "user"}], + object_permission=permission if name == "selected" else LiteLLM_ObjectPermissionTable( + object_permission_id="other-grant", mcp_servers=["slack", "linear"] + ), + ) + for name in ("selected", "other") + } + + async def get_team(team_id: str, **kwargs: object) -> LiteLLM_TeamTable: + return teams[team_id] + + monkeypatch.setattr(auth_checks, "get_user_object", AsyncMock(return_value=user)) + monkeypatch.setattr(auth_checks, "get_team_object", get_team) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + auth: Final = actor(None, delegated=True) + auth.team_id = selected_team + expected: Final = ["slack"] if selected_team and selected_grant else [] + assert await MCPRequestHandler.get_allowed_mcp_servers(auth) == expected + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == (["read"] if expected else []) + assert await MCPRequestHandler.get_allowed_tools_for_server("linear", auth) == [] + ordinary: Final = await MCPRequestHandler.reload_admitted_user("human", requires_fresh_policy=True) + assert set(await MCPRequestHandler.resolve_admitted_subject_servers(ordinary)) == {"slack", "linear"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entitlement", ("group", "toolset")) +async def test_managed_mcp_rejects_unavailable_authoritative_entitlements( + monkeypatch: pytest.MonkeyPatch, entitlement: str +) -> None: + client: Final = MagicMock() + client.writer_db.litellm_mcpservertable.find_many = AsyncMock(side_effect=RuntimeError("writer unavailable")) + client.writer_db.litellm_mcptoolsettable.find_many = AsyncMock(side_effect=RuntimeError("writer unavailable")) + monkeypatch.setattr(proxy_server, "prisma_client", client) + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="entitlements", + mcp_access_groups=["group"] if entitlement == "group" else [], + mcp_toolsets=["toolset"] if entitlement == "toolset" else [], + ) + auth: Final = actor(None) + assert auth.managed_agent_policy is not None + auth.managed_agent_policy = auth.managed_agent_policy.model_copy(update={"object_permission": permission.model_dump()}) + auth.requires_fresh_policy = True + with pytest.raises(HTTPException) as failure: + await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) + assert failure.value.status_code == 503 + client.db.litellm_mcpservertable.find_many.assert_not_called() + client.db.litellm_mcptoolsettable.find_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_managed_agent_mcp_access_is_capped_at_the_invoking_callers_grants( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The managed MCP path must honour the agent_caller ceiling the same way the unmanaged path does: + the agent's own policy grants slack and linear, but the team echoed back on the request reaches + only slack, so the agent may use slack alone.""" + from litellm.proxy._types import AgentCaller + + monkeypatch.setattr( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + AsyncMock(return_value=["slack"]), + ) + monkeypatch.setattr( + MCPRequestHandler, + "_apply_user_server_ceiling", + AsyncMock(side_effect=lambda servers, _auth: (tuple(servers), False)), + ) + + monkeypatch.setattr( + MCPRequestHandler, + "_get_team_object_permission", + AsyncMock( + return_value=LiteLLM_ObjectPermissionTable( + object_permission_id="caller-team-permissions", + mcp_servers=["slack"], + mcp_tool_permissions={"slack": ["read"]}, + ) + ), + ) + monkeypatch.setattr( + MCPRequestHandler, + "_apply_user_tool_ceiling", + AsyncMock(side_effect=lambda tools, _server_id, _auth: tools), + ) + + auth: Final = actor(("read", "write")) + auth.agent_caller = AgentCaller(user_id="alice", team_id="callers") + + assert set(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == {"slack"} + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == ["read"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fresh", [False, True]) +@pytest.mark.parametrize("caller_kind", ["team", "user"]) +async def test_caller_mcp_revocation_uses_fresh_policy( + monkeypatch: pytest.MonkeyPatch, fresh: bool, caller_kind: str, +) -> None: + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, object_permission_cache_key + from litellm.types.agents import AgentCaller + + cached_permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="caller-permission", mcp_servers=["slack", "linear"], + mcp_tool_permissions={"slack": ["read", "write"]}, + ) + current_permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="caller-permission", mcp_servers=["slack"], + mcp_tool_permissions={"slack": ["read"]}, + ) + team: Final = LiteLLM_TeamTable( + team_id="caller", object_permission_id="caller-permission", object_permission=current_permission, + ) + user: Final = LiteLLM_UserTable( + user_id="caller", teams=[], object_permission_id="caller-permission", object_permission=current_permission, + ) + database: Final = MagicMock() + database.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + database.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=user) + database.writer_db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=current_permission) + cache: Final = UserApiKeyCache() + cache.set_cache("team_id:caller", team.model_copy(update={"object_permission": cached_permission})) + cache.set_cache("caller", user.model_copy(update={"object_permission": cached_permission})) + cache.set_cache(object_permission_cache_key("caller-permission"), cached_permission) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + auth: Final = actor(("read", "write")) + auth.requires_fresh_policy = fresh + auth.agent_caller = AgentCaller(team_id="caller") if caller_kind == "team" else AgentCaller(user_id="caller") + + assert set(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == ({"slack"} if fresh else {"slack", "linear"}) + assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == (["read"] if fresh else ["read", "write"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fresh", [False, True]) +async def test_caller_team_outage_cannot_remove_authoritative_server_ceiling( + monkeypatch: pytest.MonkeyPatch, fresh: bool, +) -> None: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.types.agents import AgentCaller + + database: Final = MagicMock() + database.writer_db.litellm_teamtable.find_unique = AsyncMock(side_effect=RuntimeError("writer unavailable")) + database.db.litellm_teamtable.find_unique = AsyncMock(side_effect=RuntimeError("reader unavailable")) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + auth: Final = actor(("read",)) + auth.agent_caller = AgentCaller(team_id="caller") + auth.requires_fresh_policy = fresh + + if fresh: + with pytest.raises(HTTPException) as failure: + await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert failure.value.status_code == 503 + else: + assert set(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == {"slack", "linear"} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index fc4d7b45785..0d0c65e3650 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -369,7 +369,9 @@ class TestMCPRequestHandler: result = await MCPRequestHandler._get_allowed_mcp_servers_for_key(user_api_key_auth) assert result == ["server-a"] - mock_manager.resolve_toolset_tool_permissions.assert_awaited_once_with(toolset_ids=["toolset-1"]) + mock_manager.resolve_toolset_tool_permissions.assert_awaited_once_with( + toolset_ids=["toolset-1"], requires_fresh_policy=False + ) async def test_get_allowed_mcp_servers_for_key_skips_toolset_resolution_when_none_granted(self): user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") @@ -4147,7 +4149,7 @@ async def test_get_allowed_mcp_servers_for_team_uses_helper(): "group-server2", } - mock_get_access_group_servers.assert_called_once_with(["dev-group"]) + mock_get_access_group_servers.assert_called_once_with(["dev-group"], requires_fresh_policy=False) finally: for sid in ("direct-server1", "direct-server2"): global_mcp_server_manager.registry.pop(sid, None) @@ -4316,7 +4318,7 @@ async def test_get_allowed_mcp_servers_for_key_prefers_in_memory_permission(): assert set(result) == {"direct-server", "group-server"} mock_get_perm.assert_not_called() - mock_access_groups.assert_called_once_with(["grp-alpha"]) + mock_access_groups.assert_called_once_with(["grp-alpha"], requires_fresh_policy=False) finally: global_mcp_server_manager.registry.pop("direct-server", None) @@ -4383,7 +4385,7 @@ class TestAgentMCPPermissions: self._team_servers({"callers": ["server_2", "server_3"]}), ), patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here - MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) + MCPRequestHandler, "get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) ), patch.object( # test-quality-ok: neither the agent's owner nor the caller has a personal grant MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({}) @@ -4402,7 +4404,7 @@ class TestAgentMCPPermissions: MCPRequestHandler, "_get_allowed_mcp_servers_for_team", self._team_servers({}) ), patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here - MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) + MCPRequestHandler, "get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) ), patch.object( # test-quality-ok: same seam, keyed by which user is being asked about MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({"alice": ["server_1"]}) @@ -4421,7 +4423,7 @@ class TestAgentMCPPermissions: MCPRequestHandler, "_get_allowed_mcp_servers_for_team", self._team_servers({}) ), patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here - MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) + MCPRequestHandler, "get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) ), patch.object( # test-quality-ok: None is the resolver's own "entitlement unresolvable" signal MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({"alice": None}) @@ -4538,7 +4540,7 @@ class TestAgentMCPPermissions: ) with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key") as mock_key: with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team") as mock_team: - with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_agent") as mock_agent: + with patch.object(MCPRequestHandler, "get_allowed_mcp_servers_for_agent") as mock_agent: mock_key.return_value = ["server_1", "server_2"] mock_team.return_value = [] mock_agent.return_value = ["server_1"] @@ -4555,7 +4557,7 @@ class TestAgentMCPPermissions: ) with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key") as mock_key: with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team") as mock_team: - with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_agent") as mock_agent: + with patch.object(MCPRequestHandler, "get_allowed_mcp_servers_for_agent") as mock_agent: mock_key.return_value = ["server_1", "server_2"] mock_team.return_value = [] mock_agent.return_value = [] # no agent-level restriction @@ -4611,7 +4613,7 @@ class TestAgentMCPPermissions: ) with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key") as mock_key: with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team") as mock_team: - with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_agent") as mock_agent: + with patch.object(MCPRequestHandler, "get_allowed_mcp_servers_for_agent") as mock_agent: mock_key.return_value = ["server_1", "server_2"] mock_team.return_value = [] mock_agent.return_value = ["server_2", "server_3"] @@ -4637,7 +4639,7 @@ class TestAgentMCPPermissions: ): with patch.object( MCPRequestHandler, - "_get_agent_tool_permissions_for_server", + "get_agent_tool_permissions_for_server", new_callable=AsyncMock, return_value=["tool_a"], ) as mock_agent_tools: @@ -4669,7 +4671,7 @@ class TestAgentMCPPermissions: ): with patch.object( MCPRequestHandler, - "_get_agent_tool_permissions_for_server", + "get_agent_tool_permissions_for_server", new_callable=AsyncMock, return_value=None, ): @@ -4718,10 +4720,12 @@ class TestAgentMCPPermissions: with contextlib.ExitStack() as stack: for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): stack.enter_context(patcher) - result = await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth) + result = await MCPRequestHandler.get_allowed_mcp_servers_for_agent(user_api_key_auth) assert sorted(result) == ["server-a", "server-direct"] - mock_manager.resolve_toolset_tool_permissions.assert_awaited_once_with(toolset_ids=["toolset-1"]) + mock_manager.resolve_toolset_tool_permissions.assert_awaited_once_with( + toolset_ids=["toolset-1"], requires_fresh_policy=False + ) async def test_get_allowed_mcp_servers_toolset_only_agent_caps_key_servers(self): """Regression: an agent whose only grant is a toolset used to resolve to [] and place @@ -4760,7 +4764,7 @@ class TestAgentMCPPermissions: for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): stack.enter_context(patcher) with pytest.raises(UnloadableEntitlementError): - await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth) + await MCPRequestHandler.get_allowed_mcp_servers_for_agent(user_api_key_auth) stack.enter_context( patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here MCPRequestHandler, @@ -4789,13 +4793,13 @@ class TestAgentMCPPermissions: with contextlib.ExitStack() as stack: for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): stack.enter_context(patcher) - server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + server_a_tools = await MCPRequestHandler.get_agent_tool_permissions_for_server( "server-a", user_api_key_auth ) - server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + server_b_tools = await MCPRequestHandler.get_agent_tool_permissions_for_server( "server-b", user_api_key_auth ) - server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( + server_c_tools = await MCPRequestHandler.get_agent_tool_permissions_for_server( "server-c", user_api_key_auth ) @@ -5833,7 +5837,7 @@ def test_expand_permission_list_does_not_honor_all_proxy_sentinel(): @pytest.mark.asyncio -async def test_get_allowed_mcp_servers_for_team_expands_all_proxy_sentinel_dynamically(): +async def test_get_allowed_mcp_servers_for_team_expands_all_proxy_sentinel_dynamically(monkeypatch): """The TEAM resolver expands the all-proxy sentinel to every registered server and picks up a server registered later, so a team scoped to all-proxy tracks the live registry without any change to its stored permission. Reverting the team-side @@ -5850,6 +5854,9 @@ async def test_get_allowed_mcp_servers_for_team_expands_all_proxy_sentinel_dynam from litellm.types.mcp import MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer + monkeypatch.setattr(global_mcp_server_manager, "registry", {}) + monkeypatch.setattr(global_mcp_server_manager, "config_mcp_servers", {}) + for sid in ("srv-x", "srv-y"): global_mcp_server_manager.registry[sid] = MCPServer( server_id=sid, @@ -8305,7 +8312,7 @@ class TestUserSubjectTeamUnion: ) == ["t1"] # An admitted subject never fans out HERE: it resolves one source per team first, and each of # those pins a team_id, so this helper only ever answers the single-team question. The fan-out - # itself is _admitted_subject_sources' job, asserted below. + # itself is admitted_subject_sources' job, asserted below. with self._patch(teams_by_id={}, user_teams=["t2", "t3"]): assert await MCPRequestHandler._team_ids_for_mcp_grant(_make_admitted_subject("u")) == [] # keyless, no user_id -> nothing @@ -8868,7 +8875,7 @@ class TestUserSubjectTeamUnion: teams["t-member"].organization_id = "org-a" auth = _make_admitted_subject("sso-user") with self._patch(teams_by_id=teams, user_teams=["t-member", "t-stale"]): - sources = await MCPRequestHandler._admitted_subject_sources(auth) + sources = await MCPRequestHandler.admitted_subject_sources(auth) assert [(s.team_id, s.org_id) for s in sources] == [(None, None), ("t-member", "org-a")] # The user's own source carries their grants; a team source must NOT, or the team would be @@ -9673,7 +9680,10 @@ class TestGetUserObjectPermission: def _prisma_with_user(self, user_row): prisma_client = MagicMock() - prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + from litellm.proxy._types import LiteLLM_UserTable + + row = LiteLLM_UserTable(user_id="human", object_permission_id=user_row.object_permission_id) if user_row is not None else None + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=row) return prisma_client async def test_resolves_through_the_shared_permission_cache(self): @@ -9688,7 +9698,7 @@ class TestGetUserObjectPermission: with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))), patch( "litellm.proxy.auth.auth_checks.get_object_permission", new_callable=AsyncMock, @@ -9715,7 +9725,7 @@ class TestGetUserObjectPermission: with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))), patch("litellm.proxy.auth.auth_checks.get_object_permission", new_callable=AsyncMock) as mock_get_perm, ): assert await MCPRequestHandler._get_user_object_permission(auth) is None @@ -9734,7 +9744,7 @@ class TestGetUserObjectPermission: with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))), ): assert await MCPRequestHandler._get_user_object_permission(auth) is None @@ -9748,7 +9758,7 @@ class TestGetUserObjectPermission: with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))), ): assert await MCPRequestHandler._get_user_object_permission(auth) is None @@ -9765,7 +9775,7 @@ class TestGetUserObjectPermission: with ( patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))), patch( "litellm.proxy.auth.auth_checks.get_object_permission", new_callable=AsyncMock, @@ -10085,3 +10095,47 @@ class TestScopedSessionAdmission: def test_scope_field_cannot_be_forged_through_construction(self): forged = UserAPIKeyAuth(user_id="u1", mcp_session_resource_server_id="any-server") assert forged.mcp_session_resource_server_id is None + + +@pytest.mark.asyncio +async def test_fresh_mcp_user_permission_link_ignores_cached_and_replica_grants(monkeypatch): + from litellm.caching.dual_cache import DualCache + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_UserTable + + cached = LiteLLM_UserTable(user_id="fresh-human", object_permission_id="revoked") + current = LiteLLM_UserTable(user_id="fresh-human", object_permission_id="current") + cache = DualCache() + await cache.async_set_cache(key="fresh-human", value=cached) + database = MagicMock() + database.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=current) + database.db.litellm_usertable.find_unique = AsyncMock(return_value=cached) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + assert await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) == "current" + database.db.litellm_usertable.find_unique.assert_not_awaited() + database.writer_db.litellm_usertable.find_unique.side_effect = RuntimeError("unavailable") + with pytest.raises(HTTPException) as denied: + await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) + assert denied.value.status_code == 503 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["servers", "tools"]) +async def test_managed_agent_permission_resolution_outage_is_not_an_unrestricted_grant(monkeypatch, operation): + from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.types.agents import AgentResponse + + auth = UserAPIKeyAuth(agent_id="managed") + auth.managed_agent_policy = AgentResponse(agent_id="managed", agent_name="Managed", agent_card_params={}) + permission = LiteLLM_ObjectPermissionTable(object_permission_id="policy", mcp_toolsets=["unavailable"]) + manager = MagicMock() + manager.expand_permission_list.return_value = [] + manager.resolve_toolset_tool_permissions = AsyncMock(side_effect=RuntimeError("policy unavailable")) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + resolution = ( + MCPRequestHandler.get_allowed_mcp_servers_for_agent(auth, permission) + if operation == "servers" + else MCPRequestHandler.get_agent_tool_permissions_for_server("slack", auth, permission) + ) + with pytest.raises(RuntimeError, match="policy unavailable"): + await resolution diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index b1f0b3fa67e..4e27ec134d4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -7847,7 +7847,7 @@ async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_ca key="fresh-jwt-user", value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=[]), model_type=LiteLLM_UserTable ) prisma = MagicMock() - prisma.db.litellm_usertable.find_unique = AsyncMock( + prisma.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=["team-a"]) ) proxy_globals.user_api_key_cache = cache diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 26d4263eac1..eab57b0aee4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -211,15 +211,9 @@ def _reload_mcp_manager_module(): manager_module = sys.modules["litellm.proxy._experimental.mcp_server.mcp_server_manager"] importlib.reload(utils_module) reloaded = importlib.reload(manager_module) - # After reload, server.py still holds a stale reference to the old - # global_mcp_server_manager. Update it so tests that exercise server.py - # functions (e.g. _get_tools_from_mcp_servers) use the fresh instance. - server_module = sys.modules.get("litellm.proxy._experimental.mcp_server.server") - if server_module is not None and hasattr(server_module, "global_mcp_server_manager"): - server_module.global_mcp_server_manager = reloaded.global_mcp_server_manager - operations_module = sys.modules.get("litellm.proxy._experimental.mcp_server.operations") - if operations_module is not None: - operations_module.global_mcp_server_manager = reloaded.global_mcp_server_manager + for name, module in tuple(sys.modules.items()): + if name.startswith("litellm.proxy._experimental.mcp_server.") and hasattr(module, "global_mcp_server_manager"): + module.global_mcp_server_manager = reloaded.global_mcp_server_manager return reloaded @@ -228,6 +222,20 @@ def enable_eager_mcp_oauth_discovery(monkeypatch): monkeypatch.setenv("LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP", "1") +@pytest.fixture(autouse=True) +def restore_mcp_manager_singleton(): + """``_reload_mcp_manager_module`` rebinds ``global_mcp_server_manager`` in every MCP module, so + without this the next test file inherits a manager that has none of its servers registered.""" + bound: Final = tuple( + (module, module.global_mcp_server_manager) + for name, module in tuple(sys.modules.items()) + if name.startswith("litellm.proxy._experimental.mcp_server.") and hasattr(module, "global_mcp_server_manager") + ) + yield + for module, manager in bound: + module.global_mcp_server_manager = manager + + class TestMCPServerManager: """Test MCP Server Manager stdio functionality""" @@ -5669,9 +5677,7 @@ class TestMCPServerManager: # Mock dependencies - set object_permission and object_permission_id to None # so permission checks return None (no restrictions) - user_api_key_auth = MagicMock() - user_api_key_auth.object_permission = None - user_api_key_auth.object_permission_id = None + user_api_key_auth: Final = UserAPIKeyAuth() proxy_logging_obj = MagicMock() # Mock the async methods that pre_call_tool_check calls @@ -5738,9 +5744,7 @@ class TestMCPServerManager: # Mock dependencies - set object_permission and object_permission_id to None # so permission checks return None (no restrictions) - user_api_key_auth = MagicMock() - user_api_key_auth.object_permission = None - user_api_key_auth.object_permission_id = None + user_api_key_auth: Final = UserAPIKeyAuth() proxy_logging_obj = MagicMock() # Mock the async methods that pre_call_tool_check calls @@ -5807,9 +5811,7 @@ class TestMCPServerManager: # Mock dependencies - set object_permission and object_permission_id to None # so permission checks return None (no restrictions) - user_api_key_auth = MagicMock() - user_api_key_auth.object_permission = None - user_api_key_auth.object_permission_id = None + user_api_key_auth: Final = UserAPIKeyAuth() proxy_logging_obj = MagicMock() # Mock the async methods that pre_call_tool_check calls @@ -5844,9 +5846,7 @@ class TestMCPServerManager: # Mock dependencies - set object_permission and object_permission_id to None # so permission checks return None (no restrictions) - user_api_key_auth = MagicMock() - user_api_key_auth.object_permission = None - user_api_key_auth.object_permission_id = None + user_api_key_auth: Final = UserAPIKeyAuth() proxy_logging_obj = MagicMock() # Mock the async methods that pre_call_tool_check calls @@ -6922,9 +6922,7 @@ class TestMCPServerManager: # Mock dependencies - set object_permission and object_permission_id to None # so permission checks return None (no restrictions) - user_api_key_auth = MagicMock() - user_api_key_auth.object_permission = None - user_api_key_auth.object_permission_id = None + user_api_key_auth: Final = UserAPIKeyAuth() proxy_logging_obj = MagicMock() # Mock the async methods that pre_call_tool_check calls @@ -7007,8 +7005,7 @@ class TestMCPServerManager: # Mock _create_mcp_client to return our mock client manager._create_mcp_client = AsyncMock(return_value=mock_client) - # Real auth: the listed-tool slot identity is hashed from these fields - user_api_key_auth = UserAPIKeyAuth(api_key="sk-test") + user_api_key_auth: Final = UserAPIKeyAuth(api_key="sk-test") # Mock proxy logging proxy_logging_obj = MagicMock() @@ -11825,6 +11822,72 @@ async def test_resolve_toolset_tool_permissions_single_db_fetch_across_checks(): list_toolsets_mock.assert_awaited_once() +@pytest.mark.asyncio +async def test_resolve_toolset_tool_permissions_fresh_policy_sees_writer_revocation_past_warm_cache(): + """A managed agent's tool grant revoked in the writer DB must be gone on the very next fresh + request even though the legacy cache still holds the old grant, and the fresh read must go to + the writer, not the replica""" + from litellm.caching.caching import DualCache + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + manager = MCPServerManager() + granted = MagicMock() + granted.tools = [{"server_id": "server-a", "tool_name": "echo"}] + revoked = MagicMock() + revoked.tools = [{"server_id": "server-a", "tool_name": "other"}] + list_toolsets_mock = AsyncMock(side_effect=[[granted], [revoked]]) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.toolset_db.list_mcp_toolsets", + list_toolsets_mock, + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + ): + warm = await manager.resolve_toolset_tool_permissions(toolset_ids=["ts-1"]) + legacy_after_revoke = await manager.resolve_toolset_tool_permissions(toolset_ids=["ts-1"]) + fresh_after_revoke = await manager.resolve_toolset_tool_permissions( + toolset_ids=["ts-1"], requires_fresh_policy=True + ) + + assert warm == {"server-a": ["echo"]} + assert legacy_after_revoke == warm, "legacy callers keep the cached grant by design" + assert fresh_after_revoke == {"server-a": ["other"]} + assert list_toolsets_mock.await_count == 2 + assert list_toolsets_mock.await_args_list[0].kwargs["use_writer"] is False + assert list_toolsets_mock.await_args_list[1].kwargs["use_writer"] is True + + +@pytest.mark.asyncio +async def test_resolve_toolset_tool_permissions_fresh_policy_propagates_db_fault_instead_of_no_grants(): + """A fresh read that fails must raise so the managed-agent boundary fails closed; the legacy + path keeps its swallow-to-empty behaviour""" + from litellm.caching.caching import DualCache + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + manager = MCPServerManager() + list_toolsets_mock = AsyncMock(side_effect=RuntimeError("relation does not exist")) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.toolset_db.list_mcp_toolsets", + list_toolsets_mock, + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + ): + legacy = await manager.resolve_toolset_tool_permissions(toolset_ids=["ts-1"]) + with pytest.raises(RuntimeError, match="relation does not exist"): + await manager.resolve_toolset_tool_permissions(toolset_ids=["ts-1"], requires_fresh_policy=True) + + assert legacy == {} + + class TestMaterializeAuthHeaders: """_materialize_auth_headers drives one step of a resolved httpx.Auth's own flow to turn it into a header dict for the OpenAPI egress arm, which sends plain headers and cannot carry an @@ -12146,12 +12209,8 @@ class TestDiscoveryFailureLogging: assert "unresolved" in caplog.text -def _unrestricted_auth() -> MagicMock: - """A caller with no object_permission, so only server-level checks apply.""" - user_api_key_auth = MagicMock() - user_api_key_auth.object_permission = None - user_api_key_auth.object_permission_id = None - return user_api_key_auth +def _unrestricted_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth() def _permissive_proxy_logging() -> MagicMock: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py index ed3e5f48516..8484dfdde72 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py @@ -139,7 +139,7 @@ async def test_mint_reads_the_users_teams_from_the_database_not_a_stale_cached_r key="stale-cache-user", value=_user(user_id="stale-cache-user", teams=[]), model_type=LiteLLM_UserTable ) prisma = MagicMock() - prisma.db.litellm_usertable.find_unique = AsyncMock( + prisma.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_user(user_id="stale-cache-user", teams=["team-a"]) ) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) @@ -164,7 +164,7 @@ async def test_mint_refuses_a_user_scim_deactivated_after_the_cache_last_saw_the key="deactivated-user", value=_user(user_id="deactivated-user", teams=["team-a"]), model_type=LiteLLM_UserTable ) prisma = MagicMock() - prisma.db.litellm_usertable.find_unique = AsyncMock( + prisma.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_user(user_id="deactivated-user", teams=["team-a"], metadata={"scim_active": False}) ) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index a27dd8d04bd..0b78d88f2c7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1311,7 +1311,7 @@ class TestListToolsRestAPI: session_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user") admitted_auth = UserAPIKeyAuth(user_id="grant-user", org_id="admitted-org") - async def fake_reload(user_id): + async def fake_reload(user_id, *, requires_fresh_policy=False): assert user_id == "grant-user" return admitted_auth diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py index a5f6994b1a7..816ccc5e7e6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_ui_session_utils.py @@ -145,7 +145,7 @@ async def test_build_effective_auth_contexts_appends_admitted_user_context(monke assert contexts[-1].user_id == "user-42" and contexts[-1].team_id is None assert [ctx.team_id for ctx in contexts[:-1]] == ["team-one"] - reload_mock.assert_awaited_once_with("user-42") + reload_mock.assert_awaited_once_with("user-42", requires_fresh_policy=False) @pytest.mark.asyncio @@ -198,7 +198,7 @@ async def test_acting_user_auth_returns_admitted_subject_for_non_admin_sessions( result = await acting_user_auth(user_auth) assert result.user_id == "user-42" and result.team_id is None - reload_mock.assert_awaited_once_with("user-42") + reload_mock.assert_awaited_once_with("user-42", requires_fresh_policy=False) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py index e744e84d671..08cceb0d967 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py @@ -144,3 +144,27 @@ async def test_default_loader_returns_nothing_without_a_db(monkeypatch: pytest.M monkeypatch.setattr(proxy_server, "prisma_client", None) assert await _load_access_group("ag-1") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strict", [False, True]) +async def test_authoritative_group_ceiling_propagates_policy_outages( + monkeypatch: pytest.MonkeyPatch, strict: bool +) -> None: + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints.auth.agent_access_groups import _load_access_group + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + database: Final = MagicMock() + database.db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + database.writer_db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + if strict: + with pytest.raises(HTTPException) as failure: + await _load_access_group("group", check_db_only=True) + assert failure.value.status_code == 503 + else: + assert await _load_access_group("group") is None diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index a87716375e8..4b8d28e2406 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -67,7 +67,7 @@ class TestAgentRequestHandler: # Case 1: Both key and team have agents - intersection with patch.object( - AgentRequestHandler, "_get_allowed_agents_for_key" + AgentRequestHandler, "get_allowed_agents_for_key" ) as mock_key: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" @@ -86,7 +86,7 @@ class TestAgentRequestHandler: # Case 2: Team has agents, key has none - inherit from team with patch.object( - AgentRequestHandler, "_get_allowed_agents_for_key" + AgentRequestHandler, "get_allowed_agents_for_key" ) as mock_key: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" @@ -105,7 +105,7 @@ class TestAgentRequestHandler: # Case 3: Key has agents, team has none - key restrictions stand with patch.object( - AgentRequestHandler, "_get_allowed_agents_for_key" + AgentRequestHandler, "get_allowed_agents_for_key" ) as mock_key: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" @@ -120,7 +120,7 @@ class TestAgentRequestHandler: # Case 4: No grant anywhere - unrestricted (documented open-by-default) with patch.object( - AgentRequestHandler, "_get_allowed_agents_for_key" + AgentRequestHandler, "get_allowed_agents_for_key" ) as mock_key: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" @@ -141,7 +141,7 @@ class TestAgentRequestHandler: api_key="test-key", user_id="test-user", team_id="test-team" ) - with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "get_allowed_agents_for_key") as mock_key: with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: mock_key.return_value = RestrictedAgentAccess(frozenset({"agent-alpha"})) mock_team.return_value = RestrictedAgentAccess(frozenset({"agent-beta"})) @@ -198,7 +198,7 @@ class TestAgentRequestHandler: @staticmethod def _team_grants(grants: dict[str, AgentAccess]) -> AsyncMock: - async def by_team(user_api_key_auth: UserAPIKeyAuth | None = None) -> AgentAccess: + async def by_team(user_api_key_auth: UserAPIKeyAuth | None = None, *, strict: bool = False) -> AgentAccess: assert user_api_key_auth is not None return grants.get(user_api_key_auth.team_id or "", UnrestrictedAgentAccess()) @@ -237,6 +237,29 @@ class TestAgentRequestHandler: frozenset() ) + async def test_managed_agent_acting_for_a_user_is_capped_at_the_invoking_teams_agents(self): + """The managed path must honour the invoking team's ceiling the same way the unmanaged path does: + the agent's own policy grants alpha and beta, but the human who invoked it reaches only beta.""" + from litellm.types.agents import AgentResponse + + managed: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="actor") + managed.managed_agent_policy = AgentResponse( + agent_id="actor", + agent_name="Actor", + agent_card_params={}, + object_permission={"object_permission_id": "own", "agents": ["agent-alpha", "agent-beta"]}, + ) + managed.agent_caller = AgentCaller(user_id="alice", team_id="callers") + + with patch.object( # test-quality-ok: the team resolver reads proxy_server globals with no injection seam + AgentRequestHandler, + "_get_allowed_agents_for_team", + self._team_grants({"callers": RestrictedAgentAccess(frozenset({"agent-beta"}))}), + ): + assert await AgentRequestHandler.resolve_agent_access(managed) == RestrictedAgentAccess( + frozenset({"agent-beta"}) + ) + async def test_agent_key_acting_for_an_ungranted_caller_keeps_its_own_agents(self): agent_key: Final = self._key_granting(["agent-alpha"], agent_id="caller-agent") agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers") @@ -249,7 +272,6 @@ class TestAgentRequestHandler: frozenset({"agent-alpha"}) ) - async def test_agent_access_groups_intersect_with_key_grants(self): agent_key: Final = self._key_granting(["agent-alpha", "agent-beta"], agent_id="caller-agent") resolve, _ = self._ceiling_resolver(frozenset({"agent-beta", "agent-gamma"})) @@ -299,7 +321,7 @@ class TestAgentRequestHandler: ) as mock_groups: mock_groups.return_value = [] - assert await AgentRequestHandler._get_allowed_agents_for_key( + assert await AgentRequestHandler.get_allowed_agents_for_key( user_api_key_auth=mock_user_auth ) == RestrictedAgentAccess(frozenset()) @@ -315,7 +337,7 @@ class TestAgentRequestHandler: ) as mock_groups: mock_groups.side_effect = Exception("DB Error") - assert await AgentRequestHandler._get_allowed_agents_for_key( + assert await AgentRequestHandler.get_allowed_agents_for_key( user_api_key_auth=mock_user_auth ) == UnrestrictedAgentAccess() @@ -404,7 +426,7 @@ class TestAgentRequestHandler: ) with patch.object( - AgentRequestHandler, "_get_allowed_agents_for_key" + AgentRequestHandler, "get_allowed_agents_for_key" ) as mock_key: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" @@ -489,9 +511,9 @@ class TestAgentRequestHandler: listed: Final = await accessible_agents(session, registry.get_agent_list(), resolve_access, effective_contexts) assert {agent.agent_name for agent in listed} == {"alpha", "beta"} - async def test_get_allowed_agents_for_key_via_access_group_ids(self): + async def testget_allowed_agents_for_key_via_access_group_ids(self): """ - Test that _get_allowed_agents_for_key includes agents from key's access_group_ids + Test that get_allowed_agents_for_key includes agents from key's access_group_ids (unified access groups) when key has no native object_permission. """ mock_user_auth = UserAPIKeyAuth( @@ -508,16 +530,16 @@ class TestAgentRequestHandler: new_callable=AsyncMock, return_value=["agent-from-ag-1", "agent-from-ag-2"], ): - result = await AgentRequestHandler._get_allowed_agents_for_key( + result = await AgentRequestHandler.get_allowed_agents_for_key( user_api_key_auth=mock_user_auth ) assert result == RestrictedAgentAccess( frozenset({"agent-from-ag-1", "agent-from-ag-2"}) ) - async def test_get_allowed_agents_for_key_combines_native_and_access_groups(self): + async def testget_allowed_agents_for_key_combines_native_and_access_groups(self): """ - Test that _get_allowed_agents_for_key combines agents from native object_permission + Test that get_allowed_agents_for_key combines agents from native object_permission and key's access_group_ids (unified access groups). """ from litellm.proxy._types import LiteLLM_ObjectPermissionTable @@ -540,7 +562,7 @@ class TestAgentRequestHandler: new_callable=AsyncMock, return_value=["agent-from-ag"], ): - result = await AgentRequestHandler._get_allowed_agents_for_key( + result = await AgentRequestHandler.get_allowed_agents_for_key( user_api_key_auth=mock_user_auth ) assert result == RestrictedAgentAccess( @@ -611,7 +633,7 @@ class TestAgentRequestHandler: "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry, ): - with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "get_allowed_agents_for_key") as mock_key: with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: for key_grant, team_grant in ( ( @@ -632,3 +654,392 @@ class TestAgentRequestHandler: assert await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) == RestrictedAgentAccess(frozenset({agent.agent_id})), (key_grant, team_grant) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "state,allowed", + [ + ({}, True), + ({"enabled": False}, False), + ], +) +async def test_managed_invocation_requires_local_and_directory_admission( + monkeypatch: pytest.MonkeyPatch, state: dict[str, object], allowed: bool +) -> None: + from unittest.mock import MagicMock + + from litellm.proxy import proxy_server + from litellm.types.agents import AgentResponse + from litellm.types.proxy.agent_identity import AgentIdentityBinding + + binding: Final = AgentIdentityBinding( + agent_id="target", + provider="microsoft_entra", + tenant_id="tenant", + client_id="client", + issuer="issuer", + revision="revision", + ) + target: Final = AgentResponse( + agent_id="target", agent_name="Target", agent_card_params={}, identity=binding, identity_managed=True + ).model_copy(update=state) + client: Final = MagicMock() + client.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=target) + monkeypatch.setattr(proxy_server, "prisma_client", client) + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="human-grant", agents=["target"]) + auth: Final = UserAPIKeyAuth(user_id="human", object_permission=permission) + assert await AgentRequestHandler.is_agent_allowed("target", auth) is allowed + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delegated", [True, False]) +async def test_managed_agent_invocation_grants_intersect_verified_user_grants( + monkeypatch: pytest.MonkeyPatch, delegated: bool +) -> None: + from unittest.mock import MagicMock + + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.auth import auth_checks + from litellm.types.agents import AgentResponse + from litellm.types.proxy.agent_identity import AgentIdentityBinding, ManagedAgentContext + + database: Final = MagicMock() + monkeypatch.setattr(proxy_server, "prisma_client", database) + own: Final = LiteLLM_ObjectPermissionTable(object_permission_id="own", agents=["shared", "agent-only"]) + human_grants: Final = LiteLLM_ObjectPermissionTable(object_permission_id="human", agents=["shared", "human-only"]) + human: Final = LiteLLM_UserTable(user_id="human", teams=[], object_permission=human_grants) + monkeypatch.setattr(auth_checks, "get_user_object", AsyncMock(return_value=human)) + auth: Final = UserAPIKeyAuth(agent_id="actor", api_key="verified-jwt") + auth.managed_agent_policy = AgentResponse( + agent_id="actor", agent_name="Actor", agent_card_params={}, object_permission=own.model_dump() + ) + auth.managed_agent_context = ManagedAgentContext( + agent_id="actor", mode="delegated" if delegated else "autonomous", user_id="human" if delegated else None + ) + access: Final = await AgentRequestHandler.resolve_agent_access(auth) + assert access == RestrictedAgentAccess(frozenset({"shared"} if delegated else {"shared", "agent-only"})) + + target: Final = AgentResponse( + agent_id="shared", agent_name="Shared", agent_card_params={}, identity_managed=True, + identity=AgentIdentityBinding( + agent_id="shared", provider="microsoft_entra", tenant_id="tenant", client_id="client", issuer="issuer", revision="current" + ), + ) + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=target) + assert await AgentRequestHandler.is_agent_allowed("shared", auth) is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("revoked", ["user", "team-member", "team-grant", "team-permission", "direct-grant", "access-group"]) +async def test_delegated_grants_revoke_with_warm_user_team_and_permission_caches( + monkeypatch: pytest.MonkeyPatch, revoked: str +) -> None: + from unittest.mock import MagicMock + + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_AccessGroupTable, LiteLLM_TeamTable, LiteLLM_UserTable + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import verified_human_agent_grants + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, object_permission_cache_key + + direct: Final = revoked == "direct-grant" + grouped: Final = revoked == "access-group" + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="grant", agents=["target"]) + human: Final = LiteLLM_UserTable( + user_id="human", + teams=[] if direct else ["team"], + organization_memberships=[], + object_permission_id="grant" if direct else None, + ) + team: Final = LiteLLM_TeamTable( + team_id="team", + models=[], + members_with_roles=[{"user_id": "human", "role": "user"}], + object_permission_id=None if grouped else "grant", + access_group_ids=["group"] if grouped else [], + ) + group: Final = LiteLLM_AccessGroupTable( + access_group_id="group", access_group_name="Group", access_agent_ids=["target"] + ) + cache: Final = UserApiKeyCache() + cache.set_cache("human", human) + cache.set_cache("team_id:team", team) + cache.set_cache(object_permission_cache_key("grant"), permission) + cache.set_cache("access_group_id:group", group) + client: Final = MagicMock() + client.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=human) + client.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + client.writer_db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=permission) + client.writer_db.litellm_accessgrouptable.find_unique = AsyncMock(return_value=group) + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + assert await verified_human_agent_grants("human", "team") == frozenset({"target"}) + client.writer_db.litellm_usertable.find_unique.return_value = ( + human.model_copy(update={"teams": []}) if revoked == "user" else human + ) + client.writer_db.litellm_teamtable.find_unique.return_value = ( + team.model_copy(update={"members_with_roles": []}) + if revoked == "team-member" + else team.model_copy(update={"object_permission_id": None}) + if revoked == "team-grant" + else team + ) + client.writer_db.litellm_objectpermissiontable.find_unique.return_value = ( + permission.model_copy(update={"agents": []}) if direct or revoked == "team-permission" else permission + ) + client.writer_db.litellm_accessgrouptable.find_unique.return_value = ( + group.model_copy(update={"access_agent_ids": []}) if grouped else group + ) + assert await verified_human_agent_grants("human", "team") == frozenset() + client.db.litellm_usertable.find_unique.assert_not_called() + client.db.litellm_teamtable.find_unique.assert_not_called() + client.db.litellm_objectpermissiontable.find_unique.assert_not_called() + client.db.litellm_accessgrouptable.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_strict_legacy_group_grants_ignore_stale_replica(monkeypatch: pytest.MonkeyPatch) -> None: + from unittest.mock import MagicMock + + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints import agent_registry + from litellm.types.agents import AgentResponse + + stale: Final = AgentResponse(agent_id="revoked", agent_name="Revoked", agent_card_params={}) + registry: Final = AgentRegistry() + registry.register_agent(stale) + database: Final = MagicMock() + database.db.litellm_agentstable.find_many = AsyncMock(return_value=[stale]) + database.writer_db.litellm_agentstable.find_many = AsyncMock(return_value=[stale]) + monkeypatch.setattr(agent_registry, "global_agent_registry", registry) + monkeypatch.setattr(proxy_server, "prisma_client", database) + auth: Final = UserAPIKeyAuth( + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="permission", agent_access_groups=["group"] + ) + ) + assert await AgentRequestHandler.get_allowed_agents_for_key(auth, strict=True) == RestrictedAgentAccess( + frozenset({"revoked"}) + ) + database.writer_db.litellm_agentstable.find_many.return_value = [] + assert await AgentRequestHandler.get_allowed_agents_for_key(auth, strict=True) == RestrictedAgentAccess( + frozenset() + ) + database.db.litellm_agentstable.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("groups", [[], ["group"]]) +async def test_legacy_groups_without_database_grant_no_agents(groups: list[str]) -> None: + assert await AgentRequestHandler._get_db_agent_ids_for_access_groups(None, groups, check_db_only=True) == set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("team", [False, True]) +async def test_strict_invocation_policy_outage_denies_instead_of_allowing_all( + monkeypatch: pytest.MonkeyPatch, team: bool +) -> None: + from fastapi import HTTPException + from unittest.mock import MagicMock + from litellm.proxy import proxy_server + + database: Final = MagicMock() + database.writer_db.litellm_teamtable.find_unique = AsyncMock(side_effect=ConnectionError("writer unavailable")) + database.writer_db.litellm_agentstable.find_many = AsyncMock(side_effect=ConnectionError("writer unavailable")) + monkeypatch.setattr(proxy_server, "prisma_client", database) + auth: Final = UserAPIKeyAuth( + team_id="team" if team else None, + object_permission=None if team else LiteLLM_ObjectPermissionTable( + object_permission_id="grant", agent_access_groups=["group"] + ), + ) + with pytest.raises(HTTPException, match="policy is unavailable") as denied: + await AgentRequestHandler.resolve_key_team_agent_access(auth, strict=True) + assert denied.value.status_code == 503 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("available", [False, True]) +async def test_missing_team_cannot_grant_strict_agent_access(monkeypatch: pytest.MonkeyPatch, available: bool) -> None: + from unittest.mock import MagicMock + from litellm.proxy import proxy_server + from litellm.proxy.auth import auth_checks + + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock() if available else None) + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=None)) + assert await AgentRequestHandler._get_allowed_agents_for_team( + UserAPIKeyAuth(team_id="missing"), strict=True + ) == RestrictedAgentAccess(frozenset()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outage", [False, True]) +async def test_registered_managed_target_cannot_bypass_missing_or_unavailable_policy( + monkeypatch: pytest.MonkeyPatch, outage: bool +) -> None: + from fastapi import HTTPException + from unittest.mock import MagicMock + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints import agent_registry + from litellm.types.agents import AgentResponse + + registry: Final = AgentRegistry() + registry.register_agent(AgentResponse( + agent_id="target", agent_name="Target", agent_card_params={}, identity_managed=True + )) + monkeypatch.setattr(agent_registry, "global_agent_registry", registry) + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock( + return_value=None, side_effect=ConnectionError("unavailable") if outage else None + ) + monkeypatch.setattr(proxy_server, "prisma_client", database) + if outage: + with pytest.raises(HTTPException, match="could not be loaded") as denied: + await AgentRequestHandler.is_agent_allowed("target", UserAPIKeyAuth()) + assert denied.value.status_code == 503 + else: + assert await AgentRequestHandler.is_agent_allowed("target", UserAPIKeyAuth()) is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("grant", [False, True]) +async def test_delegation_without_a_verified_human_never_grants_agents(grant: bool) -> None: + from litellm.types.agents import AgentResponse + from litellm.types.proxy.agent_identity import ManagedAgentContext + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import verified_human_agent_grants + + auth: Final = UserAPIKeyAuth(agent_id="actor") + auth.managed_agent_policy = AgentResponse( + agent_id="actor", agent_name="Actor", agent_card_params={}, + object_permission={"object_permission_id": "own", "agents": ["target"]} if grant else None, + ) + auth.managed_agent_context = ManagedAgentContext(agent_id="actor", mode="delegated") + assert await AgentRequestHandler.resolve_agent_access(auth) == RestrictedAgentAccess(frozenset()) + assert await verified_human_agent_grants(None) == frozenset() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ("grant", "permission_reference", "groups", "team", "blocked", "expired", "deleted", "outage")) +async def test_managed_target_rechecks_authoritative_key_after_peer_revocation( + monkeypatch: pytest.MonkeyPatch, change: str +) -> None: + from unittest.mock import MagicMock + from fastapi import HTTPException + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.types.agents import AgentResponse + from litellm.types.proxy.agent_identity import AgentIdentityBinding + + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="grant", agents=["target"]) + warm: Final = UserAPIKeyAuth(api_key="a" * 64, token="a" * 64, object_permission_id="grant", object_permission=permission) + target: Final = AgentResponse( + agent_id="target", agent_name="Target", agent_card_params={}, identity_managed=True, + identity=AgentIdentityBinding(agent_id="target", provider="microsoft_entra", tenant_id="tenant", client_id="client", issuer="issuer", revision="current"), + ) + client: Final = MagicMock() + client.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=target) + client.get_data = AsyncMock(return_value=warm) + client.writer_db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=permission) + cache: Final = UserApiKeyCache() + cache.set_cache("a" * 64, warm) + monkeypatch.setattr(proxy_server, "prisma_client", client) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + assert await AgentRequestHandler.is_agent_allowed("target", warm) is True + client.get_data.return_value = warm.model_copy(update={ + "object_permission": None, + "object_permission_id": "replacement" if change == "permission_reference" else "grant", + "access_group_ids": [], + "team_id": "new-team" if change == "team" else None, + "blocked": change == "blocked", + "expires": "2000-01-01T00:00:00+00:00" if change == "expired" else None, + }) + client.writer_db.litellm_objectpermissiontable.find_unique.return_value = permission.model_copy(update={"agents": []}) + if change == "team": + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks + client.writer_db.litellm_objectpermissiontable.find_unique.return_value = permission + monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=LiteLLM_TeamTable( + team_id="new-team", object_permission=permission.model_copy(update={"agents": ["other"]}) + ))) + if change == "groups": + warm.object_permission = None + warm.access_group_ids = ["old-group"] + from litellm.proxy.auth import auth_checks + monkeypatch.setattr(auth_checks, "_get_agent_ids_from_access_groups", AsyncMock(return_value=["target"])) + if change == "deleted": + client.get_data.return_value = None + if change == "outage": + client.get_data.side_effect = RuntimeError("writer unavailable") + if change in ("blocked", "expired", "deleted", "outage"): + with pytest.raises((HTTPException, RuntimeError)): + await AgentRequestHandler.is_agent_allowed("target", warm) + else: + assert await AgentRequestHandler.is_agent_allowed("target", warm) is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ceiling", ["agent-group", "caller-team", "group-without-grant"]) +@pytest.mark.parametrize("permitted", [False, True]) +async def test_managed_target_preserves_ordinary_actor_ceilings_after_key_reload( + monkeypatch: pytest.MonkeyPatch, ceiling: str, permitted: bool +) -> None: + from unittest.mock import MagicMock + + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_AccessGroupTable, LiteLLM_TeamTable + from litellm.proxy.agent_endpoints import agent_registry + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.types.agents import AgentResponse + from litellm.types.proxy.agent_identity import AgentIdentityBinding + + target: Final = AgentResponse( + agent_id="target", agent_name="Target", agent_card_params={}, identity_managed=True, + identity=AgentIdentityBinding( + agent_id="target", provider="microsoft_entra", tenant_id="tenant", client_id="client", + issuer="issuer", revision="current", + ), + ) + actor: Final = AgentResponse( + agent_id="ordinary", agent_name="Ordinary", agent_card_params={}, + access_group_ids=["actor-group"] if ceiling != "caller-team" else [], + ) + registry: Final = AgentRegistry() + registry.register_agent(actor) + registry.register_agent(target) + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="key-grant", agents=[] if ceiling == "group-without-grant" else ["target"] + ) + persisted: Final = UserAPIKeyAuth( + api_key="a" * 64, agent_id="ordinary", object_permission_id="key-grant", object_permission=permission, + ) + auth: Final = persisted.model_copy() + auth.agent_caller = AgentCaller(team_id="caller-team") if ceiling == "caller-team" else None + group: Final = LiteLLM_AccessGroupTable( + access_group_id="actor-group", access_group_name="Actor group", + access_agent_ids=["target"] if permitted else ["other"], + ) + team: Final = LiteLLM_TeamTable( + team_id="caller-team", object_permission_id="caller-grant", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="caller-grant", agents=["target"] if permitted else ["other"], + ), + ) + database: Final = MagicMock() + database.get_data = AsyncMock(return_value=persisted) + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=target) + database.writer_db.litellm_objectpermissiontable.find_unique = AsyncMock( + side_effect=lambda where: permission if where["object_permission_id"] == "key-grant" else team.object_permission + ) + database.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + database.writer_db.litellm_accessgrouptable.find_unique = AsyncMock(return_value=group) + cache: Final = UserApiKeyCache() + cache.set_cache("access_group_id:actor-group", group.model_copy(update={"access_agent_ids": ["target"]})) + cache.set_cache("team_id:caller-team", team.model_copy(update={"object_permission": permission})) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(agent_registry, "global_agent_registry", registry) + + assert await AgentRequestHandler.is_agent_allowed("target", auth) is (permitted and ceiling != "group-without-grant") + database.get_data.assert_awaited_once() + assert auth.agent_caller == (AgentCaller(team_id="caller-team") if ceiling == "caller-team" else None) diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py new file mode 100644 index 00000000000..7747e5eff71 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py @@ -0,0 +1,285 @@ +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.managed_authorization import ( + actor_admission_failure, + admit_managed_actor, +) +from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import AgentIdentityBinding, AgentIdentityFailure, ManagedAgentContext + +BINDING: Final = AgentIdentityBinding( + agent_id="agent", + provider="microsoft_entra", + tenant_id="tenant", + client_id="client", + service_principal_id="principal", + issuer="issuer", + revision="current", +) + + +def agent(**overrides: object) -> AgentResponse: + return AgentResponse.model_validate( + { + "agent_id": "agent", + "agent_name": "Agent", + "agent_card_params": {}, + "identity": BINDING, + "identity_managed": True, + "execution_mode": "both", + **overrides, + } + ) + + +@pytest.mark.parametrize( + "state", + [ + {"enabled": False}, + {"identity": None}, + {"identity": BINDING.model_copy(update={"active": False})}, + {"execution_mode": "delegated"}, + ], +) +def test_keys_cannot_bypass_lifecycle_or_delegated_only_mode(state: dict[str, object]) -> None: + assert isinstance(actor_admission_failure(agent(**state), None), AgentIdentityFailure) + + +@pytest.mark.parametrize("mode", ["autonomous", "both", "delegated"]) +def test_keys_cannot_impersonate_an_entra_bound_agent(mode: str) -> None: + assert isinstance(actor_admission_failure(agent(execution_mode=mode), None), AgentIdentityFailure) + + +@pytest.mark.parametrize( + "context", + [ + ManagedAgentContext(agent_id="agent", binding_revision="previous", mode="autonomous"), + ManagedAgentContext(agent_id="another", binding_revision="current", mode="autonomous"), + ManagedAgentContext(agent_id="agent", binding_revision="current", mode="delegated"), + ], +) +def test_stale_binding_and_unverified_delegation_cannot_pass_admission(context: ManagedAgentContext) -> None: + assert isinstance(actor_admission_failure(agent(), context), AgentIdentityFailure) + + +@pytest.mark.asyncio +async def test_deleted_agent_key_cannot_fall_back_to_unmanaged_authentication() -> None: + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=None) + database.writer_db.litellm_retiredagent.find_unique = AsyncMock(return_value={"original_agent_id": "deleted"}) + with pytest.raises(HTTPException, match="Agent no longer exists"): + await admit_managed_actor(UserAPIKeyAuth(agent_id="deleted"), AgentIdentityStore.from_client(database)) + database.writer_db.litellm_retiredagent.find_unique.return_value = None + auth: Final = UserAPIKeyAuth(agent_id="legacy-attribution-label") + await admit_managed_actor(auth, AgentIdentityStore.from_client(database)) + assert auth.managed_agent_policy is None + database.db.litellm_agentstable.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_agent_history_outage_does_not_permit_legacy_fallback() -> None: + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=None) + database.writer_db.litellm_retiredagent.find_unique = AsyncMock(side_effect=RuntimeError("unavailable")) + with pytest.raises(HTTPException) as failure: + await admit_managed_actor(UserAPIKeyAuth(agent_id="deleted"), AgentIdentityStore.from_client(database)) + assert failure.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_agent_admission_database_outage_fails_closed() -> None: + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock(side_effect=RuntimeError("DB unavailable")) + with pytest.raises(HTTPException) as failure: + await admit_managed_actor(UserAPIKeyAuth(agent_id="agent"), AgentIdentityStore.from_client(database)) + assert failure.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_human_authentication_does_not_load_an_agent() -> None: + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock() + await admit_managed_actor(UserAPIKeyAuth(user_id="human"), AgentIdentityStore.from_client(database)) + database.writer_db.litellm_agentstable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_disabled_agent_key_is_rejected_at_admission() -> None: + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=agent(enabled=False)) + with pytest.raises(HTTPException) as failure: + await admit_managed_actor(UserAPIKeyAuth(agent_id="agent"), AgentIdentityStore.from_client(database)) + assert failure.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("permitted", [True, False]) +async def test_verified_human_still_needs_an_explicit_agent_invocation_grant( + monkeypatch: pytest.MonkeyPatch, + permitted: bool, +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_UserTable + from litellm.proxy.auth import auth_checks + + policy: Final = agent() + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=policy) + monkeypatch.setattr(proxy_server, "prisma_client", database) + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="human-grants", + agents=["agent"] if permitted else [], + ) + human: Final = LiteLLM_UserTable(user_id="human", teams=[], object_permission=permission) + monkeypatch.setattr(auth_checks, "get_user_object", AsyncMock(return_value=human)) + auth: Final = UserAPIKeyAuth(agent_id="agent") + auth.managed_agent_context = ManagedAgentContext( + agent_id="agent", + binding_revision="current", + mode="delegated", + user_id="human", + ) + if permitted: + await admit_managed_actor(auth, AgentIdentityStore.from_client(database)) + assert auth.managed_agent_policy == policy + assert auth.billing_agent_policy == policy + else: + with pytest.raises(HTTPException) as failure: + await admit_managed_actor(auth, AgentIdentityStore.from_client(database)) + assert failure.value.status_code == 403 + + +def test_execution_mode_must_match_verified_token_mode() -> None: + context: Final = ManagedAgentContext(agent_id="agent", binding_revision="current", mode="autonomous") + failure: Final = actor_admission_failure(agent(execution_mode="delegated"), context) + assert isinstance(failure, AgentIdentityFailure) + assert "execution mode" in failure.message + + +@pytest.mark.asyncio +async def test_legacy_jwt_cannot_adopt_an_agent_bound_on_another_worker() -> None: + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=agent(execution_mode="autonomous")) + auth: Final = UserAPIKeyAuth(agent_id="agent", jwt_claims={"agent": "agent", "sub": "unrelated-subject"}) + with pytest.raises(HTTPException) as denied: + await admit_managed_actor(auth, AgentIdentityStore.from_client(database)) + assert denied.value.status_code == 403 + assert auth.managed_agent_policy is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bound", [False, True]) +async def test_managed_context_or_binding_requires_database(monkeypatch: pytest.MonkeyPatch, bound: bool) -> None: + from litellm.proxy.agent_endpoints import agent_registry + from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry + + registry: Final = AgentRegistry() + registry.register_agent(agent(identity_managed=bound, identity=BINDING if bound else None)) + monkeypatch.setattr(agent_registry, "global_agent_registry", registry) + auth: Final = UserAPIKeyAuth(agent_id="agent") + if not bound: + auth.managed_agent_context = ManagedAgentContext( + agent_id="agent", binding_revision="current", mode="autonomous" + ) + with pytest.raises(HTTPException) as denied: + await admit_managed_actor(auth, None) + assert denied.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_autonomous_app_rejects_persisted_virtual_key_impersonation() -> None: + policy: Final = agent(execution_mode="autonomous") + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=policy) + auth: Final = UserAPIKeyAuth(agent_id="agent", api_key="persisted-key") + with pytest.raises(HTTPException, match="bound identity provider token") as denied: + await admit_managed_actor(auth, AgentIdentityStore.from_client(database)) + assert denied.value.status_code == 403 + assert auth.managed_agent_policy is None + assert auth.billing_agent_policy is None + + +@pytest.mark.parametrize("mode,user", [("autonomous", None), ("delegated", "verified-human")]) +def test_matching_identity_revision_and_execution_mode_pass_admission(mode: str, user: str | None) -> None: + context: Final = ManagedAgentContext.model_validate( + {"agent_id": "agent", "binding_revision": "current", "mode": mode, "user_id": user} + ) + assert actor_admission_failure(agent(), context) is None + + +@pytest.mark.asyncio +async def test_bound_autonomous_actor_is_admitted_without_a_human() -> None: + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=agent()) + auth: Final = UserAPIKeyAuth(agent_id="agent") + auth.managed_agent_context = ManagedAgentContext(agent_id="agent", binding_revision="current", mode="autonomous") + await admit_managed_actor(auth, AgentIdentityStore.from_client(database)) + assert auth.managed_agent_policy == agent() + assert auth.billing_agent_policy == agent() + assert auth.user_id is None + + +@pytest.mark.asyncio +async def test_admitted_managed_actor_requires_fresh_policy_so_revocations_bind_next_request() -> None: + """Managed MCP grants (toolsets, access groups) are read through the shared resolvers, which only + bypass the warm cache and the replica when the subject carries requires_fresh_policy""" + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=agent()) + auth: Final = UserAPIKeyAuth(agent_id="agent") + auth.managed_agent_context = ManagedAgentContext(agent_id="agent", binding_revision="current", mode="autonomous") + assert auth.requires_fresh_policy is False + await admit_managed_actor(auth, AgentIdentityStore.from_client(database)) + assert auth.requires_fresh_policy is True + + +async def test_jwt_delegation_verification_is_consumed_once_and_cannot_be_supplied_by_a_caller( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.agent_endpoints.auth import agent_permission_handler + + policy: Final = agent() + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=policy) + store: Final = AgentIdentityStore.from_client(database) + grants: Final = AsyncMock(return_value=frozenset()) + monkeypatch.setattr(agent_permission_handler, "verified_human_agent_grants", grants) + auth: Final = UserAPIKeyAuth.model_validate({"agent_id": "agent", "_managed_delegation_verified": True}) + assert auth._managed_delegation_verified is False + auth.managed_agent_context = ManagedAgentContext( + agent_id="agent", binding_revision="current", mode="delegated", user_id="human" + ) + auth._managed_delegation_verified = True + assert "_managed_delegation_verified" not in auth.model_dump() + await admit_managed_actor(auth, store) + grants.assert_not_awaited() + assert auth._managed_delegation_verified is False + with pytest.raises(HTTPException) as failure: + await admit_managed_actor(auth, store) + assert failure.value.status_code == 403 + grants.assert_awaited_once_with("human", None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("database_available", (False, True)) +async def test_ordinary_agent_admission_preserves_legacy_authentication( + monkeypatch: pytest.MonkeyPatch, database_available: bool +) -> None: + from litellm.proxy.agent_endpoints import agent_registry + + registry: Final = agent_registry.AgentRegistry() + ordinary: Final = agent(identity_managed=False, identity=None) + registry.register_agent(ordinary) + monkeypatch.setattr(agent_registry, "global_agent_registry", registry) + database: Final = MagicMock() + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=ordinary) + auth: Final = UserAPIKeyAuth(agent_id="agent") + await admit_managed_actor(auth, AgentIdentityStore.from_client(database) if database_available else None) + assert auth.agent_id == "agent" + assert auth.managed_agent_policy is None + assert auth.requires_fresh_policy is False diff --git a/tests/test_litellm/proxy/agent_endpoints/test_identity.py b/tests/test_litellm/proxy/agent_endpoints/test_identity.py new file mode 100644 index 00000000000..c9d803fdae7 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_identity.py @@ -0,0 +1,27 @@ +from collections.abc import Mapping + +import pytest +from fastapi import HTTPException + +from litellm.proxy.agent_endpoints.identity import has_legacy_identity, reject_legacy_identity + +TENANT = "11111111-1111-4111-8111-111111111111" +CLIENT = "22222222-2222-4222-8222-222222222222" + + +@pytest.mark.parametrize("params", [None, {}, {"model": "gpt-4o", "api_key": "sk-test"}]) +def test_runtime_params_without_identity_are_accepted(params: Mapping[str, object] | None) -> None: + assert has_legacy_identity(params) is False + reject_legacy_identity(params) + + +@pytest.mark.parametrize( + "identity", [None, {}, {"provider": "microsoft_entra", "tenant_id": TENANT, "client_id": CLIENT}] +) +def test_legacy_litellm_params_identity_is_rejected(identity: object) -> None: + params: Mapping[str, object] = {"model": "gpt-4o", "identity": identity} + assert has_legacy_identity(params) is True + with pytest.raises(HTTPException) as failure: + reject_legacy_identity(params) + assert failure.value.status_code == 400 + assert "top-level identity field" in failure.value.detail diff --git a/tests/test_litellm/proxy/agent_endpoints/test_identity_store.py b/tests/test_litellm/proxy/agent_endpoints/test_identity_store.py new file mode 100644 index 00000000000..005f0b4c074 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_identity_store.py @@ -0,0 +1,450 @@ +from datetime import datetime, timezone +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException +from prisma.models import LiteLLM_VerifiedSubject + +from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore, resolve_managed_agent +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.repositories.table_repositories import ( + AgentIdentityRepository, + AgentsRepository, + VerifiedSubjectRepository, +) +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import ( + AgentIdentityBinding, + AgentIdentityFailure, + ManagedAgentContext, + MicrosoftInteractiveSubject, +) + +TENANT: Final = "11111111-1111-4111-8111-111111111111" +CLIENT: Final = "22222222-2222-4222-8222-222222222222" +PRINCIPAL: Final = "33333333-3333-4333-8333-333333333333" +HUMAN: Final = "44444444-4444-4444-8444-444444444444" +ISSUER: Final = f"https://login.microsoftonline.com/{TENANT}/v2.0" +BINDING: Final = AgentIdentityBinding( + agent_id="agent-one", + provider="microsoft_entra", + tenant_id=TENANT, + client_id=CLIENT, + service_principal_id=PRINCIPAL, + issuer=ISSUER, + required_roles=("Agent.Invoke",), + revision="revision-one", +) +CLAIMS: Final = {"iss": ISSUER, "tid": TENANT, "azp": CLIENT, "oid": PRINCIPAL, "roles": ["Agent.Invoke"]} + + +def stored_agent(**overrides: object) -> AgentResponse: + return AgentResponse.model_validate( + { + "agent_id": "agent-one", + "agent_name": "Research", + "agent_card_params": {}, + "identity": BINDING, + "identity_managed": True, + "execution_mode": "both", + **overrides, + } + ) + + +def setup_store( + agent: AgentResponse | None = stored_agent(), + human: LiteLLM_VerifiedSubject | None = None, + cache: UserApiKeyCache | None = None, +) -> tuple[AgentIdentityStore, AsyncMock, AsyncMock, AsyncMock]: + agents: Final = AsyncMock() + identities: Final = AsyncMock() + humans: Final = AsyncMock() + agents.find_unique.return_value = agent + identities.find_unique.return_value = BINDING + identities.update_many.return_value = 1 + humans.find_unique.return_value = human + db: Final = SimpleNamespace( + db=SimpleNamespace( + litellm_agentstable=agents, + litellm_agentidentity=identities, + litellm_verifiedsubject=humans, + ) + ) + return ( + AgentIdentityStore(AgentsRepository(db), AgentIdentityRepository(db), VerifiedSubjectRepository(db), cache=cache), + agents, + identities, + humans, + ) + + +@pytest.mark.asyncio +async def test_application_authentication_has_no_fabricated_human() -> None: + store, _, _, humans = setup_store() + result: Final = await store.resolve_verified_claims(CLAIMS) + assert isinstance(result, ManagedAgentContext) + assert result.agent_id == "agent-one" + assert result.mode == "autonomous" + assert result.user_id is None + humans.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_shared_binding_lookup_cache_keeps_policy_reads_authoritative() -> None: + cache: Final = UserApiKeyCache() + store, agents, identities, _ = setup_store(cache=cache) + other: Final = AgentIdentityStore(store.agents, store.identities, store.humans, cache=cache) + assert isinstance(await store.resolve_verified_claims(CLAIMS), ManagedAgentContext) + assert isinstance(await other.resolve_verified_claims(CLAIMS), ManagedAgentContext) + identities.find_unique.assert_awaited_once() + assert agents.find_unique.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "changed", + [ + None, + stored_agent(enabled=False), + stored_agent(identity=None), + stored_agent(identity_managed=False), + stored_agent(execution_mode="delegated"), + stored_agent(identity=BINDING.model_copy(update={"active": False})), + stored_agent(identity=BINDING.model_copy(update={"client_id": HUMAN, "revision": "new-binding"})), + stored_agent(identity=BINDING.model_copy(update={"required_roles": ("New.Role",), "revision": "new-policy"})), + ], +) +async def test_lifecycle_is_read_on_every_request_without_cached_allow(changed: AgentResponse | None) -> None: + store, agents, identities, _ = setup_store(cache=UserApiKeyCache()) + agents.find_unique.side_effect = [stored_agent(), changed] + assert isinstance(await store.resolve_verified_claims(CLAIMS), ManagedAgentContext) + denial: Final = await store.resolve_verified_claims(CLAIMS) + assert isinstance(denial, AgentIdentityFailure) + assert denial.code == "identity_denied" + identities.find_unique.assert_awaited_once() + assert agents.find_unique.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unavailable_table", ["agents", "identities", "humans"]) +async def test_identity_store_failure_never_becomes_a_legacy_allow(unavailable_table: str) -> None: + store, agents, identities, humans = setup_store() + table: Final = {"agents": agents, "identities": identities, "humans": humans}[unavailable_table] + table.find_unique.side_effect = RuntimeError("database unavailable") + result: Final = await store.resolve_verified_claims({**CLAIMS, "oid": HUMAN, "scp": "user_impersonation"}) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "policy_unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unavailable_table", ["agents", "humans"]) +async def test_cached_binding_cannot_hide_authoritative_storage_failure(unavailable_table: str) -> None: + store, agents, identities, humans = setup_store(cache=UserApiKeyCache()) + assert isinstance(await store.resolve_verified_claims(CLAIMS), ManagedAgentContext) + table: Final = {"agents": agents, "humans": humans}[unavailable_table] + table.find_unique.side_effect = ConnectionError("writer unavailable") + result: Final = await store.resolve_verified_claims({**CLAIMS, "oid": HUMAN, "scp": "user_impersonation"}) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "policy_unavailable" + identities.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_unclassified_delegated_subject_cannot_authenticate_as_a_user() -> None: + store, _, _, _ = setup_store() + result: Final = await store.resolve_verified_claims( + {**CLAIMS, "oid": HUMAN, "scp": "user_impersonation", "idtyp": "user"} + ) + assert isinstance(result, AgentIdentityFailure) + assert "first sign in" in result.message + + +@pytest.mark.asyncio +async def test_delegated_subject_uses_canonical_sso_user_not_email_claim() -> None: + human: Final = LiteLLM_VerifiedSubject( + kind="human", + subject_id="subject-one", + issuer=ISSUER, + tenant_id=TENANT, + oid=HUMAN, + user_id="canonical-user", + verified_via="sso_interactive", + verified_at=datetime.now(timezone.utc), + ) + store, _, identities, humans = setup_store(human=human, cache=UserApiKeyCache()) + result: Final = await store.resolve_verified_claims( + { + **CLAIMS, + "oid": HUMAN, + "scp": "user_impersonation", + "email": "untrusted-alias@example.com", + } + ) + assert isinstance(result, ManagedAgentContext) + assert result.mode == "delegated" + assert result.user_id == "canonical-user" + humans.find_unique.assert_awaited_once_with( + where={"issuer_tenant_id_oid": {"issuer": ISSUER, "tenant_id": TENANT, "oid": HUMAN}} + ) + humans.find_unique.return_value = None + denied: Final = await store.resolve_verified_claims({**CLAIMS, "oid": HUMAN, "scp": "user_impersonation"}) + assert isinstance(denied, AgentIdentityFailure) + assert denied.code == "identity_denied" + identities.find_unique.assert_awaited_once() + assert humans.find_unique.await_count == 2 + + +@pytest.mark.asyncio +async def test_rebinding_during_authentication_does_not_mark_new_identity_verified() -> None: + store, _, identities, _ = setup_store() + identities.update_many.return_value = 0 + context: Final = ManagedAgentContext(agent_id="agent-one", binding_revision="old-revision", mode="autonomous") + result: Final = await store.record_authentication(context) + assert isinstance(result, AgentIdentityFailure) + assert "changed" in result.message + assert identities.update_many.call_args.kwargs["where"] == { + "agent_id": "agent-one", + "revision": "old-revision", + "active": True, + "agent": {"is": {"enabled": True, "identity_managed": True}}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("agent", [None, stored_agent(identity=None), stored_agent(identity_managed=False)]) +async def test_stale_binding_cannot_bypass_lifecycle(agent: AgentResponse | None) -> None: + store, _, _, _ = setup_store(agent=agent) + assert isinstance(await store.resolve_verified_claims(CLAIMS), AgentIdentityFailure) + + +@pytest.mark.asyncio +async def test_unrelated_non_entra_claims_do_not_query_identity_store() -> None: + store, agents, identities, _ = setup_store() + assert await store.resolve_verified_claims({"sub": "ordinary-user"}) is None + identities.find_unique.assert_not_awaited() + agents.find_unique.assert_not_awaited() + + +HUMAN_CLAIMS: Final = {"iss": ISSUER, "tid": TENANT, "azp": CLIENT, "oid": HUMAN, "scp": "user_impersonation"} + + +@pytest.mark.asyncio +async def test_bound_agents_and_policy_failures_are_never_served_from_the_miss_cache() -> None: + store, _, identities, _ = setup_store() + assert isinstance(await store.resolve_verified_claims(CLAIMS), ManagedAgentContext) + assert isinstance(await store.resolve_verified_claims(CLAIMS), ManagedAgentContext) + assert identities.find_unique.await_count == 2 + identities.find_unique.side_effect = ConnectionError("database down") + assert isinstance(await store.resolve_verified_claims(CLAIMS), AgentIdentityFailure) + assert isinstance(await store.resolve_verified_claims(CLAIMS), AgentIdentityFailure) + assert identities.find_unique.await_count == 4 + + +@pytest.mark.asyncio +async def test_retired_client_cannot_fall_back_to_ordinary_user_authentication() -> None: + from prisma.models import LiteLLM_RetiredAgentIdentity + + from litellm.repositories.table_repositories import RetiredAgentIdentityRepository + + identities: Final = AsyncMock() + identities.find_unique.return_value = None + retired: Final = AsyncMock() + retired.find_unique.return_value = LiteLLM_RetiredAgentIdentity( + binding_id="retired", + agent_id="agent-one", + provider="microsoft_entra", + issuer=ISSUER, + tenant_id=TENANT, + client_id=CLIENT, + ) + db: Final = SimpleNamespace( + db=SimpleNamespace( + litellm_agentidentity=identities, + litellm_retiredagentidentity=retired, + litellm_agentstable=AsyncMock(), + litellm_verifiedsubject=AsyncMock(), + ) + ) + store: Final = AgentIdentityStore( + AgentsRepository(db), + AgentIdentityRepository(db), + VerifiedSubjectRepository(db), + RetiredAgentIdentityRepository(db), + ) + result: Final = await store.resolve_verified_claims({**CLAIMS, "oid": HUMAN, "scp": "user_impersonation"}) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "identity_denied" + assert "retired" in result.message + + +@pytest.mark.asyncio +async def test_missing_revision_cannot_create_entra_authentication_evidence() -> None: + store, _, identities, _ = setup_store() + result: Final = await store.record_authentication(ManagedAgentContext(agent_id="agent-one", mode="autonomous")) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "identity_denied" + identities.update_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_authentication_evidence_write_failure_is_not_success() -> None: + store, _, identities, _ = setup_store() + identities.update_many.side_effect = RuntimeError("writer unavailable") + result: Final = await store.record_authentication( + ManagedAgentContext(agent_id="agent-one", binding_revision="revision-one", mode="autonomous") + ) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "policy_unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unavailable", [True, False]) +async def test_retired_binding_denies_and_history_outage_cannot_become_legacy_fallback(unavailable: bool) -> None: + + database: Final = MagicMock() + database.writer_db.litellm_agentidentity.find_unique = AsyncMock(return_value=None) + database.writer_db.litellm_verifiedsubject.find_unique = AsyncMock(return_value=None) + database.writer_db.litellm_retiredagentidentity.find_unique = AsyncMock( + return_value={"client_id": CLIENT}, side_effect=RuntimeError("unavailable") if unavailable else None + ) + result: Final = await AgentIdentityStore.from_client(database).resolve_verified_claims(CLAIMS) + assert isinstance(result, AgentIdentityFailure) + assert result.code == ("policy_unavailable" if unavailable else "identity_denied") + assert result.message == ( + "Retired agent identity could not be checked" if unavailable else "This agent identity binding has been retired" + ) + + +@pytest.mark.asyncio +async def test_new_binding_is_enforced_after_another_worker_commits_it() -> None: + _, agents, identities, humans = setup_store() + identities.find_unique.return_value = None + retired: Final = AsyncMock() + retired.find_unique.return_value = None + db: Final = SimpleNamespace( + writer_db=SimpleNamespace( + litellm_agentstable=agents, + litellm_agentidentity=identities, + litellm_verifiedsubject=humans, + litellm_retiredagentidentity=retired, + litellm_retiredagent=retired, + ) + ) + worker: Final = AgentIdentityStore.from_client(db, cache=UserApiKeyCache()) + claims: Final = {**CLAIMS, "oid": "55555555-5555-4555-8555-555555555555"} + assert await worker.resolve_verified_claims(claims) is None + identities.find_unique.return_value = BINDING + denied: Final = await worker.resolve_verified_claims(claims) + assert isinstance(denied, AgentIdentityFailure) + assert denied.code == "identity_denied" + assert "Application token contradicts" in denied.message + assert identities.find_unique.await_count == 2 + + +@pytest.mark.asyncio +async def test_non_string_subject_does_not_query_directory_ownership() -> None: + store, _, _, humans = setup_store() + assert await store.subject(ISSUER, TENANT, None) is None + humans.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("configured", [False, True]) +async def test_missing_or_unavailable_retirement_history_fails_closed(configured: bool) -> None: + database: Final = MagicMock() + database.writer_db.litellm_retiredagent.find_unique = AsyncMock(side_effect=RuntimeError("history unavailable")) + store: Final = AgentIdentityStore.from_client(database) if configured else setup_store()[0] + result: Final = await store.retired_agent("deleted-agent") + assert isinstance(result, AgentIdentityFailure) + assert result.code == "policy_unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("owner", ["canonical-user", "another-user"]) +async def test_interactive_enrollment_preserves_existing_subject_ownership(owner: str) -> None: + store, _, _, humans = setup_store() + humans.upsert.return_value = LiteLLM_VerifiedSubject( + subject_id="subject-one", + issuer=ISSUER, + tenant_id=TENANT, + oid=HUMAN, + user_id=owner, + kind="human", + verified_via="sso_interactive", + verified_at=datetime.now(timezone.utc), + ) + result: Final = await store.enroll_interactive_human( + MicrosoftInteractiveSubject(issuer=ISSUER, tenant_id=TENANT, oid=HUMAN), "canonical-user" + ) + if owner == "canonical-user": + assert result is None + else: + assert isinstance(result, AgentIdentityFailure) + assert result.code == "identity_denied" + assert humans.upsert.call_args.kwargs["data"]["update"] == {} + assert humans.upsert.call_args.kwargs["data"]["create"]["user_id"] == "canonical-user" + + +@pytest.mark.asyncio +async def test_interactive_enrollment_outage_fails_closed() -> None: + store, _, _, humans = setup_store() + humans.upsert.side_effect = ConnectionError("writer unavailable") + result: Final = await store.enroll_interactive_human( + MicrosoftInteractiveSubject(issuer=ISSUER, tenant_id=TENANT, oid=HUMAN), "canonical-user" + ) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "policy_unavailable" + + +@pytest.mark.asyncio +async def test_matching_revision_records_successful_authentication() -> None: + store, _, identities, _ = setup_store() + assert ( + await store.record_authentication( + ManagedAgentContext(agent_id="agent-one", binding_revision="revision-one", mode="autonomous") + ) + is None + ) + identities.update_many.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outage", [False, True]) +async def test_resolver_maps_denials_and_outages_to_public_errors(outage: bool) -> None: + database: Final = MagicMock() + database.writer_db.litellm_agentidentity.find_unique = AsyncMock( + return_value=BINDING, side_effect=ConnectionError("unavailable") if outage else None + ) + database.writer_db.litellm_verifiedsubject.find_unique = AsyncMock(return_value=None) + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=stored_agent(enabled=False)) + with pytest.raises(HTTPException) as exc: + await resolve_managed_agent(CLAIMS, database) + assert exc.value.status_code == (503 if outage else 403) + + +@pytest.mark.asyncio +async def test_resolver_preserves_unconfigured_and_unrelated_authentication() -> None: + assert await resolve_managed_agent(CLAIMS, None) is None + assert await resolve_managed_agent({"sub": "ordinary-user"}, MagicMock()) is None + store, _, identities, _ = setup_store() + identities.find_unique.return_value = None + assert await store.resolve_verified_claims(CLAIMS) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("registered", [True, False]) +async def test_application_and_unregistered_clients_do_not_depend_on_human_subject_storage(registered: bool) -> None: + store, _, identities, humans = setup_store() + identities.find_unique.return_value = BINDING if registered else None + humans.find_unique.side_effect = RuntimeError("subject database unavailable") + result: Final = await store.resolve_verified_claims(CLAIMS) + if registered: + assert isinstance(result, ManagedAgentContext) + assert result.mode == "autonomous" + assert result.user_id is None + else: + assert result is None + humans.find_unique.assert_not_awaited() diff --git a/tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py b/tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py new file mode 100644 index 00000000000..45fe4b0655f --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py @@ -0,0 +1,257 @@ +from typing import Final + +import pytest + +from litellm.proxy.agent_endpoints.managed_identity import classify_agent_subject, managed_write_fields +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import ( + AgentExecutionMode, + AgentIdentityBinding, + AgentIdentityFailure, + AgentSubject, +) + +TENANT: Final = "11111111-1111-4111-8111-111111111111" +CLIENT: Final = "22222222-2222-4222-8222-222222222222" +PRINCIPAL: Final = "33333333-3333-4333-8333-333333333333" +HUMAN: Final = "44444444-4444-4444-8444-444444444444" +ISSUER: Final = f"https://login.microsoftonline.com/{TENANT}/v2.0" +BINDING: Final = AgentIdentityBinding( + agent_id="agent-one", + provider="microsoft_entra", + tenant_id=TENANT, + client_id=CLIENT, + service_principal_id=PRINCIPAL, + issuer=ISSUER, + required_roles=("Agent.Invoke",), + required_scopes=("user_impersonation",), + revision="binding-one", +) + + +def claims(**overrides: object) -> dict[str, object]: + return {"iss": ISSUER, "tid": TENANT, "azp": CLIENT, "oid": PRINCIPAL, "roles": ["Agent.Invoke"], **overrides} + + +def test_autonomous_identity_needs_no_human_and_checks_the_pinned_principal() -> None: + result: Final = classify_agent_subject(BINDING, claims(), "autonomous") + assert result == AgentSubject(kind="application", oid=PRINCIPAL, mode="autonomous") + assert isinstance(classify_agent_subject(BINDING, claims(oid=HUMAN), "autonomous"), AgentIdentityFailure) + + +@pytest.mark.parametrize( + "overrides", + [ + {"iss": "https://untrusted.example"}, + {"tid": CLIENT}, + {"azp": TENANT}, + {"roles": []}, + {"idtyp": "user"}, + {"scp": "user_impersonation"}, + {"scp": 1}, + {"oid": None}, + ], +) +def test_application_rejects_mismatched_or_contradictory_verified_claims(overrides: dict[str, object]) -> None: + assert isinstance(classify_agent_subject(BINDING, claims(**overrides), "both"), AgentIdentityFailure) + + +def test_delegated_profile_identifies_a_subject_without_asserting_that_it_is_human() -> None: + result: Final = classify_agent_subject(BINDING, claims(oid=HUMAN, scp="user_impersonation"), "delegated") + assert result == AgentSubject(kind="delegated_subject", oid=HUMAN, mode="delegated") + + +@pytest.mark.parametrize( + "overrides", + [ + {"scp": "unrelated"}, + {"scp": ""}, + {"idtyp": "app"}, + {"xms_sub_fct": "2 13 15"}, + {"xms_sub_fct": [13]}, + ], +) +def test_delegated_profile_rejects_unknown_scope_and_known_nonhuman_subjects(overrides: dict[str, object]) -> None: + assert isinstance( + classify_agent_subject(BINDING, claims(**{"oid": HUMAN, "scp": "user_impersonation", **overrides}), "both"), + AgentIdentityFailure, + ) + + +def test_allowed_mode_cannot_be_selected_by_the_caller() -> None: + assert isinstance(classify_agent_subject(BINDING, claims(), "delegated"), AgentIdentityFailure) + assert isinstance( + classify_agent_subject(BINDING, claims(oid=HUMAN, scp="user_impersonation"), "autonomous"), + AgentIdentityFailure, + ) + + +def test_native_facet_absence_does_not_establish_human_identity() -> None: + result: Final = classify_agent_subject( + BINDING, claims(oid=HUMAN, scp="user_impersonation", xms_sub_fct="113"), "both" + ) + assert isinstance(result, AgentSubject) + assert result.kind == "delegated_subject" + + +def managed_agent() -> AgentResponse: + return AgentResponse( + agent_id="agent-one", agent_name="Research", agent_card_params={}, identity=BINDING, identity_managed=True + ) + + +def test_unbinding_keeps_managed_state_and_disables_agent() -> None: + result: Final = managed_write_fields({"identity": None, "enabled": True}, managed_agent(), "admin") + assert not isinstance(result, AgentIdentityFailure) + assert result["identity_managed"] is True + assert result["enabled"] is False + assert result["identity"]["update"]["active"] is False + assert result["identity"]["update"]["last_authenticated_at"] is None + assert result["identity"]["update"]["revision"] != BINDING.revision + + +def test_rename_does_not_rewrite_binding_or_evidence() -> None: + assert managed_write_fields({"agent_name": "Renamed"}, managed_agent(), "admin") == {} + + +def test_autonomous_binding_requires_enterprise_application_object_id() -> None: + result: Final = managed_write_fields( + {"identity": {"provider": "microsoft_entra", "tenant_id": TENANT, "client_id": CLIENT}}, None, "admin" + ) + assert isinstance(result, AgentIdentityFailure) + assert "service-principal" in result.message + + +def test_rebinding_clears_evidence_and_uses_atomic_nested_write() -> None: + result: Final = managed_write_fields( + { + "identity": { + "provider": "microsoft_entra", + "tenant_id": TENANT, + "client_id": CLIENT, + "service_principal_id": PRINCIPAL, + } + }, + managed_agent(), + "admin", + ) + assert not isinstance(result, AgentIdentityFailure) + assert result["identity_managed"] is True + assert "upsert" in result["identity"] + assert result["identity"]["upsert"]["update"]["revision"] != BINDING.revision + assert result["identity"]["upsert"]["update"]["last_authenticated_at"] is None + + +def test_unbound_identity_can_be_reactivated_with_the_same_application() -> None: + disabled: Final = managed_agent().model_copy( + update={"identity": BINDING.model_copy(update={"active": False}), "enabled": False} + ) + configuration: Final = BINDING.model_dump( + exclude={"agent_id", "issuer", "revision", "last_authenticated_at", "active"} + ) + result: Final = managed_write_fields({"identity": configuration, "enabled": True}, disabled, "admin") + assert not isinstance(result, AgentIdentityFailure) + assert result["enabled"] is True + assert result["identity"]["upsert"]["update"]["active"] is True + assert result["identity"]["upsert"]["update"]["revision"] != BINDING.revision + + +def test_each_application_binding_records_its_history_atomically() -> None: + configuration: Final = BINDING.model_dump( + exclude={"agent_id", "issuer", "revision", "last_authenticated_at", "active"} + ) + created: Final = managed_write_fields({"identity": configuration}, None, "admin") + assert not isinstance(created, AgentIdentityFailure) + assert created["retired_identities"]["connectOrCreate"]["create"]["client_id"] == CLIENT + replacement: Final = managed_write_fields( + {"identity": {**configuration, "client_id": HUMAN}}, managed_agent(), "admin" + ) + assert not isinstance(replacement, AgentIdentityFailure) + assert replacement["retired_identities"]["connectOrCreate"]["create"]["client_id"] == HUMAN + + +def test_unchanged_binding_preserves_revision_and_authentication_evidence() -> None: + configuration: Final = BINDING.model_dump( + exclude={"agent_id", "issuer", "revision", "last_authenticated_at", "active"} + ) + assert managed_write_fields({"identity": configuration}, managed_agent(), "admin") == {} + + +@pytest.mark.parametrize("identity", [None, BINDING.model_copy(update={"active": False})]) +def test_enabling_unbound_or_inactive_identity_requires_rebinding(identity: AgentIdentityBinding | None) -> None: + agent: Final = managed_agent().model_copy(update={"identity": identity, "enabled": False}) + result: Final = managed_write_fields({"enabled": True}, agent, "admin") + assert isinstance(result, AgentIdentityFailure) + assert "Bind an identity" in result.message + + +@pytest.mark.parametrize("mode", ["delegated", "both"]) +def test_explicit_empty_scope_requirements_can_be_registered_and_preserved(mode: str) -> None: + from litellm.types.proxy.agent_identity import EntraIdentityConfig + + configuration: Final = EntraIdentityConfig( + provider="microsoft_entra", + tenant_id=TENANT, + client_id=CLIENT, + service_principal_id=PRINCIPAL, + required_scopes=(), + ) + created: Final = managed_write_fields( + {"identity": configuration.model_dump(), "execution_mode": mode}, None, "admin" + ) + assert not isinstance(created, AgentIdentityFailure) + assert created["identity"]["create"]["required_scopes"] == () + agent: Final = managed_agent().model_copy(update={"identity": BINDING.model_copy(update={"required_scopes": ()})}) + updated: Final = managed_write_fields({"execution_mode": mode}, agent, "admin") + assert not isinstance(updated, AgentIdentityFailure) + assert updated["execution_mode"] == mode + + +@pytest.mark.parametrize( + "incoming", + [ + {"identity": {"provider": "microsoft_entra", "tenant_id": "invalid", "client_id": CLIENT}}, + {"execution_mode": "unknown"}, + ], +) +def test_invalid_identity_configuration_returns_a_public_validation_failure(incoming: dict[str, object]) -> None: + result: Final = managed_write_fields(incoming, None, "admin") + assert isinstance(result, AgentIdentityFailure) + assert result.code == "identity_denied" + assert result.message.startswith("Invalid agent identity configuration:") + + +@pytest.mark.parametrize("roles", ["Agent.Invoke", [42], None]) +def test_malformed_application_roles_are_rejected(roles: object) -> None: + result: Final = classify_agent_subject(BINDING, claims(roles=roles), "autonomous") + assert isinstance(result, AgentIdentityFailure) + assert "Invalid application roles" in result.message + + +def test_entra_binding_normalizes_identifiers_and_rejects_invalid_configuration() -> None: + from pydantic import ValidationError + + from litellm.types.proxy.agent_identity import EntraIdentityConfig + + identifier = "ABCDEF00-1234-4234-9234-123456789ABC" + config = EntraIdentityConfig(provider="microsoft_entra", tenant_id=identifier, client_id=identifier) + assert config.tenant_id == identifier.lower() + assert config.client_id == identifier.lower() + assert config.service_principal_id is None + assert config.issuer == f"https://login.microsoftonline.com/{config.tenant_id}/v2.0" + with pytest.raises(ValidationError): + EntraIdentityConfig(provider="microsoft_entra", tenant_id="invalid", client_id=identifier) + + +@pytest.mark.parametrize("mode", ["delegated", "both"]) +def test_empty_required_scopes_allow_valid_delegated_scope(mode: AgentExecutionMode) -> None: + binding: Final = BINDING.model_copy(update={"required_scopes": ()}) + result: Final = classify_agent_subject(binding, claims(oid=HUMAN, scp="custom_scope"), mode) + assert result == AgentSubject(kind="delegated_subject", oid=HUMAN, mode="delegated") + + +@pytest.mark.parametrize("scope", [None, "", " \t ", 42]) +def test_empty_requirements_do_not_make_a_scope_less_human_token_valid(scope: object) -> None: + binding: Final = BINDING.model_copy(update={"required_scopes": ()}) + result: Final = classify_agent_subject(binding, claims(oid=HUMAN, scp=scope), "both") + assert isinstance(result, AgentIdentityFailure) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d3cb8e4a645..353249dddf0 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1,5 +1,7 @@ import asyncio +import base64 import json +import re import sys import time from collections.abc import Iterator, Mapping @@ -38,6 +40,7 @@ from litellm.proxy._types import ( from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver from litellm.types.agents import AgentCaller from litellm.proxy.auth.auth_checks import ( + LITELLM_SESSION_TOKEN_PREFIX, ExperimentalUIJWTToken, _cache_management_object, _can_object_call_model, @@ -76,7 +79,9 @@ from litellm.constants import ( TAG_REGISTRY_MAX_SIZE, ) from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.proxy.auth.user_api_key_auth import check_api_key_for_custom_headers_or_pass_through_endpoints +from litellm.proxy import proxy_server +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_bearer_token, encrypt_value_helper from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from prisma.errors import DataError from litellm.proxy.common_utils.user_api_key_cache import ( @@ -149,7 +154,7 @@ def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_v token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) # Check that decrypted_token is not None before using json.loads assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -175,7 +180,7 @@ def test_get_cli_jwt_auth_token_includes_team_alias(valid_sso_user_defined_value team_alias="test-team", ) - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -202,7 +207,7 @@ def test_get_cli_jwt_auth_token_carries_team_grants_not_user_allowlist( team_model_aliases={"team-fast": "gpt-4.1-mini"}, ) - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -219,7 +224,7 @@ def test_get_cli_jwt_auth_token_keeps_user_allowlist_when_no_team( """A session token with no team bound still carries the user's own allowlist.""" token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -233,7 +238,7 @@ def test_get_experimental_ui_login_jwt_auth_token_uses_10_min_expiry( ): """Test that Experimental UI token uses fixed 10-minute expiry (does not use LITELLM_UI_SESSION_DURATION).""" token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) @@ -251,7 +256,7 @@ def test_experimental_ui_token_ignores_litellm_ui_session_duration( was incorrectly wired to the experimental flow.""" # Default LITELLM_UI_SESSION_DURATION is "24h" - token must still expire in ~10 min token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values) - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) @@ -288,6 +293,51 @@ def test_get_key_object_from_ui_hash_key_valid(valid_sso_user_defined_values, mo assert key_object.max_budget == litellm.max_ui_session_budget +@pytest.mark.parametrize("encryption_algorithm", ["xsalsa20-poly1305", "aes-256-gcm"]) +def test_get_key_object_from_ui_hash_key_accepts_only_minted_session_tokens( + valid_sso_user_defined_values, monkeypatch, encryption_algorithm +): + monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": encryption_algorithm}) + session_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + stored_value = encrypt_value_helper(json.dumps({"user_role": LitellmUserRoles.PROXY_ADMIN.value})) + + key_object = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(session_token) + assert key_object is not None + assert key_object.user_role == LitellmUserRoles.PROXY_ADMIN + reshaped = LITELLM_SESSION_TOKEN_PREFIX + stored_value.removeprefix("v2:gcm:").rstrip("=") + for candidate in (stored_value, reshaped): + assert ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(candidate) is None + + +def test_session_tokens_are_header_safe_and_never_look_like_virtual_keys(valid_sso_user_defined_values): + for token in ( + ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(valid_sso_user_defined_values), + ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values), + ): + assert re.fullmatch(r"litellm_login_[A-Za-z0-9_-]+", token), token + assert ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(token) is not None + + +@pytest.mark.asyncio +async def test_session_token_survives_langfuse_basic_auth_parsing(valid_sso_user_defined_values): + session_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + basic_credentials = base64.b64encode(f"{session_token}:sk-lf-secret".encode()).decode() + request = MagicMock() + request.headers = {} + + api_key = await check_api_key_for_custom_headers_or_pass_through_endpoints( + request=request, + route="/api/public/ingestion", + pass_through_endpoints=[ + {"path": "/api/public/ingestion", "target": "https://example.com", "custom_auth_parser": "langfuse"} + ], + api_key=f"Basic {basic_credentials}", + ) + + assert api_key == session_token + assert ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(session_token) is not None + + def test_get_key_object_from_ui_hash_key_invalid(): """Test getting key object from invalid UI hash key""" # Test with invalid token @@ -801,7 +851,7 @@ def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -841,7 +891,7 @@ def test_get_cli_jwt_auth_token_custom_expiration(valid_sso_user_defined_values, token = auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) # Decrypt and verify token contents - decrypted_token = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted_token = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted_token is not None token_data = json.loads(decrypted_token) @@ -859,7 +909,7 @@ def test_get_cli_jwt_auth_token_unique_per_session(valid_sso_user_defined_values from litellm.constants import CLI_SESSION_KEY_PREFIX def _decode(token: str) -> dict: - decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted is not None return json.loads(decrypted) @@ -879,7 +929,7 @@ def test_get_cli_jwt_auth_token_applies_fallback_budget(valid_sso_user_defined_v token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( valid_sso_user_defined_values, max_budget=litellm.max_ui_session_budget ) - decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted is not None assert json.loads(decrypted).get("max_budget") == litellm.max_ui_session_budget @@ -888,7 +938,7 @@ def test_get_cli_jwt_auth_token_no_fallback_when_budget_provided( valid_sso_user_defined_values, ): token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values, max_budget=None) - decrypted = decrypt_value_helper(token, key="ui_hash_key", exception_type="debug") + decrypted = decrypt_bearer_token(token, prefix=LITELLM_SESSION_TOKEN_PREFIX) assert decrypted is not None assert json.loads(decrypted).get("max_budget") is None @@ -1091,7 +1141,7 @@ async def test_get_user_object_check_db_only_ignores_recent_miss(monkeypatch): monkeypatch.setitem(auth_checks.last_db_access_time, f"user_id:{user_id}", (None, time.time())) db_row = LiteLLM_UserTable(user_id=user_id, user_email=None, user_role="internal_user") mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=db_row) + mock_prisma_client.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=db_row) result = await get_user_object( user_id=user_id, @@ -1103,7 +1153,7 @@ async def test_get_user_object_check_db_only_ignores_recent_miss(monkeypatch): assert result is not None assert result.user_id == user_id - mock_prisma_client.db.litellm_usertable.find_unique.assert_awaited_once() + mock_prisma_client.writer_db.litellm_usertable.find_unique.assert_awaited_once() @pytest.mark.asyncio @@ -3058,7 +3108,7 @@ async def test_get_team_object_raises_404_when_not_found(): mock_prisma_client = MagicMock() mock_db = AsyncMock() mock_prisma_client.db = mock_db - mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=None) mock_cache = MagicMock() mock_cache.async_get_cache = AsyncMock(return_value=None) @@ -3076,11 +3126,40 @@ async def test_get_team_object_raises_404_when_not_found(): assert "Team doesn't exist in db" in str(exc_info.value.detail) +@pytest.mark.asyncio +async def test_get_team_object_check_db_only_reads_writer_through_the_shared_loader(): + """Management endpoints mock ``_get_team_object_from_user_api_key_cache`` and expect + ``check_db_only`` to still flow through it; only the table it reads moves to the writer.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import get_team_object + + row = {"team_id": "team-writer", "models": ["gpt-4o"], "object_permission_id": None} + prisma = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=SimpleNamespace(dict=lambda: row)) + prisma.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=SimpleNamespace(dict=lambda: row)) + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=None) + cache.async_set_cache = AsyncMock() + shared_loader = AsyncMock(wraps=auth_checks._get_team_object_from_user_api_key_cache) + + with patch.object(auth_checks, "_get_team_object_from_user_api_key_cache", shared_loader): + team = await get_team_object("team-writer", prisma, cache, check_db_only=True) + + assert team.team_id == "team-writer" + assert shared_loader.await_args.kwargs["use_writer"] is True + prisma.writer_db.litellm_teamtable.find_unique.assert_awaited_once() + prisma.db.litellm_teamtable.find_unique.assert_not_awaited() + cache.async_set_cache.assert_awaited_once() + + def _mock_prisma_for_team_lookup(find_unique): from unittest.mock import MagicMock mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_teamtable.find_unique = find_unique + mock_prisma_client.writer_db.litellm_teamtable.find_unique = find_unique return mock_prisma_client @@ -5621,7 +5700,8 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): team_table = LiteLLM_TeamTableCachedObj(**base_team_row) cache = MagicMock() cache.async_set_cache = AsyncMock() - cache.delete_cache = MagicMock() + cache.async_delete_cache = AsyncMock() + cache.async_delete_cache_pre_call = AsyncMock(return_value=None) # no request pipeline open logging_obj = MagicMock() logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() @@ -5642,9 +5722,9 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): written_value = cache.async_set_cache.await_args.kwargs.get("value") or cache.async_set_cache.await_args.args[1] assert written_value is team_table - # (2) team_alias-keyed entry is deleted in BOTH the in-memory cache - # and the Redis dual cache (mirrors _delete_cache_key_object pattern). - cache.delete_cache.assert_called_once_with(key="team_alias:H-Capacity") + # (2) team_alias-keyed entry is deleted in BOTH the in-memory cache and the Redis dual cache, on the + # async path: a Redis DEL must never run synchronously on the event loop. + cache.async_delete_cache.assert_awaited_once_with(key="team_alias:H-Capacity") # (4) internal usage cache: team_id entry deleted BEFORE the fresh # write, alias entry deleted as before. @@ -5658,7 +5738,7 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): aliasless = LiteLLM_TeamTableCachedObj(**{**base_team_row, "team_alias": None}) cache2 = MagicMock() cache2.async_set_cache = AsyncMock() - cache2.delete_cache = MagicMock() + cache2.async_delete_cache = AsyncMock() logging_obj2 = MagicMock() logging_obj2.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() @@ -5669,7 +5749,7 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): proxy_logging_obj=logging_obj2, ) - cache2.delete_cache.assert_not_called() + cache2.async_delete_cache.assert_not_awaited() logging_obj2.internal_usage_cache.dual_cache.async_delete_cache.assert_awaited_once_with( key="team_id:team-no-alias" ) @@ -10003,3 +10083,196 @@ def test_can_object_call_model_allows_listed_model_for_key(): ) assert result is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("allowed", [True, False]) +async def test_authoritative_access_group_reads_writer_despite_stale_allow_cache(allowed: bool) -> None: + from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.auth.auth_checks import get_access_object + + stale: Final = LiteLLM_AccessGroupTable(access_group_id="group", access_group_name="Policy", access_model_names=["old"]) + current: Final = stale.model_copy(update={"access_model_names": ["new"] if allowed else []}) + client: Final = MagicMock() + client.writer_db.litellm_accessgrouptable.find_unique = AsyncMock(return_value=current) + client.db.litellm_accessgrouptable.find_unique = AsyncMock(return_value=stale) + cache: Final = MagicMock() + cache.async_get_cache = AsyncMock(return_value=stale) + cache.async_set_cache = AsyncMock() + result: Final = await get_access_object("group", client, cache, check_db_only=True) + assert result.access_model_names == (["new"] if allowed else []) + cache.async_get_cache.assert_not_awaited() + client.db.litellm_accessgrouptable.find_unique.assert_not_awaited() + client.writer_db.litellm_accessgrouptable.find_unique.assert_awaited_once_with(where={"access_group_id": "group"}) + + +@pytest.mark.asyncio +async def test_authoritative_access_group_outage_does_not_use_cached_grants() -> None: + from fastapi import HTTPException + + from litellm.proxy.auth.auth_checks import get_access_object + + client: Final = MagicMock() + client.writer_db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=RuntimeError("writer unavailable")) + cache: Final = MagicMock() + cache.async_get_cache = AsyncMock() + with pytest.raises(HTTPException) as failure: + await get_access_object("group", client, cache, check_db_only=True) + assert failure.value.status_code == 503 + assert failure.value.detail == "Access group policy is unavailable" + cache.async_get_cache.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_authoritative_team_permission_outage_cannot_drop_the_teams_restrictions() -> None: + from fastapi import HTTPException + + from litellm.proxy.auth.auth_checks import get_team_object + + row: Final = LiteLLM_TeamTable(team_id="team-policy-outage", object_permission_id="team-permission") + client: Final = MagicMock() + client.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=row) + client.writer_db.litellm_objectpermissiontable.find_unique = AsyncMock(side_effect=RuntimeError("unavailable")) + cache: Final = MagicMock() + cache.async_get_cache = AsyncMock() + cache.async_set_cache = AsyncMock() + with pytest.raises(HTTPException) as failure: + await get_team_object(row.team_id, client, cache, check_db_only=True) + assert failure.value.status_code == 404 + client.writer_db.litellm_objectpermissiontable.find_unique.assert_awaited_once() + cache.async_set_cache.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strict", [True, False]) +@pytest.mark.parametrize("missing", [True, False]) +async def test_referenced_permission_failures_preserve_legacy_behavior_and_deny_strict_reads(strict, missing): + from fastapi import HTTPException + + from litellm.proxy.auth.auth_checks import get_object_permission + + client = MagicMock() + lookup = AsyncMock(return_value=None, side_effect=None if missing else RuntimeError("unavailable")) + client.writer_db.litellm_objectpermissiontable.find_unique = lookup + client.db.litellm_objectpermissiontable.find_unique = lookup + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=None) + if strict: + with pytest.raises(HTTPException if missing else RuntimeError): + await get_object_permission("referenced", client, cache, check_db_only=True) + cache.async_get_cache.assert_not_awaited() + else: + assert await get_object_permission("referenced", client, cache) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "models,key_aliases,team_aliases,allowed", + [ + (["fast"], {}, {}, True), + ([], {}, {}, False), + (["other"], {}, {}, False), + (["target"], {"fast": "target"}, {}, True), + (["target"], {}, {"fast": "target"}, True), + (["fast"], {}, {"fast": "forbidden"}, False), + ], +) +async def test_managed_agent_model_policy_checks_dispatched_model( + models: list[str], key_aliases: dict[str, str], team_aliases: dict[str, str], allowed: bool +) -> None: + from fastapi import HTTPException + + from litellm.proxy.auth.auth_checks import common_checks + from litellm.types.agents import AgentResponse + + agent: Final = AgentResponse( + agent_id="managed", agent_name="Managed", agent_card_params={}, object_permission={"models": models} + ) + auth: Final = UserAPIKeyAuth( + token="test-token", team_id="team", aliases=key_aliases, team_model_aliases=team_aliases + ) + auth.managed_agent_policy = agent + checks: Final = common_checks( + request_body={"model": "fast", "messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=auth, + request=MagicMock(spec=Request), + ) + if allowed: + assert await checks is True + else: + with pytest.raises((HTTPException, ModelAccessDeniedProxyException)) as failure: + await checks + assert str(getattr(failure.value, "status_code", getattr(failure.value, "code", None))) == "403" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reconnect", (False, True)) +async def test_authoritative_key_load_bypasses_warm_key_and_permission_caches(reconnect: bool) -> None: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, object_permission_cache_key + + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="current", agents=["allowed"]) + stale: Final = UserAPIKeyAuth(token="hash", team_id="old-team", object_permission_id="old") + current: Final = UserAPIKeyAuth(token="hash", team_id="new-team", object_permission_id="current") + cache: Final = UserApiKeyCache() + cache.set_cache("hash", stale) + cache.set_cache(object_permission_cache_key("current"), permission.model_copy(update={"agents": ["revoked"]})) + database: Final = MagicMock() + database.get_data = AsyncMock(side_effect=[httpx.ConnectError("reset"), current] if reconnect else [current]) + database.attempt_db_reconnect = AsyncMock(return_value=True) + database.writer_db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=permission) + fresh: Final = await get_key_object("hash", database, cache, check_db_only=True) + assert fresh.team_id == "new-team" + assert fresh.object_permission == permission + assert all(call.kwargs["use_writer"] is True for call in database.get_data.await_args_list) + database.db.litellm_objectpermissiontable.find_unique.assert_not_called() + cached: Final = await get_key_object("hash", database, cache) + assert cached.team_id == "old-team" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("missing", (False, True)) +async def test_authoritative_key_cannot_keep_grants_when_permission_is_unavailable(missing: bool) -> None: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + database: Final = MagicMock() + database.get_data = AsyncMock(return_value=UserAPIKeyAuth( + object_permission_id="grant", object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="grant", agents=["allowed"]) + )) + database.writer_db.litellm_objectpermissiontable.find_unique = AsyncMock( + return_value=None, side_effect=None if missing else RuntimeError("writer unavailable") + ) + with pytest.raises(Exception, match=r"does not exist|unavailable"): + await get_key_object("hash", database, UserApiKeyCache(), check_db_only=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strict", [False, True]) +async def test_authoritative_group_grants_propagate_policy_outages( + monkeypatch: pytest.MonkeyPatch, strict: bool +) -> None: + from unittest.mock import AsyncMock, MagicMock + + from fastapi import HTTPException + + from litellm.proxy import proxy_server + from litellm.proxy.auth.auth_checks import _get_agent_ids_from_access_groups + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + database: Final = MagicMock() + database.db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + database.writer_db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache()) + if strict: + with pytest.raises(HTTPException): + await _get_agent_ids_from_access_groups(["group"], check_db_only=True) + else: + assert await _get_agent_ids_from_access_groups(["group"]) == [] diff --git a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py index 0fd0dda3017..ffac95d6815 100644 --- a/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py +++ b/tests/test_litellm/proxy/auth/test_auth_object_prefetch.py @@ -18,13 +18,18 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ( + get_end_user_object, get_org_object, get_team_membership, get_team_object, get_user_object, ) -from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects, prefetch_identity_keys +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + end_user_cache_key, + end_user_restricted_registry_cache_key, +) USER_ID = "prefetch-user" TEAM_ID = "prefetch-team" @@ -336,3 +341,29 @@ async def test_no_redis_goes_straight_to_one_query(): assert prisma.db.query_first.await_count == 1 assert cache.in_memory_cache.get_cache(f"team_membership:{USER_ID}:{TEAM_ID}") is not None + + +@pytest.mark.asyncio +async def test_identity_prefetch_warms_the_end_user_so_its_getter_needs_neither_redis_nor_the_database(): + end_user_key = end_user_cache_key("eu-1") + redis = CountingRedis({end_user_key: json.dumps({"user_id": "eu-1", "blocked": False, "spend": 0.0})}) + cache = _cache(redis) + prisma = _prisma() + + await prefetch_identity_keys([end_user_key, end_user_restricted_registry_cache_key()], cache) + end_user = await get_end_user_object(end_user_id="eu-1", prisma_client=prisma, user_api_key_cache=cache) + + assert end_user is not None and end_user.user_id == "eu-1" + assert redis.commands == [f"MGET {end_user_key} {end_user_restricted_registry_cache_key()}"] + assert prisma.db.mock_calls == [] + + +@pytest.mark.asyncio +async def test_identity_prefetch_does_not_cache_an_absent_entry_as_present(): + redis = CountingRedis({}) + cache = _cache(redis) + + await prefetch_identity_keys([end_user_cache_key("eu-absent")], cache) + + assert redis.round_trips == 1 + assert cache.in_memory_cache.get_cache(end_user_cache_key("eu-absent")) is None diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 78b281d5c78..b2da7f30926 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -9354,3 +9354,59 @@ async def test_admission_and_budget_reservation_read_the_key_spend_counter_with_ ], "admission and reservation share one snapshot, and read-then-write callers go to Redis once it closes" assert redis.async_batch_get_cache.await_count == 1 assert "spend:key:hashed" in redis.async_batch_get_cache.await_args.kwargs["key_list"] + + +def test_identity_prefetch_keys_match_what_auth_reads_for_the_request(): + from litellm.proxy.auth.user_api_key_auth import _identity_cache_keys + from litellm.proxy.common_utils.user_api_key_cache import ( + end_user_cache_key, + end_user_restricted_registry_cache_key, + model_access_group_registry_cache_key, + ) + from litellm.proxy.utils import hash_token + + assert _identity_cache_keys("sk-1234", end_user_id="eu-1", key_is_resolved=False) == ( + hash_token("sk-1234"), + end_user_cache_key("eu-1"), + end_user_restricted_registry_cache_key(), + model_access_group_registry_cache_key(), + ) + assert _identity_cache_keys("a" * 64, end_user_id=None, key_is_resolved=False) == ( + hash_token("a" * 64), + model_access_group_registry_cache_key(), + ) + master_key_keys = _identity_cache_keys("my-master-key", end_user_id=None, key_is_resolved=False) + assert master_key_keys == (hash_token("my-master-key"), model_access_group_registry_cache_key()) + assert "my-master-key" not in master_key_keys, "a bearer that is not an sk- key must not be sent to Redis as is" + assert _identity_cache_keys("sk-1234", end_user_id=None, key_is_resolved=True) == ( + model_access_group_registry_cache_key(), + ) + + +@pytest.mark.asyncio +async def test_virtual_key_cannot_enter_checks_as_an_identity_managed_actor(monkeypatch: pytest.MonkeyPatch) -> None: + from typing import Final + from litellm.proxy import proxy_server + from litellm.proxy.auth import user_api_key_auth as auth_module + from litellm.types.agents import AgentResponse + from litellm.types.proxy.agent_identity import AgentIdentityBinding + + target: Final = AgentResponse( + agent_id="bound", agent_name="Bound", agent_card_params={}, identity_managed=True, + identity=AgentIdentityBinding( + agent_id="bound", provider="microsoft_entra", tenant_id="tenant", client_id="client", issuer="issuer", revision="current" + ), + ) + client: Final = MagicMock() + client.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=target) + monkeypatch.setattr(proxy_server, "prisma_client", client) + checks: Final = AsyncMock() + monkeypatch.setattr(auth_module, "_run_centralized_common_checks", checks) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock(return_value=None))) + data: Final = {"model": "allowed", "messages": [{"role": "user", "content": "hello"}]} + request: Final = _alias_request("/v1/chat/completions", data) + with pytest.raises(ProxyException): + await auth_module._authorize_authenticated_request( + UserAPIKeyAuth(agent_id="bound"), request, data, "/v1/chat/completions", "sk-test" + ) + checks.assert_not_awaited() diff --git a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py index 9c07242bd23..5b7d35c3b46 100644 --- a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py @@ -7,14 +7,17 @@ gate, and the backward-compatibility guarantees that let legacy XSalsa20-Poly130 """ import base64 +import re import pytest from litellm.proxy import proxy_server from litellm.proxy.common_utils.encrypt_decrypt_utils import ( _V2_GCM_PREFIX, + decrypt_bearer_token, decrypt_if_encrypted_with, decrypt_value_helper, + encrypt_bearer_token, encrypt_value, encrypt_value_helper, ) @@ -236,3 +239,30 @@ def test_explicit_key_decrypt_supports_the_empty_master_key(): written_with_empty_key = encrypt_value(value="stored-secret", signing_key="") assert decrypt_if_encrypted_with(base64.urlsafe_b64encode(written_with_empty_key).decode(), "") == "stored-secret" + + +def test_bearer_token_opens_only_under_its_own_prefix(): + token = encrypt_bearer_token("session", prefix="kind_a_") + relabeled = "kind_b_" + token.removeprefix("kind_a_") + + assert decrypt_bearer_token(token, prefix="kind_a_") == "session" + assert decrypt_bearer_token(token, prefix="kind_b_") is None + assert decrypt_bearer_token(relabeled, prefix="kind_b_") is None + + +@pytest.mark.parametrize("use_aes", [False, True]) +def test_stored_value_is_not_a_bearer_token_even_when_reshaped(monkeypatch, use_aes: bool): + if use_aes: + _use_aes(monkeypatch) + stored = encrypt_value_helper("stored-secret") + + for candidate in (stored, "kind_a_" + stored.removeprefix(_V2_GCM_PREFIX).rstrip("=")): + assert decrypt_bearer_token(candidate, prefix="kind_a_") is None + + +@pytest.mark.parametrize("length", range(6)) +def test_bearer_token_uses_only_header_safe_characters(length: int): + token = encrypt_bearer_token("x" * length, prefix="kind_a_") + + assert re.fullmatch(r"kind_a_[A-Za-z0-9_-]+", token), token + assert decrypt_bearer_token(token, prefix="kind_a_") == "x" * length diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 7abb6e1ef92..7b160c055d2 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1638,6 +1638,45 @@ async def test_endpoint_field_is_correctly_mapped_from_call_type(): assert transaction["custom_llm_provider"] == "openai" +@pytest.mark.asyncio +async def test_endpoint_field_maps_retrieve_batch_spend_row_to_batches_endpoint(): + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + mock_prisma.get_request_status = MagicMock(return_value="success") + + payload = { + "request_id": "req-retrieve-batch", + "user": "test-user", + "call_type": "aretrieve_batch", + "startTime": "2024-01-01T12:00:00", + "api_key": "test-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "model_group": "gpt-4-group", + "prompt_tokens": 15, + "completion_tokens": 10, + "spend": 0.0175, + "metadata": '{"usage_object": {}}', + } + + writer.daily_spend_update_queue.add_update = AsyncMock() + + await writer.add_spend_log_transaction_to_daily_user_transaction( + payload=payload, + prisma_client=mock_prisma, + ) + + writer.daily_spend_update_queue.add_update.assert_called_once() + + call_args = writer.daily_spend_update_queue.add_update.call_args[1] + update_dict = call_args["update"] + assert len(update_dict) == 1 + + for key, transaction in update_dict.items(): + assert key == "test-user_2024-01-01_test-key_gpt-4_openai_/batches" + assert transaction["endpoint"] == "/batches" + + @pytest.mark.asyncio async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index aa6be328f4a..a5d2828dd9c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -20291,7 +20291,11 @@ async def test_update_key_row_with_soft_budget_updates_budget_and_key_in_transac existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) created_row = MagicMock(budget_id="budget-new") updated_row = MagicMock() - updated_row.model_dump.return_value = {"token": "hashed", "budget_id": "budget-new"} + updated_row.model_dump.return_value = { + "token": "hashed", + "budget_id": "budget-new", + "object_permission": {"mcp_servers": ["srv-1"], "mcp_tool_permissions": {"srv-1": ["read"]}}, + } tx = MagicMock() tx.litellm_budgettable.create = AsyncMock(return_value=created_row) tx.litellm_verificationtoken.update = AsyncMock(return_value=updated_row) @@ -20312,10 +20316,15 @@ async def test_update_key_row_with_soft_budget_updates_budget_and_key_in_transac ) assert set(result) == {"token", "data"} - assert result["data"] == {"token": "hashed", "budget_id": "budget-new"} + assert result["data"] == { + "token": "hashed", + "budget_id": "budget-new", + "object_permission": {"mcp_servers": ["srv-1"], "mcp_tool_permissions": {"srv-1": ["read"]}}, + } tx.litellm_verificationtoken.update.assert_awaited_once() update_call = tx.litellm_verificationtoken.update.await_args assert update_call.kwargs["where"] == {"token": result["token"]} + assert update_call.kwargs["include"] == {"object_permission": True} assert update_call.kwargs["data"]["budget_id"] == "budget-new" assert "soft_budget" not in update_call.kwargs["data"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 160cf8be4e0..f5fc5ae24d4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -7679,7 +7679,7 @@ class TestConnectedAppViewAnnotation: flags = {server.server_id: server.connected_app_reachable for server in result} assert flags == {"server-1": True, "server-2": False} - reload_mock.assert_awaited_once_with("test_user_id") + reload_mock.assert_awaited_once_with("test_user_id", requires_fresh_policy=False) mock_manager.get_allowed_mcp_servers.assert_awaited_once_with(admitted_auth) @pytest.mark.asyncio @@ -9653,7 +9653,7 @@ class TestMCPServerResolutionCharacterization: server_id: str, ) -> tuple[MagicMock, MCPServerManager, UserAPIKeyAuth]: team_id: Final = UI_SESSION_TOKEN_TEAM_ID if grant_route == "direct user object_permission" else "lit3974_team" - user_id: Final = "lit3974_direct_user" + user_id: Final = f"{server_id}:{grant_route}:user" key_permission: Final = LiteLLM_ObjectPermissionTable( object_permission_id=f"lit3974_{grant_route}_key_permission", mcp_servers=None, diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index b066b3b80e6..a53894fcd1b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4355,7 +4355,7 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many) prisma_client.db.litellm_teamtable.count = AsyncMock(side_effect=count) prisma_client.db.litellm_verificationtoken.group_by = AsyncMock(return_value=[]) - prisma_client.db.litellm_usertable.find_unique = AsyncMock( + prisma_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=LiteLLM_UserTable( user_id="org_admin_user", teams=["team_in_org_A", "team_in_org_B"], @@ -4394,11 +4394,11 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs( assert await list_teams(None) == own_view assert await list_teams("org_admin_user", search="team_in_org_B") == ["team_in_org_B"] assert await list_teams("other_user") == ["other_team_in_org_A"] - prisma_client.db.litellm_usertable.find_unique.assert_awaited_with( + prisma_client.writer_db.litellm_usertable.find_unique.assert_awaited_with( where={"user_id": "org_admin_user"}, include={"organization_memberships": True} ) - prisma_client.db.litellm_usertable.find_unique.side_effect = RuntimeError("db down") + prisma_client.writer_db.litellm_usertable.find_unique.side_effect = RuntimeError("db down") with pytest.raises(ValueError, match="db down"): await list_teams("org_admin_user") @@ -15813,7 +15813,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) mock_db_client.db.query_raw = AsyncMock(return_value=[]) - mock_db_client.db.litellm_usertable.find_unique = AsyncMock( + mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("alice", ["team-alpha"]) ) @@ -15835,7 +15835,7 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) mock_db_client.db.query_raw = AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)]) - mock_db_client.db.litellm_usertable.find_unique = AsyncMock( + mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("bob", ["team-alpha"]) ) @@ -15856,7 +15856,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_USER) mock_db_client.db.query_raw = AsyncMock(return_value=[]) - mock_db_client.db.litellm_usertable.find_unique = AsyncMock( + mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("bob", ["team-alpha"]) ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py index 74f3a2248c7..bc0bbd4dd38 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py +++ b/tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py @@ -1,19 +1,27 @@ import asyncio +import re import time -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta +from pathlib import Path from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock +import litellm_proxy_extras +import psycopg import pytest from prisma.errors import PrismaError +from psycopg.rows import dict_row +from psycopg.types.json import Jsonb +from pytest_postgresql import factories from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( SPEND_LOG_KEY_METADATA_CACHE_TTL, SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL, SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS, + SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE, ) from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_details, @@ -588,12 +596,314 @@ async def test_recover_key_metadata_from_spend_logs_bounds_the_scan_with_a_state await recover_key_metadata_from_spend_logs(mock_prisma, {digest}, window, cache=InMemoryCache()) - assert calls == [f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}", "scan"] + assert calls == [ + f"SET LOCAL statement_timeout = {SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS}", + "SET LOCAL enable_bitmapscan = off", + "scan", + ] assert mock_prisma.db.tx.call_args.kwargs["timeout"] == timedelta( milliseconds=2 * SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS ) +_spend_logs_postgresql_proc: Final = factories.postgresql_proc() +_spend_logs_postgresql: Final = factories.postgresql("_spend_logs_postgresql_proc") + +_SPEND_LOGS_DDL: Final = """ + CREATE TABLE "LiteLLM_SpendLogs" ( + request_id TEXT PRIMARY KEY, + api_key TEXT NOT NULL DEFAULT '', + "startTime" TIMESTAMP(3) NOT NULL, + "user" TEXT DEFAULT '', + team_id TEXT, + metadata JSONB DEFAULT '{}' + ) +""" + +_API_KEY_START_TIME_INDEX_MIGRATION: Final = ( + Path(litellm_proxy_extras.__file__).parent + / "migrations" + / "20260823000000_add_spend_logs_api_key_starttime_index" + / "migration.sql" +) + +_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL: Final = """ + SELECT COALESCE(seq_tup_read, 0) + COALESCE(idx_tup_fetch, 0) AS rows_read + FROM pg_stat_xact_user_tables + WHERE relname = 'LiteLLM_SpendLogs' +""" + + +def _create_spend_logs_table(conn: psycopg.Connection) -> None: + conn.execute(_SPEND_LOGS_DDL) # pyright: ignore[reportArgumentType] # DDL literal + conn.execute(_API_KEY_START_TIME_INDEX_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # migration file + + +def _psycopg_prisma(conn: psycopg.Connection) -> MagicMock: + async def query_raw(sql: str, *params: object) -> list[dict[str, object]]: + with conn.cursor(row_factory=dict_row) as cur: + cur.execute( + re.sub(r"\$(\d+)", r"%(p\1)s", sql), # pyright: ignore[reportArgumentType] # proxy SQL is not a literal + {f"p{i}": v for i, v in enumerate(params, start=1)}, + ) + return cur.fetchall() + + async def execute_raw(sql: str) -> int: + conn.execute(sql) # pyright: ignore[reportArgumentType] # proxy SQL is not a literal + return 0 + + mock_prisma: Final = MagicMock() + transaction: Final = MagicMock() + transaction.query_raw = AsyncMock(side_effect=query_raw) + transaction.execute_raw = AsyncMock(side_effect=execute_raw) + mock_prisma.db.tx.return_value.__aenter__.return_value = transaction + return mock_prisma + + +def _commit_and_vacuum(conn: psycopg.Connection) -> None: + conn.commit() + conn.set_autocommit(True) + conn.execute('VACUUM (ANALYZE) "LiteLLM_SpendLogs"') + conn.set_autocommit(False) + + +def _insert_nameless_spend_logs(conn: psycopg.Connection, digest: str, rows: int) -> None: + conn.execute( + """ + INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime") + SELECT %(digest)s || '-' || g, %(digest)s, %(start)s + g * interval '1 minute' + FROM generate_series(1, %(rows)s) g + """, + {"digest": digest, "start": datetime(2026, 9, 7), "rows": rows}, + ) + + +def _named_spend_log( + digest: str, logged_at: datetime, alias: str | None, user: str | None, team: str | None = None +) -> tuple[str, str, datetime, str, str | None, Jsonb]: + return ( + f"{digest}-{logged_at.isoformat()}", + digest, + logged_at, + user or "", + team, + Jsonb({"user_api_key_alias": alias} if alias else {}), + ) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_names_a_key_by_its_oldest_and_newest_named_rows_in_the_window( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + unnamed_edges, owner_logged_late, reowned, outside_window, never_named = ( + hash_token(f"cli-session-{name}") for name in ("edges", "late", "reowned", "window", "never") + ) + with conn.cursor() as cur: + cur.executemany( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", team_id, metadata)' + " VALUES (%s, %s, %s, %s, %s, %s)", + ( + _named_spend_log(unnamed_edges, datetime(2026, 9, 7, 1), None, None), + _named_spend_log(unnamed_edges, datetime(2026, 9, 8), "cli-a", "alice", "team-a"), + _named_spend_log(unnamed_edges, datetime(2026, 9, 9), "cli-a", "alice", "team-a"), + _named_spend_log(unnamed_edges, datetime(2026, 9, 9, 23), None, None), + _named_spend_log(owner_logged_late, datetime(2026, 9, 7, 1), "cli-b", None), + _named_spend_log(owner_logged_late, datetime(2026, 9, 9), "cli-b", "bob"), + _named_spend_log(reowned, datetime(2026, 9, 7, 1), "cli-c", "carol"), + _named_spend_log(reowned, datetime(2026, 9, 9), "cli-c", "dave"), + _named_spend_log(outside_window, datetime(2026, 9, 6), "stale-alias", "erin"), + _named_spend_log(outside_window, datetime(2026, 9, 8), "cli-d", "erin"), + _named_spend_log(outside_window, datetime(2026, 9, 10), "later-alias", "erin"), + _named_spend_log(never_named, datetime(2026, 9, 8), None, None), + ), + ) + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), + {unnamed_edges, owner_logged_late, reowned, outside_window, never_named}, + (datetime(2026, 9, 7), datetime(2026, 9, 10)), + cache=InMemoryCache(), + ) + + assert dict(result) == { + unnamed_edges: {"key_alias": "cli-a", "team_id": "team-a", "user_id": "alice"}, + owner_logged_late: {"key_alias": "cli-b", "team_id": None, "user_id": "bob"}, + reowned: {"key_alias": "cli-c", "team_id": None, "user_id": None}, + outside_window: {"key_alias": "cli-d", "team_id": None, "user_id": "erin"}, + } + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_reads_two_rows_per_key_however_many_the_key_logged( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + owners: Final[Mapping[str, str]] = {hash_token(f"cli-session-busy-{i}"): f"user-{i}" for i in range(5)} + for digest, owner in owners.items(): + conn.execute( + """ + INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", metadata) + SELECT %(digest)s || '-' || g, %(digest)s, %(start)s + g * interval '1 minute', %(owner)s, + jsonb_build_object('user_api_key_alias', 'cli-session-' || %(owner)s) + FROM generate_series(1, 2000) g + """, + {"digest": digest, "owner": owner, "start": datetime(2026, 9, 7)}, + ) + conn.execute('ANALYZE "LiteLLM_SpendLogs"') + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), frozenset(owners), (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=InMemoryCache() + ) + + assert {digest: meta.get("user_id") for digest, meta in result.items()} == owners + rows_read: Final = conn.execute(_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL).fetchone() # pyright: ignore[reportArgumentType] # SQL literal + assert rows_read is not None and rows_read[0] <= 2 * len(owners) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_walks_a_bounded_number_of_nameless_rows_per_key( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + named_late: Final[Mapping[str, str]] = {hash_token(f"cli-session-late-{i}"): f"user-{i}" for i in range(3)} + never_named: Final = frozenset(hash_token(f"cli-session-never-{i}") for i in range(3)) + for digest in (*named_late, *never_named): + _insert_nameless_spend_logs(conn, digest, 3 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE) + for digest, owner in named_late.items(): + conn.execute( + """ + INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", metadata) + VALUES (%(digest)s || '-newest', %(digest)s, %(logged_at)s, %(owner)s, + jsonb_build_object('user_api_key_alias', 'cli-session-' || %(owner)s)) + """, + {"digest": digest, "owner": owner, "logged_at": datetime(2026, 9, 9)}, + ) + conn.execute('ANALYZE "LiteLLM_SpendLogs"') + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), + frozenset(named_late) | never_named, + (datetime(2026, 9, 7), datetime(2026, 9, 10)), + cache=InMemoryCache(), + ) + + assert {digest: meta.get("user_id") for digest, meta in result.items()} == named_late + rows_read: Final = conn.execute(_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL).fetchone() # pyright: ignore[reportArgumentType] # SQL literal + assert rows_read is not None + assert rows_read[0] <= 3 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE * (len(named_late) + len(never_named)) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_reads_a_short_nameless_key_once( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + rows_per_key: Final = SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE // 2 + never_named: Final = frozenset(hash_token(f"cli-session-short-{i}") for i in range(20)) + for digest in never_named: + _insert_nameless_spend_logs(conn, digest, rows_per_key) + _commit_and_vacuum(conn) + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), never_named, (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=InMemoryCache() + ) + + assert dict(result) == {} + rows_read: Final = conn.execute(_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL).fetchone() # pyright: ignore[reportArgumentType] # SQL literal + assert rows_read is not None + assert rows_read[0] <= rows_per_key * len(never_named) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_bounds_a_busy_nameless_key_among_short_keys_before_any_vacuum( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + busy: Final = frozenset(hash_token(f"cli-session-busy-nameless-{i}") for i in range(3)) + for short_key in range(200): + _insert_nameless_spend_logs( + conn, hash_token(f"cli-session-short-{short_key}"), SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE // 5 + ) + for digest in busy: + _insert_nameless_spend_logs(conn, digest, 30 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE) + conn.execute('ANALYZE "LiteLLM_SpendLogs"') + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), busy, (datetime(2026, 9, 7), datetime(2026, 9, 10)), cache=InMemoryCache() + ) + + assert dict(result) == {} + rows_read: Final = conn.execute(_SPEND_LOG_ROWS_READ_IN_THIS_TRANSACTION_SQL).fetchone() # pyright: ignore[reportArgumentType] # SQL literal + assert rows_read is not None + assert rows_read[0] <= 3 * SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE * len(busy) + + +@pytest.mark.asyncio +async def test_recover_key_metadata_from_spend_logs_finds_a_name_logged_where_the_oldest_probe_stopped( + _spend_logs_postgresql: psycopg.Connection, +): + conn: Final = _spend_logs_postgresql + _create_spend_logs_table(conn) + start: Final = datetime(2026, 9, 7) + past_the_stop, tied_with_the_stop = (hash_token(f"cli-session-{name}") for name in ("past", "tied")) + same_millisecond: Final = tuple( + start + timedelta(minutes=SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE, microseconds=n) for n in (100, 200, 300) + ) + with conn.cursor() as cur: + cur.executemany( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, api_key, "startTime", "user", team_id, metadata)' + " VALUES (%s, %s, %s, %s, %s, %s)", + ( + *( + _named_spend_log(past_the_stop, start + timedelta(minutes=minute), None, None) + for minute in range(1, SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 20) + ), + _named_spend_log( + past_the_stop, + start + timedelta(minutes=SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 20), + "cli-p", + "pat", + ), + *( + _named_spend_log(past_the_stop, start + timedelta(minutes=minute), None, None) + for minute in range( + SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 21, SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE + 51 + ) + ), + *( + _named_spend_log(tied_with_the_stop, start + timedelta(minutes=minute), None, None) + for minute in range(1, SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE) + ), + _named_spend_log(tied_with_the_stop, same_millisecond[0], None, None), + _named_spend_log(tied_with_the_stop, same_millisecond[1], None, None), + _named_spend_log(tied_with_the_stop, same_millisecond[2], "cli-t", "tess"), + ), + ) + conn.commit() + + result = await recover_key_metadata_from_spend_logs( + _psycopg_prisma(conn), + {past_the_stop, tied_with_the_stop}, + (start, datetime(2026, 9, 10)), + cache=InMemoryCache(), + ) + + assert dict(result) == { + past_the_stop: {"key_alias": "cli-p", "team_id": None, "user_id": "pat"}, + tied_with_the_stop: {"key_alias": "cli-t", "team_id": None, "user_id": "tess"}, + } + + @pytest.mark.asyncio async def test_recover_cli_session_key_metadata_names_the_owner_only_when_the_suffix_is_a_real_user(): mock_prisma = MagicMock() diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 347adc421a2..3b265653b12 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -57,7 +57,7 @@ def _filter_logs_by_date_range(logs, where): _SEARCH_CLAUSE_RE = re.compile( r'\(request_id = \$(\d+) OR \("startTime" >= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) ' r'AND "startTime" <= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) ' - r'AND \(api_key = \$\1 OR team_id = \$\1 OR "user" = \$\1 OR end_user = \$\1 ' + r'AND \(litellm_call_id = \$\1 OR api_key = \$\1 OR team_id = \$\1 OR "user" = \$\1 OR end_user = \$\1 ' r"OR session_id = \$\1 OR model_id = \$\1\)\)\)" ) @@ -68,7 +68,7 @@ def _matches_spend_log_search(log, search): return True if not _filter_logs_by_date_range([log], {"startTime": {"gte": search["gte"], "lte": search["lte"]}}): return False - columns = ("api_key", "team_id", "user", "end_user", "session_id", "model_id") + columns = ("litellm_call_id", "api_key", "team_id", "user", "end_user", "session_id", "model_id") return any(log.get(col) == search["value"] for col in columns) @@ -2986,7 +2986,7 @@ def test_build_spend_log_search_condition_windows_every_branch_except_request_id assert condition.sql == ( "(request_id = $3 OR (\"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC') " "AND \"startTime\" <= ($5::timestamptz AT TIME ZONE 'UTC') " - 'AND (api_key = $3 OR team_id = $3 OR "user" = $3 OR end_user = $3 OR session_id = $3 OR model_id = $3)))' + 'AND (litellm_call_id = $3 OR api_key = $3 OR team_id = $3 OR "user" = $3 OR end_user = $3 OR session_id = $3 OR model_id = $3)))' ) assert condition.params == ("key-hash-7", start, end) @@ -3012,6 +3012,8 @@ def _search_fixture_logs(today): {**base, "request_id": "req-user", "user": "user-7", "startTime": recent}, {**base, "request_id": "req-end-user", "end_user": "cust-7", "startTime": recent}, {**base, "request_id": "req-model", "model_id": "mdl-7", "startTime": recent}, + {**base, "request_id": "chatcmpl-x", "litellm_call_id": "call-recent", "startTime": recent}, + {**base, "request_id": "chatcmpl-old", "litellm_call_id": "call-old", "startTime": old}, ] @@ -3046,6 +3048,8 @@ def _five_day_window(today): ("user-7", {"req-user"}), ("cust-7", {"req-end-user"}), ("mdl-7", {"req-model"}), + ("call-recent", {"chatcmpl-x"}), + ("call-old", set()), ("no-such-id", set()), ], ) @@ -3762,7 +3766,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"actor_agent_id": null, "target_agent_id": null, "billing_agent_id": null, "agent_execution_mode": null, "verified_human_user_id": null, "applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 00223f192ec..f3991c0e494 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -612,7 +612,7 @@ def test_sanitize_request_body_for_spend_logs_payload_mixed_types(): request_body = { "text": long_string, "number": 42, - "nested": {"list": ["short", long_string], "dict": {"key": long_string}}, + "nested": {"list": ["short", long_string], "dict": {"value": long_string}}, } sanitized = _sanitize_request_body_for_spend_logs_payload(request_body) @@ -631,7 +631,7 @@ def test_sanitize_request_body_for_spend_logs_payload_mixed_types(): assert sanitized["number"] == 42 assert sanitized["nested"]["list"][0] == "short" assert len(sanitized["nested"]["list"][1]) == expected_length - assert len(sanitized["nested"]["dict"]["key"]) == expected_length + assert len(sanitized["nested"]["dict"]["value"]) == expected_length def test_sanitize_request_body_for_spend_logs_payload_uses_runtime_env_override( @@ -1207,7 +1207,7 @@ def test_get_logging_payload_placeholders_the_metadata_copied_into_the_stored_re stored_request_body: Final = json.loads(payload["proxy_server_request"]) assert stored_request_body["metadata"]["model_group"] == expected_stored_model_group assert stored_request_body["metadata"]["error_information"]["error_message"] == expected_stored_error_message - assert stored_request_body["metadata"]["user_api_key"] == "sk-test" + assert stored_request_body["metadata"]["user_api_key"] == REDACTED_BY_LITELM_STRING assert ("medical records" in payload["proxy_server_request"]) == bool(deployment_info) @@ -2691,6 +2691,104 @@ def test_sanitize_request_body_strips_secret_fields(): assert sanitized["messages"] == [{"role": "user", "content": "hi"}] +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") +def test_proxy_server_request_payload_strips_nested_aws_credentials(mock_should_store: MagicMock) -> None: + mock_should_store.return_value = True + credentials: Final = { + "aws_access_key_id": "AKIA-canary", + "aws_secret_access_key": "secret-canary", + "aws_session_token": "token-canary", + "aws_web_identity_token": "wit-canary", + } + tool_parameters: Final = {"type": "object", "properties": {"aws_secret_access_key": {"type": "string"}}} + litellm_params: Final = { + "proxy_server_request": { + "body": { + "model": "bedrock-claude", + "messages": [{"role": "user", "content": "hello"}], + "fallbacks": [{"model": "bedrock-b", "aws_region_name": "us-west-2", **credentials}], + "extra_body": {"aws_role_name": "arn:aws:iam::123456789012:role/r", **credentials}, + "tools": [{"type": "function", "function": {"name": "f", "parameters": tool_parameters}}], + **credentials, + } + } + } + + parsed: Final = json.loads( + _get_proxy_server_request_for_spend_logs_payload(metadata={}, litellm_params=litellm_params, kwargs={}) + ) + + assert "canary" not in json.dumps(parsed) + masked: Final = dict.fromkeys(credentials, REDACTED_BY_LITELM_STRING) + assert parsed["fallbacks"] == [{"model": "bedrock-b", "aws_region_name": "us-west-2", **masked}] + assert parsed["extra_body"] == {"aws_role_name": "arn:aws:iam::123456789012:role/r", **masked} + assert {name: parsed[name] for name in credentials} == masked + assert parsed["tools"][0]["function"]["parameters"] == tool_parameters + assert parsed["messages"] == [{"role": "user", "content": "hello"}] + + +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") +def test_proxy_server_request_payload_redacts_provider_credentials(mock_should_store: MagicMock) -> None: + mock_should_store.return_value = True + credentials: Final = { + "azure_password": "canary-azure-password", + "client_secret": "canary-client-secret", + "azure_ad_token": "canary-azure-ad-token", + "vertex_credentials": "canary-vertex-credentials", + "s3_secret_access_key": "canary-s3-secret", + "token": "canary-watsonx-token", + "apikey": "canary-watsonx-apikey", + "zen_api_key": "canary-zen-api-key", + "gemini_api_key": "canary-gemini-api-key", + "gigachat_access_token": "canary-gigachat-token", + "oci_key": "canary-oci-key", + } + metadata: Final = {"user_api_key": "custom-auth-raw-key", "requester_ip_address": "10.0.0.1"} + tool_parameters: Final = {"type": "object", "properties": {"client_secret": {"type": "string"}}} + litellm_params: Final = { + "proxy_server_request": { + "body": { + "model": "azure-gpt", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + "prompt_cache_key": "user-123-cache", + "vertex_credentials": {"private_key": "canary-private-key", "client_email": "sa@example.com"}, + "extra_headers": {"Authorization": "Bearer canary-extra-header"}, + "tools": [ + {"type": "function", "function": {"name": "f", "parameters": tool_parameters}}, + {"type": "mcp", "server_url": "https://mcp.example.com", "headers": {"Authorization": "canary-mcp"}}, + ], + "fallbacks": [{"model": "azure-b", **credentials}], + "metadata": metadata, + **credentials, + } + } + } + + parsed: Final = json.loads( + _get_proxy_server_request_for_spend_logs_payload(metadata={}, litellm_params=litellm_params, kwargs={}) + ) + + assert "canary" not in json.dumps(parsed) + assert {name: parsed[name] for name in credentials} == dict.fromkeys(credentials, REDACTED_BY_LITELM_STRING) + assert parsed["vertex_credentials"] == REDACTED_BY_LITELM_STRING + assert parsed["extra_headers"] == {"Authorization": REDACTED_BY_LITELM_STRING} + assert parsed["tools"][0]["function"]["parameters"] == tool_parameters + assert parsed["tools"][1]["server_url"] == "https://mcp.example.com" + assert parsed["metadata"] == {"user_api_key": REDACTED_BY_LITELM_STRING, "requester_ip_address": "10.0.0.1"} + assert parsed["max_tokens"] == 10 + assert parsed["prompt_cache_key"] == REDACTED_BY_LITELM_STRING + assert parsed["messages"] == [{"role": "user", "content": "hello"}] + + +def test_sanitize_response_redacts_credential_named_fields() -> None: + response: Final = {"access_token": "canary-oauth-token", "usage": {"prompt_tokens": 1}} + + assert _sanitize_request_body_for_spend_logs_payload({"response": response}) == { + "response": {"access_token": REDACTED_BY_LITELM_STRING, "usage": {"prompt_tokens": 1}} + } + + @patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_proxy_server_request_payload_excludes_secret_fields(mock_should_store): """ diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index b43a75d3323..adc3bc04bdf 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -20,6 +20,14 @@ from litellm.proxy._types import ( ) SERVER_ONLY_MARKERS = ( + "requires_fresh_policy", + "mcp_explicit_grants_only", + "managed_agent_context", + "managed_agent_policy", + "invoked_agent_id", + "invoked_agent_policy", + "agent_invocation_cost", + "billing_agent_policy", "mcp_admitted_user_subject", "mcp_source_team_rpm_limits", "mcp_session_resource_server_id", diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index 672dd1eb674..05c4f9d8a67 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -634,3 +634,28 @@ async def test_query_first_with_cached_plan_fallback_reports_the_reader_generati "reader_served_the_query": 2, "writer_served_the_query": 0, } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rotated", (False, True)) +async def test_authoritative_combined_key_view_uses_writer_through_rotation( + prisma_client: PrismaClient, rotated: bool +) -> None: + writer: Final = MagicMock() + reader: Final = MagicMock() + active: Final = { + "token": "current-token", "team_id": "current-team", "team_models": None, + "team_blocked": None, "team_members_with_roles": None, "user_id": None, "expires": None, + } + writer.query_first = AsyncMock(side_effect=[None, active] if rotated else [active]) + reader.query_first = AsyncMock(return_value={**active, "team_id": "stale-team"}) + writer.litellm_deprecatedverificationtoken.find_first = AsyncMock(return_value=SimpleNamespace( + active_token_id="current-token", revoke_at=datetime.now(timezone.utc) + timedelta(hours=1) + )) + prisma_client.db = RoutingPrismaWrapper(writer=writer, reader=reader) + response: Final = await prisma_client.get_data(token="original-token", table_name="combined_view", use_writer=True) + assert isinstance(response, LiteLLM_VerificationTokenView) + assert response.team_id == "current-team" + assert response.token == "current-token" + reader.query_first.assert_not_awaited() + assert writer.query_first.await_count == (2 if rotated else 1) diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py index dd241397e87..6e69444a1b5 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py @@ -155,6 +155,7 @@ async def test_update_data_token_hashes_and_updates( "token": hashlib.sha256(token.encode()).hexdigest(), "spend": 1.0, "user_id": "u1", + "object_permission": {"mcp_servers": ["srv-1"]}, }, ) prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=response) @@ -167,15 +168,22 @@ async def test_update_data_token_hashes_and_updates( actual = { "result": result, "where": update_kwargs["where"], + "include": update_kwargs["include"], "data_token": update_kwargs["data"]["token"], "data_spend": update_kwargs["data"]["spend"], } assert actual == { "result": { "token": hashed, - "data": {"token": hashed, "spend": 1.0, "user_id": "u1"}, + "data": { + "token": hashed, + "spend": 1.0, + "user_id": "u1", + "object_permission": {"mcp_servers": ["srv-1"]}, + }, }, "where": {"token": hashed}, + "include": {"object_permission": True}, "data_token": hashed, "data_spend": 1.0, } diff --git a/tests/unit/caching/test_dual_cache.py b/tests/unit/caching/test_dual_cache.py index eb2f19ac377..521fda31b58 100644 --- a/tests/unit/caching/test_dual_cache.py +++ b/tests/unit/caching/test_dual_cache.py @@ -2,6 +2,7 @@ import asyncio import logging import time import uuid +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -136,6 +137,35 @@ def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): assert "shared_a" not in dual_cache.last_redis_batch_access_time +def test_reserve_redis_batch_reads_reserves_memory_misses_and_can_be_rolled_back(): + mock_redis: Final = MagicMock(spec=RedisCache) + dual_cache: Final = DualCache( + in_memory_cache=InMemoryCache(), + redis_cache=mock_redis, + default_redis_batch_cache_expiry=10, + ) + dual_cache.in_memory_cache.set_cache("memory_key", "memory_value") + + reserved, previous_access_times = dual_cache.reserve_redis_batch_reads(["memory_key", "missing_key"]) + + assert reserved == ["missing_key"] + assert previous_access_times == {"missing_key": None} + assert dual_cache.reserve_redis_batch_reads(["memory_key", "missing_key"]) == ([], {}) + + dual_cache._rollback_redis_batch_key_reservations(previous_access_times) + + assert dual_cache.reserve_redis_batch_reads(["memory_key", "missing_key"]) == ( + ["missing_key"], + {"missing_key": None}, + ) + + +def test_reserve_redis_batch_reads_returns_empty_without_redis(): + dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=None) + + assert dual_cache.reserve_redis_batch_reads(["missing_key"]) == ([], {}) + + def test_dual_cache_batch_get_cache_returns_memory_only_when_redis_read_is_throttled(): mock_redis = _redis_mock_for_sync_batch({"throttled_key": "redis_value"}) dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) diff --git a/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index 9433aeac524..93206efc80f 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -58,6 +58,10 @@ class FakePipeline: self.commands.append(("SET", name, value, None if ex is None else int(ex.total_seconds()))) return self + def delete(self, *names: str) -> FakePipeline: + self.commands.append(("DEL", *names)) + return self + async def execute(self, raise_on_error: bool = True) -> list[Any]: assert raise_on_error is False self.executed = True @@ -101,6 +105,14 @@ class FakeRedisCache(RedisCache): self.store[key] = float(self.store.get(key, 0.0)) + value return self.store[key] + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # fake, no server + self.alone.append(("SET", key, value)) + self.store[key] = value + + async def async_delete_cache(self, key: str) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # records the direct delete + self.alone.append(("DEL", key)) + self.store.pop(key, None) + async def async_set_cache_pipeline_with_ttls(self, cache_list: Sequence[tuple[str, object, float | None]]) -> None: self.alone.append(("SET_PIPELINE", tuple(cache_list))) for key, value, _ttl in cache_list: @@ -124,6 +136,8 @@ def replies(command: tuple[Any, ...]) -> Any: return 1 case "SET": return True + case "DEL": + return 1 raise AssertionError(command) @@ -297,3 +311,51 @@ def test_request_scope_hands_out_one_batch_per_backend_and_nests() -> None: assert active_request_redis_batch(cache_a) is first assert len(batches.batches) == 2 assert active_request_redis_batch(cache_a) is None + + +@pytest.mark.asyncio +async def test_a_key_an_mget_read_as_absent_stays_known_missing_until_something_sets_it() -> None: + cache, client = make() + batch = RedisBatch(cache) + values = await batch.mget(["a-hit", "b-miss"]) + assert values == {"a-hit": {"k": "a-hit"}, "b-miss": None} + assert batch.read_as_missing("b-miss") is True + assert batch.read_as_missing("a-hit") is False + assert batch.read_as_missing("never-read") is False + batch.set("b-miss", "now-present") + assert batch.read_as_missing("b-miss") is False + + +@pytest.mark.asyncio +async def test_a_delete_rides_the_pipeline_under_the_namespace_and_reads_as_missing_afterwards() -> None: + cache, client = make(namespace="ns") + batch = RedisBatch(cache) + gone = batch.delete("team_alias:x") + got = batch.mget(["a-hit"]) + assert await gone is None + assert await got == {"a-hit": {"k": "ns:a-hit"}} + assert len(client.pipelines) == 1 + assert client.pipelines[0].commands[0] == ("DEL", "ns:team_alias:x") + assert batch.read_as_missing("team_alias:x") is True + assert cache.alone == [] + + +@pytest.mark.asyncio +async def test_a_delete_on_a_cluster_cache_runs_as_its_own_del() -> None: + client = FakeClient(replies) + cache = FakeClusterCache(client) + cache.store["team_alias:x"] = "stale" + batch = RedisBatch(cache) + assert await batch.delete("team_alias:x") is None + assert cache.alone == [("DEL", "team_alias:x")] + assert "team_alias:x" not in cache.store + assert client.pipelines == [] + + +@pytest.mark.asyncio +async def test_a_failed_mget_marks_nothing_as_missing() -> None: + cache, client = make(fail=ConnectionError("down")) + batch = RedisBatch(cache) + with pytest.raises(ConnectionError): + await batch.mget(["b-miss"]) + assert batch.read_as_missing("b-miss") is False diff --git a/tests/unit/caching/test_request_redis_batch_pre_call.py b/tests/unit/caching/test_request_redis_batch_pre_call.py index c0834974f26..d4388110131 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -6,16 +6,19 @@ from __future__ import annotations import asyncio import hashlib import json +from itertools import chain from typing import Any, Final -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest from litellm import Router +import litellm.caching.dual_cache as dual_cache_module from litellm.caching.dual_cache import DualCache from litellm.caching.redis_batch import active_request_redis_batches, request_redis_batch_scope -from litellm.proxy._types import LiteLLM_UserTable -from litellm.proxy.auth.auth_object_prefetch import _CacheEntry, _write_back +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable +from litellm.proxy.auth.auth_checks import _cache_team_object +from litellm.proxy.auth.auth_object_prefetch import _CacheEntry, _write_back, prefetch_identity_keys from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.parallel_request_limiter_v3 import ( CHECK_AND_INCREMENT_BY_N_SCRIPT, @@ -27,7 +30,7 @@ from litellm.proxy.utils import InternalUsageCache from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.routing_read_batch import RoutingPrefetch -from .test_redis_batch import FakeClient, FakeRedisCache +from .test_redis_batch import FakeClient, FakeRedisCache, replies _MODEL_GROUP = "claude" _FAR_FUTURE = 4_102_444_800.0 # 2100-01-01, a cooldown stamped then is still active @@ -352,7 +355,12 @@ async def test_a_prefetch_that_does_not_cover_the_routing_keys_is_ignored_and_ro router.arm_routing_read_prefetch(_MODEL_GROUP, {}) armed = request.prefetched["routing_read"] assert isinstance(armed, RoutingPrefetch) - request.prefetched["routing_read"] = RoutingPrefetch(keys=frozenset({"other"}), result=armed.result) + request.prefetched["routing_read"] = RoutingPrefetch( + keys=frozenset({"other"}), + fetched=armed.fetched, + result=armed.result, + reservations=armed.reservations, + ) deployment = await router.async_get_available_deployment( model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} ) @@ -362,6 +370,32 @@ async def test_a_prefetch_that_does_not_cover_the_routing_keys_is_ignored_and_ro assert len(redis_cache.alone) == 1 # the shared cooldown+usage read, one round trip as in P1 +@pytest.mark.asyncio +async def test_a_prefetch_with_incomplete_usage_keys_releases_cooldown_reservations(): + client: Final = FakeClient(_lua_ok_replies) + redis_cache: Final = FakeRedisCache(client) + router: Final = _router(redis_cache) + cooldown_keys: Final = frozenset( + { + CooldownCache.get_cooldown_cache_key("dep-a"), + CooldownCache.get_cooldown_cache_key("dep-b"), + } + ) + + with request_redis_batch_scope(): + RoutingPrefetch.arm(router, router.lowesttpm_logger_v2, router.model_list[:1]) + deployment: Final = await router.async_get_available_deployment( + model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} + ) + + fallback_cooldown_mgets: Final = tuple( + keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_keys.issubset(keys) + ) + + assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} + assert len(fallback_cooldown_mgets) == 1 + + @pytest.mark.asyncio async def test_a_failed_prefetch_falls_back_to_the_shared_read(): client = FakeClient(_lua_ok_replies, fail=ConnectionError("redis down")) @@ -376,6 +410,217 @@ async def test_a_failed_prefetch_falls_back_to_the_shared_read(): assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} assert len(redis_cache.alone) == 1 + cooldown_keys: Final = frozenset( + { + CooldownCache.get_cooldown_cache_key("dep-a"), + CooldownCache.get_cooldown_cache_key("dep-b"), + } + ) + fallback_cooldown_mgets: Final = tuple( + keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_keys.issubset(keys) + ) + assert len(fallback_cooldown_mgets) == 1 + + +@pytest.mark.asyncio +async def test_an_abandoned_prefetch_still_backfills_the_cooldown_it_read(monkeypatch): + clock: Final = 1_000_000.0 + monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock) + cooldown_key: Final = CooldownCache.get_cooldown_cache_key("dep-a") + active_cooldown: Final = { + "exception_received": "429", + "status_code": "429", + "timestamp": _FAR_FUTURE, + "cooldown_time": 60, + } + + def replies(command: tuple[Any, ...]) -> Any: + if command[0] == "MGET": + return [json.dumps(active_cooldown) if key == cooldown_key else None for key in command[1:]] + return _lua_ok_replies(command) + + client: Final = FakeClient(replies) + redis_cache: Final = FakeRedisCache(client) + redis_cache.store[cooldown_key] = active_cooldown + router: Final = _router(redis_cache, routing_strategy="simple-shuffle") + limiter: Final = _limiter(redis_cache) + + with request_redis_batch_scope(): + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + await limiter.atomic_check_and_increment_by_n( + descriptors=[_descriptor("api_key", "k1", 10)], + increments=[{"requests": 1}], + ) + + pipeline_count: Final = len(client.pipelines) + first_cooldown_mgets: Final = tuple( + command + for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[:pipeline_count]) + if command[0] == "MGET" and cooldown_key in command[1:] + ) + with request_redis_batch_scope(): + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + deployment: Final = await router.async_get_available_deployment( + model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} + ) + cooldowns: Final = await router.cooldown_cache.async_get_active_cooldowns(["dep-a"], parent_otel_span=None) + + second_cooldown_mgets: Final = tuple( + command + for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[pipeline_count:]) + if command[0] == "MGET" and cooldown_key in command[1:] + ) + + assert deployment["model_info"]["id"] == "dep-b" + assert [model_id for model_id, _ in cooldowns] == ["dep-a"] + assert len(first_cooldown_mgets) == 1 + assert second_cooldown_mgets == () + assert redis_cache.alone == [] + + +@pytest.mark.asyncio +async def test_prefetch_settlement_keeps_newer_memory_values_and_backfills_misses(monkeypatch): + clock: Final = 1_000_000.0 + monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock) + dep_a_key: Final = CooldownCache.get_cooldown_cache_key("dep-a") + dep_b_key: Final = CooldownCache.get_cooldown_cache_key("dep-b") + old_cooldown: Final = { + "exception_received": "429", + "status_code": "429", + "timestamp": _FAR_FUTURE, + "cooldown_time": 60, + } + newer_memory_cooldown: Final = { + "exception_received": "429", + "status_code": "429", + "timestamp": _FAR_FUTURE + 1, + "cooldown_time": 60, + } + redis_only_cooldown: Final = { + "exception_received": "429", + "status_code": "429", + "timestamp": _FAR_FUTURE + 2, + "cooldown_time": 60, + } + + def replies(command: tuple[Any, ...]) -> Any: + if command[0] == "MGET": + return [json.dumps(redis_cache.store[key]) if key in redis_cache.store else None for key in command[1:]] + return _lua_ok_replies(command) + + client: Final = FakeClient(replies) + redis_cache: Final = FakeRedisCache(client) + redis_cache.store[dep_a_key] = old_cooldown + redis_cache.store[dep_b_key] = redis_only_cooldown + router: Final = _router(redis_cache, routing_strategy="simple-shuffle") + memory_cache: Final = router.cooldown_cache.cooldown_store.in_memory_cache + assert memory_cache is not None + + with request_redis_batch_scope() as request: + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + memory_cache.set_cache(dep_a_key, newer_memory_cooldown) + await request.flush_all() + + prefetched_mgets: Final = tuple(command for command in client.pipelines[0].commands if command[0] == "MGET") + + assert len(prefetched_mgets) == 1 + assert frozenset(prefetched_mgets[0][1:]) == frozenset({dep_a_key, dep_b_key}) + assert memory_cache.get_cache(dep_a_key) == newer_memory_cooldown + assert memory_cache.get_cache(dep_b_key) == redis_only_cooldown + + +@pytest.mark.asyncio +async def test_an_abandoned_prefetch_whose_mget_fails_releases_its_reservation(monkeypatch): + clock: Final = 1_000_000.0 + monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock) + cooldown_key: Final = CooldownCache.get_cooldown_cache_key("dep-a") + active_cooldown: Final = { + "exception_received": "429", + "status_code": "429", + "timestamp": _FAR_FUTURE, + "cooldown_time": 60, + } + mget_replies: Final = iter((ConnectionError("redis down"), None)) + + def replies(command: tuple[Any, ...]) -> Any: + if command[0] == "MGET": + response: Final = next(mget_replies) + if isinstance(response, Exception): + return response + return [json.dumps(active_cooldown) if key == cooldown_key else None for key in command[1:]] + return _lua_ok_replies(command) + + client: Final = FakeClient(replies) + redis_cache: Final = FakeRedisCache(client) + redis_cache.store[cooldown_key] = active_cooldown + router: Final = _router(redis_cache, routing_strategy="simple-shuffle") + + with request_redis_batch_scope() as request: + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + await request.flush_all() + + pipeline_count: Final = len(client.pipelines) + first_cooldown_mgets: Final = tuple( + command + for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[:pipeline_count]) + if command[0] == "MGET" and cooldown_key in command[1:] + ) + with request_redis_batch_scope(): + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + deployment: Final = await router.async_get_available_deployment( + model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} + ) + cooldowns: Final = await router.cooldown_cache.async_get_active_cooldowns(["dep-a"], parent_otel_span=None) + + second_cooldown_mgets: Final = tuple( + command + for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[pipeline_count:]) + if command[0] == "MGET" and cooldown_key in command[1:] + ) + + assert deployment["model_info"]["id"] == "dep-b" + assert [model_id for model_id, _ in cooldowns] == ["dep-a"] + assert len(first_cooldown_mgets) == 1 + assert len(second_cooldown_mgets) == 1 + + +@pytest.mark.asyncio +async def test_a_cooldown_that_leaves_memory_before_routing_is_read_again(monkeypatch): + clock: Final = 1_000_000.0 + monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock) + cooldown_key: Final = CooldownCache.get_cooldown_cache_key("dep-a") + active_cooldown: Final = { + "exception_received": "429", + "status_code": "429", + "timestamp": _FAR_FUTURE, + "cooldown_time": 60, + } + client: Final = FakeClient(_lua_ok_replies) + redis_cache: Final = FakeRedisCache(client) + redis_cache.store[cooldown_key] = active_cooldown + router: Final = _router(redis_cache, routing_strategy="simple-shuffle") + cooldown_store: Final = router.cooldown_cache.cooldown_store + memory_cache: Final = cooldown_store.in_memory_cache + assert memory_cache is not None + memory_cache.set_cache(cooldown_key, active_cooldown) + + with request_redis_batch_scope() as request: + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + await request.flush_all() + memory_cache.delete_cache(cooldown_key) + deployment: Final = await router.async_get_available_deployment( + model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} + ) + + prefetched_mgets: Final = tuple(command for command in client.pipelines[0].commands if command[0] == "MGET") + fallback_cooldown_mgets: Final = tuple( + keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_key in keys + ) + + assert len(prefetched_mgets) == 1 + assert prefetched_mgets[0][1:] == (CooldownCache.get_cooldown_cache_key("dep-b"),) + assert deployment["model_info"]["id"] == "dep-b" + assert fallback_cooldown_mgets == ((cooldown_key,),) @pytest.mark.asyncio @@ -422,6 +667,182 @@ async def test_simple_shuffle_prefetches_only_its_cooldown_read_into_the_admissi assert armed.keys == {CooldownCache.get_cooldown_cache_key("dep-a")} # no usage counters for shuffle +@pytest.mark.asyncio +@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing-v2"]) +@pytest.mark.parametrize("with_limiter", [True, False]) +async def test_requests_within_the_cooldown_read_interval_read_cooldowns_from_redis_once( + routing_strategy: str, with_limiter: bool +): + client = FakeClient(_lua_ok_replies) + redis_cache = FakeRedisCache(client) + router = _router(redis_cache, routing_strategy=routing_strategy) + limiter = _limiter(redis_cache) + request_round_trips: list[tuple[int, int]] = [] + + for _ in range(3): + pipeline_count = len(client.pipelines) + alone_count = len(redis_cache.alone) + with request_redis_batch_scope(): + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + if with_limiter: + await limiter.atomic_check_and_increment_by_n( + descriptors=[_descriptor("api_key", "k1", 10)], + increments=[{"requests": 1}], + ) + await router.async_get_available_deployment( + model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} + ) + request_round_trips.append((len(client.pipelines) - pipeline_count, len(redis_cache.alone) - alone_count)) + + pipeline_mgets = [command for pipeline in client.pipelines for command in pipeline.commands if command[0] == "MGET"] + alone_mgets = [keys for command, keys in redis_cache.alone if command == "MGET"] + cooldown_keys = { + CooldownCache.get_cooldown_cache_key("dep-a"), + CooldownCache.get_cooldown_cache_key("dep-b"), + } + cooldown_mgets = [command[1:] for command in pipeline_mgets if cooldown_keys.intersection(command[1:])] + [ + keys for keys in alone_mgets if cooldown_keys.intersection(keys) + ] + + assert len(cooldown_mgets) == 1 + if not with_limiter: + assert request_round_trips[1:] == [(0, 0), (0, 0)] + + +@pytest.mark.asyncio +async def test_concurrent_requests_share_one_cooldown_read_per_interval(): + client: Final = FakeClient(_lua_ok_replies) + redis_cache: Final = FakeRedisCache(client) + router: Final = _router(redis_cache) + first_armed: Final = asyncio.Event() + both_armed: Final = asyncio.Event() + + async def route_after_both_requests_arm(): + with request_redis_batch_scope(): + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + if first_armed.is_set(): + both_armed.set() + else: + first_armed.set() + await both_armed.wait() + return await router.async_get_available_deployment( + model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} + ) + + deployments: Final = await asyncio.gather(route_after_both_requests_arm(), route_after_both_requests_arm()) + cooldown_keys: Final = frozenset( + { + CooldownCache.get_cooldown_cache_key("dep-a"), + CooldownCache.get_cooldown_cache_key("dep-b"), + } + ) + cooldown_mgets: Final = tuple( + command + for pipeline in client.pipelines + for command in pipeline.commands + if command[0] == "MGET" and cooldown_keys.intersection(command[1:]) + ) + + assert all(deployment["model_info"]["id"] in {"dep-a", "dep-b"} for deployment in deployments) + assert len(cooldown_mgets) == 1 + + +@pytest.mark.asyncio +async def test_the_prefetch_reads_cooldowns_again_once_the_read_interval_elapses(monkeypatch): + first_time: Final = 1_000_000.0 + monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time) + active_cooldown = { + "exception_received": "429", + "status_code": "429", + "timestamp": _FAR_FUTURE, + "cooldown_time": 60, + } + mget_results = iter((None, active_cooldown)) + + def replies(command: tuple[Any, ...]) -> Any: + if command[0] == "MGET": + result = next(mget_results) + return [ + None if result is None or key != CooldownCache.get_cooldown_cache_key("dep-a") else json.dumps(result) + for key in command[1:] + ] + return _lua_ok_replies(command) + + client = FakeClient(replies) + redis_cache = FakeRedisCache(client) + router = _router(redis_cache, routing_strategy="simple-shuffle") + cooldown_store = router.cooldown_cache.cooldown_store + + with request_redis_batch_scope(): + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + await router.async_get_available_deployment( + model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} + ) + + monkeypatch.setattr( + dual_cache_module.time, + "time", + lambda: first_time + cooldown_store.redis_batch_cache_expiry + 1, + ) + with request_redis_batch_scope(): + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + deployment = await router.async_get_available_deployment( + model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} + ) + + cooldown_keys = { + CooldownCache.get_cooldown_cache_key("dep-a"), + CooldownCache.get_cooldown_cache_key("dep-b"), + } + cooldown_mgets = [ + command + for pipeline in client.pipelines + for command in pipeline.commands + if command[0] == "MGET" and cooldown_keys.intersection(command[1:]) + ] + assert len(cooldown_mgets) == 2 + assert deployment["model_info"]["id"] == "dep-b" + + +@pytest.mark.asyncio +async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(monkeypatch): + first_time: Final = 1_000_000.0 + monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time) + client = FakeClient(_lua_ok_replies) + redis_cache = FakeRedisCache(client) + router = _router(redis_cache, routing_strategy="usage-based-routing-v2") + cooldown_store = router.cooldown_cache.cooldown_store + usage_cache = router.lowesttpm_logger_v2.router_cache + time_offset = cooldown_store.redis_batch_cache_expiry + 0.5 + + assert time_offset < usage_cache.redis_batch_cache_expiry + + with request_redis_batch_scope(): + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + await router.async_get_available_deployment( + model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} + ) + + monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time + time_offset) + with request_redis_batch_scope(): + router.arm_routing_read_prefetch(_MODEL_GROUP, {}) + await router.async_get_available_deployment( + model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={} + ) + + cooldown_keys = frozenset( + { + CooldownCache.get_cooldown_cache_key("dep-a"), + CooldownCache.get_cooldown_cache_key("dep-b"), + } + ) + second_pipeline_mgets = tuple(command for command in client.pipelines[1].commands if command[0] == "MGET") + + assert len(client.pipelines) == 2 + assert len(second_pipeline_mgets) == 1 + assert frozenset(second_pipeline_mgets[0][1:]) == cooldown_keys + + @pytest.mark.asyncio async def test_two_backends_flush_concurrently_one_pipeline_each(): a_client, b_client = FakeClient(_lua_ok_replies), FakeClient(_lua_ok_replies) @@ -528,3 +949,91 @@ async def test_auth_write_back_outside_a_scope_writes_through_as_before(): assert [(op[0], [(key, ttl) for key, _value, ttl in op[1]]) for op in redis_cache.alone] == [ ("SET_PIPELINE", [("user-1", 42)]) ] + + +@pytest.mark.asyncio +async def test_a_key_the_request_mget_read_as_absent_is_not_read_again_by_a_per_key_get(): + client = FakeClient(_lua_ok_replies) + redis_cache = FakeRedisCache(client) + cache = UserApiKeyCache(redis_cache=redis_cache) + with request_redis_batch_scope() as request: + assert await request.batch(redis_cache).mget(["absent-key"]) == {"absent-key": None} + assert await cache.async_get_cache("absent-key") is None + assert redis_cache.alone == [] and len(client.pipelines) == 1 + await cache.async_set_cache("absent-key", {"v": 1}, ttl=5) + await request.flush_all() + assert [c[:2] for c in client.pipelines[1].commands] == [("SET", "absent-key")] + + +@pytest.mark.asyncio +async def test_management_object_writes_inside_a_request_ride_its_pipeline_and_write_through_outside(): + client = FakeClient(_lua_ok_replies) + redis_cache = FakeRedisCache(client) + cache = UserApiKeyCache(redis_cache=redis_cache) + with request_redis_batch_scope() as request: + await cache.async_set_cache("team_id:t1", {"team_id": "t1"}, ttl=60) + await cache.async_set_cache("hashed-key-object", {"token": "hashed-key-object"}, ttl=60) + assert client.pipelines == [] + assert cache.in_memory_cache.get_cache("team_id:t1") == {"team_id": "t1"} + assert await cache.async_get_cache("hashed-key-object") == {"token": "hashed-key-object"} + await request.flush_all() + assert sorted((c[0], c[1], c[3]) for c in client.pipelines[0].commands) == [ + ("SET", "hashed-key-object", 60), + ("SET", "team_id:t1", 60), + ] + await cache.async_set_cache("team_id:t2", {"team_id": "t2"}, ttl=60) + assert len(client.pipelines) == 1 + assert redis_cache.alone == [("SET", "team_id:t2", {"team_id": "t2"})] + + +@pytest.mark.asyncio +async def test_a_team_refresh_inside_a_request_sends_its_set_and_alias_del_in_one_pipeline_before_returning(): + client = FakeClient(_lua_ok_replies) + redis_cache = FakeRedisCache(client) + cache = UserApiKeyCache(redis_cache=redis_cache) + usage_cache = DualCache(redis_cache=redis_cache) + usage_cache.in_memory_cache.set_cache("team_id:t1", "stale team") + usage_cache.in_memory_cache.set_cache("team_alias:alpha", "stale alias") + cache.in_memory_cache.set_cache("team_alias:alpha", "stale alias") + proxy_logging_obj = MagicMock() + proxy_logging_obj.internal_usage_cache = InternalUsageCache(dual_cache=usage_cache) + team = LiteLLM_TeamTableCachedObj(team_id="t1", team_alias="alpha") + with request_redis_batch_scope() as request: + await _cache_team_object("t1", team, cache, proxy_logging_obj) + assert [c[:2] for c in client.pipelines[0].commands] == [("SET", "team_id:t1"), ("DEL", "team_alias:alpha")], ( + "the alias DEL must reach Redis before the refresh returns, or another request can refill memory from it" + ) + assert redis_cache.alone == [] + assert usage_cache.in_memory_cache.get_cache("team_id:t1") is None + assert usage_cache.in_memory_cache.get_cache("team_alias:alpha") is None + assert cache.in_memory_cache.get_cache("team_alias:alpha") is None + assert cache.in_memory_cache.get_cache("team_id:t1")["team_id"] == "t1" + await request.flush_all() + assert len(client.pipelines) == 1 and redis_cache.alone == [] + + +@pytest.mark.asyncio +async def test_a_pipelined_management_write_without_a_ttl_expires_in_redis_like_the_direct_path(): + client = FakeClient(_lua_ok_replies) + redis_cache = FakeRedisCache(client) + cache = UserApiKeyCache(redis_cache=redis_cache) + cache.update_cache_ttl(default_in_memory_ttl=5, default_redis_ttl=None) + with request_redis_batch_scope() as request: + await cache.async_set_cache("team_id:t1", {"team_id": "t1"}) + await request.flush_all() + assert [(c[0], c[1], c[3]) for c in client.pipelines[0].commands] == [("SET", "team_id:t1", 5)] + + +@pytest.mark.asyncio +async def test_identity_prefetch_is_one_mget_after_which_hits_and_misses_alike_cost_no_read(): + client = FakeClient(replies) + redis_cache = FakeRedisCache(client) + cache = UserApiKeyCache(redis_cache=redis_cache) + with request_redis_batch_scope(): + await prefetch_identity_keys(["key-hit", "end_user_id:eu-miss", "key-hit"], cache) + assert [c[0] for c in client.pipelines[0].commands] == ["MGET"] + assert sorted(client.pipelines[0].commands[0][1:]) == ["end_user_id:eu-miss", "key-hit"] + assert await cache.async_get_cache("key-hit") == {"k": "key-hit"} + assert await cache.async_get_cache("end_user_id:eu-miss") is None + assert len(client.pipelines) == 1 and redis_cache.alone == [] + assert cache.in_memory_cache.get_cache("end_user_id:eu-miss") is None diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 45fc93f04c1..0375ff14852 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1879,6 +1879,20 @@ class TestEncryptedReasoningReplay: assert messages[0] == {"role": "user", "content": "question"} assert messages[2] == {"role": "user", "content": [{"type": "text", "text": "follow-up"}]} + def test_strip_uses_predicate_to_keep_selected_encrypted_blocks(self): + kept_signature = encrypted_reasoning_signature("keep") + stripped_signature = encrypted_reasoning_signature("strip") + content = [ + {"type": "thinking", "thinking": "keep", "signature": kept_signature}, + {"type": "thinking", "thinking": "strip", "signature": stripped_signature}, + ] + messages = [{"role": "assistant", "content": content}] + + strip_encrypted_reasoning_from_messages(messages, should_strip=lambda block: block.get("thinking") == "strip") + + assert messages[0]["content"] is content + assert content == [{"type": "thinking", "thinking": "keep", "signature": kept_signature}] + @pytest.mark.parametrize( "messages", [ diff --git a/tests/unit/llms/anthropic/test_anthropic_common_utils.py b/tests/unit/llms/anthropic/test_anthropic_common_utils.py index 52b53769457..b80129a55bf 100644 --- a/tests/unit/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/unit/llms/anthropic/test_anthropic_common_utils.py @@ -2345,3 +2345,34 @@ class TestMalformedContentListItems: api_key=FAKE_REGULAR_KEY, max_tokens=5, ) + + +@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") +@pytest.mark.parametrize("nested_output_config", [False, True]) +@pytest.mark.parametrize("explicit_beta", [False, True]) +@pytest.mark.parametrize("output_config", [{}, {"effort": "high"}, {"format": {"type": "text"}}]) +def test_validate_environment_adds_mid_conversation_output_config_beta( + nested_output_config: bool, explicit_beta: bool, output_config: dict[str, object] +) -> None: + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + from litellm.types.llms.anthropic import ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER + + beta: Final = ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER + + messages: Final = [ + {"role": "user", "content": "Hello"}, + *([{"role": "system", "content": [], "output_config": output_config}] if nested_output_config else []), + {"role": "user", "content": "Reply with OK"}, + ] + + headers: Final = AnthropicModelInfo().validate_environment( + headers={"anthropic-beta": beta} if explicit_beta else {}, + model="claude-fable-5-1", + messages=messages, + optional_params={"output_config": {"effort": "high"}}, + litellm_params={}, + api_key=FAKE_REGULAR_KEY, + ) + + assert headers.get("anthropic-beta", "").split(",").count(beta) == int(nested_output_config or explicit_beta) + assert headers["x-api-key"] == FAKE_REGULAR_KEY diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 79207ece259..f92de7370bd 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3494,3 +3494,56 @@ async def test_get_async_streaming_response_iterator_yields_small_frame_before_u remaining: Final = tuple([chunk async for chunk in iterator]) assert any(chunk.startswith(b"event: message_stop\n") for chunk in remaining), remaining await iterator.aclose() + + +@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") +@pytest.mark.parametrize("nested_output_config", [False, True]) +@pytest.mark.parametrize("explicit_beta", [False, True]) +@pytest.mark.parametrize("output_config", [{}, {"effort": "high"}, {"format": {"type": "text"}}]) +def test_bedrock_messages_mid_conversation_output_config_beta( + nested_output_config: bool, explicit_beta: bool, output_config: dict[str, object] +) -> None: + from litellm.types.llms.anthropic import ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER + from litellm.types.router import GenericLiteLLMParams + + beta: Final = ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER + + messages: Final = [ + {"role": "user", "content": "Hello"}, + *([{"role": "system", "content": [], "output_config": output_config}] if nested_output_config else []), + {"role": "user", "content": "Reply with OK"}, + ] + + result: Final = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="global.anthropic.claude-fable-5-1", + messages=messages, + anthropic_messages_optional_request_params={"max_tokens": 1024, "output_config": {"effort": "high"}}, + litellm_params=GenericLiteLLMParams(), + headers={"anthropic-beta": beta} if explicit_beta else {}, + ) + + assert result.get("anthropic_beta", []).count(beta) == int(nested_output_config or explicit_beta) + assert result["messages"] == messages + assert result["output_config"] == {"effort": "high"} + + +@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") +@pytest.mark.parametrize("explicit_beta", [False, True]) +def test_bedrock_messages_removed_output_config_does_not_add_beta(explicit_beta: bool) -> None: + from litellm.types.llms.anthropic import ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER + from litellm.types.router import GenericLiteLLMParams + + beta: Final = ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER + result: Final = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="global.anthropic.claude-fable-5-1", + messages=[ + {"role": "system", "content": "Answer briefly", "output_config": {"effort": "high"}}, + {"role": "user", "content": "Reply with OK"}, + ], + anthropic_messages_optional_request_params={"max_tokens": 1024}, + litellm_params=GenericLiteLLMParams(), + headers={"anthropic-beta": beta} if explicit_beta else {}, + ) + + assert result["messages"] == [{"role": "user", "content": "Reply with OK"}] + assert result.get("anthropic_beta", []).count(beta) == int(explicit_beta) diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm.py b/tests/unit/router_strategy/test_lowest_tpm_rpm.py index 0fa11cda20c..625f648bec4 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm.py @@ -72,11 +72,48 @@ async def test_v2_async_selection_uses_prefetched_counters_only_when_they_cover_ keys = tpm_keys + rpm_keys covering = PrefetchedUsage(keys=frozenset(keys), values=dict(zip(keys, [10, 100, None, None]))) - chosen = await strategy.async_get_available_deployments(model_group="g", healthy_deployments=deployments, prefetched_usage=covering) + with PrefetchedUsage.scoped(covering): + chosen: Final = await strategy.async_get_available_deployments(model_group="g", healthy_deployments=deployments) assert chosen["model_info"]["id"] == "a", "the prefetched counters say a is the lowest" router_cache.async_batch_get_cache.assert_not_awaited() stale = PrefetchedUsage(keys=frozenset(keys[:1]), values={keys[0]: 10}) - chosen = await strategy.async_get_available_deployments(model_group="g", healthy_deployments=deployments, prefetched_usage=stale) - assert chosen["model_info"]["id"] == "b", "counters that do not cover this minute's keys are read again" + with PrefetchedUsage.scoped(stale): + chosen_stale: Final = await strategy.async_get_available_deployments( + model_group="g", healthy_deployments=deployments + ) + assert chosen_stale["model_info"]["id"] == "b", "counters that do not cover this minute's keys are read again" router_cache.async_batch_get_cache.assert_awaited_once_with(keys=keys) + + +@pytest.mark.asyncio +async def test_v2_subclass_overriding_async_get_available_deployments_with_the_old_signature_still_routes() -> None: + class OldSignatureV2(LowestTPMLoggingHandler_v2): + async def async_get_available_deployments( + self, + model_group: str, + healthy_deployments: list, + messages: list[dict[str, str]] | None = None, + input: str | list | None = None, + ): + return await super().async_get_available_deployments( + model_group=model_group, + healthy_deployments=healthy_deployments, + messages=messages, + input=input, + ) + + router: Final = Router( + model_list=[_deployment(HIGH_USAGE_DEPLOYMENT_ID), _deployment(LOW_USAGE_DEPLOYMENT_ID)], + routing_strategy="usage-based-routing-v2", + ) + router.lowesttpm_logger_v2 = OldSignatureV2(router_cache=router.cache, routing_args={}) + + response: Final = await router.acompletion( + model=MODEL_GROUP, messages=[{"role": "user", "content": "x"}] + ) + + assert response.choices[0].message.content in { + f"from {HIGH_USAGE_DEPLOYMENT_ID}", + f"from {LOW_USAGE_DEPLOYMENT_ID}", + } diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 836049c88a2..3a92aa221e5 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -1961,6 +1961,315 @@ class TestStripEncryptedReasoningFromInput: ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) assert request_input == before + def test_strips_only_items_selected_by_predicate(self): + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") + request_input = [ + {"type": "reasoning", "id": "keep", "encrypted_content": wrapped, "summary": "keep"}, + {"type": "reasoning", "id": "strip", "encrypted_content": wrapped, "summary": "strip"}, + ] + + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input( + request_input, should_strip=lambda item: item.get("id") == "strip" + ) + + assert request_input == [ + {"type": "reasoning", "id": "keep", "encrypted_content": wrapped, "summary": "keep"}, + {"type": "reasoning", "summary": "strip"}, + ] + + +@pytest.mark.asyncio +async def test_real_router_selection_keeps_origin_reasoning_and_strips_foreign_origin(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-openai", + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_base": "https://api.openai.com/v1", + "api_key": "key-openai", + }, + "model_info": {"id": "dep-openai"}, + }, + { + "model_name": "gpt-azure", + "litellm_params": { + "model": "azure/gpt-5.1-codex", + "api_base": "https://res-b.openai.azure.com/", + "api_key": "key-azure", + "api_version": "2025-04-01-preview", + }, + "model_info": {"id": "dep-azure"}, + }, + ], + optional_pre_call_checks=["encrypted_content_affinity"], + num_retries=0, + ) + openai_item_id = ResponsesAPIRequestUtils._build_encrypted_item_id("dep-openai", "rs-openai") + azure_item_id = ResponsesAPIRequestUtils._build_encrypted_item_id("dep-azure", "rs-azure") + openai_wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("blob-openai", "dep-openai") + azure_wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("blob-azure", "dep-azure") + request_input = [ + {"type": "message", "role": "user", "content": "first question"}, + { + "type": "reasoning", + "id": openai_item_id, + "encrypted_content": openai_wrapped, + "summary": [{"type": "summary_text", "text": "openai summary"}], + }, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "first answer"}]}, + {"type": "message", "role": "user", "content": "second question"}, + { + "type": "reasoning", + "id": azure_item_id, + "encrypted_content": azure_wrapped, + "summary": [{"type": "summary_text", "text": "azure summary"}], + }, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "second answer"}]}, + {"type": "message", "role": "user", "content": "third question"}, + ] + + request_kwargs = {"input": request_input, "store": False} + try: + deployment = await router.async_get_available_deployment( + model="gpt-openai", request_kwargs=request_kwargs, input=request_kwargs["input"] + ) + + assert deployment["model_info"]["id"] == "dep-openai" + assert deployment["litellm_params"]["model"] == "openai/gpt-5.1-codex" + assert deployment["litellm_params"]["api_base"] == "https://api.openai.com/v1" + assert request_kwargs["input"] == [ + {"type": "message", "role": "user", "content": "first question"}, + { + "type": "reasoning", + "id": openai_item_id, + "encrypted_content": openai_wrapped, + "summary": [{"type": "summary_text", "text": "openai summary"}], + }, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "first answer"}]}, + {"type": "message", "role": "user", "content": "second question"}, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "azure summary"}], + }, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "second answer"}]}, + {"type": "message", "role": "user", "content": "third question"}, + ] + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_affinity_keeps_mixed_origins_on_the_same_encryption_boundary(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + shared_api_base = "https://account-a.openai.azure.com/" + shared_api_key = "shared-key" + origin_d2 = _make_originating_mock(shared_api_base, shared_api_key) + mock_router = _make_router_mock_with_cooldown(origin_d2, cooldown_entries=[], routed_group_model_ids=["d1", "d2"]) + deployment_d1 = { + "model_info": {"id": "d1"}, + "litellm_params": {"api_base": shared_api_base, "api_key": shared_api_key}, + } + deployment_d2 = { + "model_info": {"id": "d2"}, + "litellm_params": {"api_base": shared_api_base, "api_key": shared_api_key}, + } + d2_item = { + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("blob-d2", "d2"), + "summary": [{"type": "summary_text", "text": "second origin"}], + } + request_kwargs = { + "input": [ + { + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("blob-d1", "d1"), + "summary": [{"type": "summary_text", "text": "first origin"}], + }, + d2_item.copy(), + ] + } + mock_router.get_deployment.side_effect = lambda model_id: origin_d2 if model_id == "d2" else None + check = EncryptedContentAffinityCheck(router=mock_router) + + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=[deployment_d1, deployment_d2], + messages=None, + request_kwargs=request_kwargs, + ) + + assert result == [deployment_d1] + assert request_kwargs["input"][1] == d2_item + + +@pytest.mark.asyncio +async def test_boundary_pin_strips_reasoning_from_a_different_origin(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + origin_a = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") + origin_b = _make_originating_mock("https://account-b.openai.azure.com/", "key-b") + mock_router = _make_router_mock_with_cooldown(origin_a, cooldown_entries=[], routed_group_model_ids=["peer-a"]) + mock_router.get_deployment.side_effect = lambda model_id: {"origin-a": origin_a, "origin-b": origin_b}.get(model_id) + peer_a = { + "model_info": {"id": "peer-a"}, + "litellm_params": {"api_base": "https://account-a.openai.azure.com/", "api_key": "key-a"}, + } + request_kwargs = { + "input": [ + { + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + "blob-origin-a", "origin-a" + ), + "summary": [{"type": "summary_text", "text": "origin A summary"}], + }, + { + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + "blob-origin-b", "origin-b" + ), + "summary": [{"type": "summary_text", "text": "origin B summary"}], + }, + ] + } + check = EncryptedContentAffinityCheck(router=mock_router) + + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=[peer_a], + messages=None, + request_kwargs=request_kwargs, + ) + + assert result == [peer_a] + assert request_kwargs["input"] == [ + { + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + "blob-origin-a", "origin-a" + ), + "summary": [{"type": "summary_text", "text": "origin A summary"}], + }, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "origin B summary"}]}, + ] + + +@pytest.mark.asyncio +async def test_affinity_keeps_only_anthropic_reasoning_from_the_pinned_origin(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + origin_a = _make_originating_mock("https://account-a.openai.azure.com/", "key-a") + origin_b = _make_originating_mock("https://account-b.openai.azure.com/", "key-b") + mock_router = _make_router_mock_with_cooldown(origin_b, cooldown_entries=[], routed_group_model_ids=["origin-a"]) + mock_router.get_deployment.side_effect = lambda model_id: {"origin-a": origin_a, "origin-b": origin_b}.get(model_id) + deployment_a = { + "model_info": {"id": "origin-a"}, + "litellm_params": {"api_base": "https://account-a.openai.azure.com/", "api_key": "key-a"}, + } + deployment_b = { + "model_info": {"id": "origin-b"}, + "litellm_params": {"api_base": "https://account-b.openai.azure.com/", "api_key": "key-b"}, + } + messages = _bridge_replayed_anthropic_messages(minted_by="origin-a") + foreign_messages = _bridge_replayed_anthropic_messages(minted_by="origin-b") + assistant_content = messages[1]["content"] + assistant_content.insert(3, foreign_messages[1]["content"][1]) + check = EncryptedContentAffinityCheck(router=mock_router) + + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=[deployment_a, deployment_b], + messages=messages, + request_kwargs={"model": "gpt-5.4"}, + ) + + assert result == [deployment_a] + assert messages[1]["content"] is assistant_content + assert assistant_content == [ + {"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"}, + { + "type": "redacted_thinking", + "data": ( + "litellm_encrypted_reasoning:" + f"{ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id('gAAAAA_turn_one', 'origin-a')}" + ), + }, + { + "type": "thinking", + "thinking": "The bridge packed this one", + "signature": ( + "litellm_encrypted_reasoning:" + f"{ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id('gAAAAA_turn_one', 'origin-a')}" + ), + }, + {"type": "text", "text": "The zebra owner lives in the green house."}, + ] + + +@pytest.mark.asyncio +async def test_affinity_strips_unknown_origins_but_leaves_unmarked_encrypted_content(): + from unittest.mock import MagicMock + + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + mock_router = MagicMock() + mock_router.get_deployment.return_value = None + deployment_a = { + "model_info": {"id": "origin-a"}, + "litellm_params": {"api_base": "https://account-a.openai.azure.com/", "api_key": "key-a"}, + } + openai_item = { + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("blob-a", "origin-a"), + "summary": [{"type": "summary_text", "text": "origin A"}], + } + request_kwargs = { + "input": [ + openai_item.copy(), + { + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + "blob-removed", "origin-removed" + ), + "summary": [{"type": "summary_text", "text": "removed origin"}], + }, + { + "type": "reasoning", + "encrypted_content": "raw-encrypted-content", + "summary": [{"type": "summary_text", "text": "unmarked content"}], + }, + ] + } + check = EncryptedContentAffinityCheck(router=mock_router) + + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=[deployment_a], + messages=None, + request_kwargs=request_kwargs, + ) + + assert result == [deployment_a] + assert request_kwargs["input"] == [ + openai_item, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "removed origin"}]}, + { + "type": "reasoning", + "encrypted_content": "raw-encrypted-content", + "summary": [{"type": "summary_text", "text": "unmarked content"}], + }, + ] + def _cross_group_request_kwargs(): wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a") diff --git a/tests/unit/router_utils/test_routing_read_batch.py b/tests/unit/router_utils/test_routing_read_batch.py index 73be5fd4a3a..a74e3e24f99 100644 --- a/tests/unit/router_utils/test_routing_read_batch.py +++ b/tests/unit/router_utils/test_routing_read_batch.py @@ -6,6 +6,7 @@ Before `RoutingReadBatch`, `async_get_available_deployment` issued one MGET for """ import time +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -13,6 +14,7 @@ import pytest import litellm from litellm import Router from litellm.caching.redis_cache import RedisCache +from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 _MODEL_GROUP = "claude" _MESSAGES = [{"role": "user", "content": "ping"}] @@ -79,6 +81,46 @@ async def test_usage_based_routing_reads_cooldowns_and_counters_in_one_redis_rou ], "cooldown state and usage counters must arrive in one MGET" +@pytest.mark.asyncio +async def test_usage_based_routing_still_batches_when_the_strategy_is_a_fixed_signature_subclass(): + class OldSignatureV2(LowestTPMLoggingHandler_v2): + async def async_get_available_deployments( + self, + model_group: str, + healthy_deployments: list, + messages: list[dict[str, str]] | None = None, + input: str | list | None = None, + ): + return await super().async_get_available_deployments( + model_group=model_group, + healthy_deployments=healthy_deployments, + messages=messages, + input=input, + ) + + redis: Final = _redis_answering({}) + router: Final = _router(redis, "usage-based-routing-v2") + router.lowesttpm_logger_v2 = OldSignatureV2(router_cache=router.cache) + router.cache.async_batch_get_cache = AsyncMock(wraps=router.cache.async_batch_get_cache) + + deployment: Final = await router.async_get_available_deployment( + model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES + ) + + assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} + assert _redis_key_families(redis) == [ + [ + "dep-a:anthropic/claude-x:rpm", + "dep-a:anthropic/claude-x:tpm", + "dep-b:anthropic/claude-x:rpm", + "dep-b:anthropic/claude-x:tpm", + "deployment:dep-a:cooldown", + "deployment:dep-b:cooldown", + ] + ], "the subclassed strategy must still get the batched read, not a second MGET" + router.cache.async_batch_get_cache.assert_not_awaited() + + @pytest.mark.asyncio async def test_simple_shuffle_still_reads_only_cooldowns(): redis = _redis_answering({}) diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 2fe09e6de98..36e188e82d6 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -26,6 +26,7 @@ from litellm.types.utils import ( CacheCreationTokenDetails, CallTypes, Choices, + EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, @@ -160,6 +161,80 @@ def test_cost_calculator_with_response_cost_in_additional_headers(): assert result == 1000 +def test_response_cost_calculator_keeps_optional_params_out_of_hidden_params(): + class MockResponse(BaseModel): + pass + + response = MockResponse() + response._hidden_params = {"custom_llm_provider": "openai"} + optional_params = { + "dimensions": 256, + "extra_headers": {"x-goog-api-key": "goog-secret"}, + "aws_session_token": "session-secret", + } + + response_cost_calculator( + response_object=response, + model="text-embedding-3-small", + custom_llm_provider="openai", + call_type="embedding", + optional_params=optional_params, + ) + + assert response._hidden_params == {"custom_llm_provider": "openai"} + assert optional_params["extra_headers"] == {"x-goog-api-key": "goog-secret"} + assert optional_params["aws_session_token"] == "session-secret" + + +def test_embedding_success_logging_and_spend_log_carry_no_forwarded_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_proxy_server_request_for_spend_logs_payload + + monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True}) + shared_metadata: dict[str, object] = {"user_api_key_alias": "alias"} + proxy_server_request: Final = {"body": {"model": "emb", "input": "hi", "metadata": shared_metadata}} + shared_optional_params: dict[str, object] = {"encoding_format": "float"} + logging_obj = Logging( + model="text-embedding-3-small", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="aembedding", + start_time=datetime.datetime.now(), + litellm_call_id="embedding-hidden-params", + function_id="f", + ) + logging_obj.update_environment_variables( + model="text-embedding-3-small", + litellm_params={"metadata": shared_metadata, "proxy_server_request": proxy_server_request}, + optional_params=shared_optional_params, + custom_llm_provider="openai", + ) + shared_optional_params["extra_headers"] = {"x-goog-api-key": "goog-secret"} + response = EmbeddingResponse(model="text-embedding-3-small", data=[], usage=Usage(prompt_tokens=3, total_tokens=3)) + response._hidden_params = {"custom_llm_provider": "openai"} + + logging_obj._process_hidden_params_and_response_cost( + response, + start_time=datetime.datetime.now(), + end_time=datetime.datetime.now(), + ) + + litellm_params = logging_obj.model_call_details["litellm_params"] + stored_request: Final = _get_proxy_server_request_for_spend_logs_payload( + metadata=shared_metadata, + litellm_params=litellm_params, + kwargs=logging_obj.model_call_details, + ) + hidden_params = litellm_params["metadata"]["hidden_params"] + assert isinstance(hidden_params, dict) + assert "optional_params" not in hidden_params + assert '"hidden_params"' in stored_request + assert "goog-secret" not in stored_request + assert "goog-secret" not in str(logging_obj.model_call_details["standard_logging_object"]) + assert logging_obj.model_call_details["response_cost"] is not None + assert logging_obj.optional_params["extra_headers"] == {"x-goog-api-key": "goog-secret"} + + diff --git a/tests/unit/test_openai_service_tier_long_context_pricing.py b/tests/unit/test_openai_service_tier_long_context_pricing.py index 9b3a1e57169..9777af1af70 100644 --- a/tests/unit/test_openai_service_tier_long_context_pricing.py +++ b/tests/unit/test_openai_service_tier_long_context_pricing.py @@ -1,10 +1,13 @@ import json from functools import lru_cache from pathlib import Path +from typing import Final import pytest import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.types.utils import PromptTokensDetailsWrapper, Usage REPO_ROOT = Path(__file__).parents[2] MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" @@ -72,7 +75,23 @@ PRIORITY_LONG_CONTEXT = { }, } -EXPECTED = {**FLEX_LONG_CONTEXT, **PRIORITY_LONG_CONTEXT} +ULTRAFAST_LONG_CONTEXT = { + "gpt-6-astra": { + "input_cost_per_token_above_272k_tokens_ultrafast": 0.00012, + "output_cost_per_token_above_272k_tokens_ultrafast": 0.00045, + "cache_read_input_token_cost_above_272k_tokens_ultrafast": 1.2e-05, + "cache_creation_input_token_cost_above_272k_tokens_ultrafast": 0.00015, + } +} + +EXPECTED: Final = { + model: { + **FLEX_LONG_CONTEXT.get(model, {}), + **PRIORITY_LONG_CONTEXT.get(model, {}), + **ULTRAFAST_LONG_CONTEXT.get(model, {}), + } + for model in {**FLEX_LONG_CONTEXT, **PRIORITY_LONG_CONTEXT, **ULTRAFAST_LONG_CONTEXT} +} NO_PUBLISHED_PRIORITY_LONG_CONTEXT = ("gpt-5.4", "gpt-5.5") @@ -102,6 +121,85 @@ TIERED_COST_CASES = [ ("gpt-5.6-terra", "priority", 8e-06, 3.6e-05), ("gpt-5.6-luna", "priority", 8e-07, 3.6e-06), ("gpt-6-astra", "priority", 4e-05, 0.00015), + ("gpt-6-astra", "ultrafast", 0.00012, 0.00045), ("gpt-6-sol", "priority", 8e-06, 3e-05), ("gpt-6-luna", "priority", 4e-07, 1.5e-06), ] + + +@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) +def test_catalogs_contain_expected_tiered_long_context_rates(path: Path) -> None: + catalog: Final = _load(path) + + assert {model: {key: catalog[model][key] for key in rates} for model, rates in EXPECTED.items()} == EXPECTED, ( + "gpt-6-astra ultrafast rates per https://developers.openai.com/api/docs/pricing (2026-09-29)" + ) + + +def test_get_model_info_preserves_expected_tiered_long_context_rates() -> None: + assert { + model: {key: litellm.get_model_info(model)[key] for key in rates} for model, rates in EXPECTED.items() + } == EXPECTED + + +@pytest.mark.parametrize(("model", "service_tier", "input_rate", "output_rate"), TIERED_COST_CASES) +def test_tiered_long_context_cost_uses_catalog_rates( + model: str, service_tier: str, input_rate: float, output_rate: float +) -> None: + usage: Final = Usage( + prompt_tokens=LONG_CONTEXT_PROMPT_TOKENS, + completion_tokens=COMPLETION_TOKENS, + total_tokens=LONG_CONTEXT_PROMPT_TOKENS + COMPLETION_TOKENS, + ) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="openai", + service_tier=service_tier, + ) + + assert prompt_cost == pytest.approx(LONG_CONTEXT_PROMPT_TOKENS * input_rate) + assert completion_cost == pytest.approx(COMPLETION_TOKENS * output_rate) + + +def test_gpt_6_astra_ultrafast_long_context_costs_and_controls() -> None: + ultrafast_usage: Final = Usage( + prompt_tokens=300_000, + completion_tokens=1_000, + total_tokens=301_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100, cache_creation_tokens=200), + ) + ultrafast_prompt_cost, ultrafast_completion_cost = generic_cost_per_token( + model="gpt-6-astra", + usage=ultrafast_usage, + custom_llm_provider="openai", + service_tier="ultrafast", + ) + standard_prompt_cost, standard_completion_cost = generic_cost_per_token( + model="gpt-6-astra", + usage=ultrafast_usage, + custom_llm_provider="openai", + ) + below_threshold_usage: Final = Usage( + prompt_tokens=271_000, + completion_tokens=1_000, + total_tokens=272_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100, cache_creation_tokens=200), + ) + below_threshold_prompt_cost, below_threshold_completion_cost = generic_cost_per_token( + model="gpt-6-astra", + usage=below_threshold_usage, + custom_llm_provider="openai", + service_tier="ultrafast", + ) + + assert (ultrafast_prompt_cost, ultrafast_completion_cost) == pytest.approx( + (299_700 * 0.00012 + 100 * 1.2e-05 + 200 * 0.00015, 1_000 * 0.00045) + ) + assert ultrafast_prompt_cost + ultrafast_completion_cost == pytest.approx(36.4452) + assert (standard_prompt_cost, standard_completion_cost) == pytest.approx( + (299_700 * 0.00002 + 100 * 2e-06 + 200 * 2.5e-05, 1_000 * 7.5e-05) + ) + assert (below_threshold_prompt_cost, below_threshold_completion_cost) == pytest.approx( + (270_700 * 6e-05 + 100 * 6e-06 + 200 * 7.5e-05, 1_000 * 0.0003) + ) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index e4e65f8904c..96dddf15869 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -48,6 +48,7 @@ from litellm.router import ( _is_retriable_anthropic_status, _responses_stream_holds_event, _without_line_breaks, + Span, ) from litellm.router_strategy import simple_shuffle from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit @@ -18837,3 +18838,47 @@ def test_a_failed_routing_read_prefetch_logs_the_request_model_without_its_line_ assert messages == [ "routing read prefetch not armed for gpt-4ERROR forged entry: no deployments for gpt-4ERROR forged entry" ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "routing_strategy", + ["simple-shuffle", "usage-based-routing-v2", "least-busy", "latency-based-routing"], +) +async def test_router_subclass_overriding_async_get_healthy_deployments_with_the_old_signature_still_routes( + routing_strategy: str, +) -> None: + class OldSignatureRouter(litellm.Router): + async def async_get_healthy_deployments( + self, + model: str, + request_kwargs: dict, + messages: list[dict[str, str]] | None = None, + input: str | list | None = None, + specific_deployment: bool | None = False, + parent_otel_span: Span | None = None, + health_check_probe: bool = False, + ): + return await super().async_get_healthy_deployments( + model=model, + request_kwargs=request_kwargs, + messages=messages, + input=input, + specific_deployment=specific_deployment, + parent_otel_span=parent_otel_span, + health_check_probe=health_check_probe, + ) + + router: Final = OldSignatureRouter( + model_list=[ + { + "model_name": "m", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "x", "mock_response": "hi"}, + } + ], + routing_strategy=routing_strategy, + ) + + response: Final = await router.acompletion(model="m", messages=[{"role": "user", "content": "x"}]) + + assert response.choices[0].message.content == "hi" diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index d73f5efa96b..86206da16a1 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -1829,6 +1829,34 @@ def test_register_deployment_in_model_cost_writes_both_key_families(): _restore_model_cost_entries(model_keys) +def test_router_registration_keeps_ultrafast_long_context_deployment_pricing() -> None: + model_id: Final = "ultrafast-long-context-pricing-id" + backend_key: Final = "openai/gpt-6-astra" + rates: Final = { + "input_cost_per_token_above_272k_tokens_ultrafast": 0.00012, + "output_cost_per_token_above_272k_tokens_ultrafast": 0.00045, + "cache_read_input_token_cost_above_272k_tokens_ultrafast": 1.2e-05, + "cache_creation_input_token_cost_above_272k_tokens_ultrafast": 0.00015, + } + model_cost_entries: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (model_id, backend_key, "gpt-6-astra") + } + try: + Router( + model_list=[ + { + "model_name": "ultrafast-long-context-pricing", + "litellm_params": {"model": backend_key, **rates}, + "model_info": {"id": model_id}, + } + ] + ) + + assert {key: litellm.model_cost[model_id][key] for key in rates} == rates + finally: + _restore_model_cost_entries(model_cost_entries) + + def test_reload_keeps_custom_pricing_configured_on_litellm_params_for_a_db_model(): """ A deployment added at runtime, which is what /model/new does, configures its diff --git a/tests/unit/test_router_silent_experiment.py b/tests/unit/test_router_silent_experiment.py index ab65e09e133..722a76fa7ef 100644 --- a/tests/unit/test_router_silent_experiment.py +++ b/tests/unit/test_router_silent_experiment.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.router import Router from litellm.router import _silent_experiment_kwargs_snapshot from litellm.router import _silent_experiment_targets @@ -30,8 +31,20 @@ class _RecordingLogger(CustomLogger): ] +async def _settle_shared_logging_worker() -> None: + try: + await GLOBAL_LOGGING_WORKER.flush() + finally: + await GLOBAL_LOGGING_WORKER.stop() + + @pytest.fixture def recording_logger(): + settle_loop: Final = asyncio.new_event_loop() + try: + settle_loop.run_until_complete(_settle_shared_logging_worker()) + finally: + settle_loop.close() original_callbacks: Final = litellm.callbacks logger: Final = _RecordingLogger() litellm.callbacks = [logger] diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index afc449d22c4..52b1714b4d9 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -768,12 +768,14 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_creation_input_token_cost_above_256k_tokens": {"type": "number"}, "cache_creation_input_token_cost_above_272k_tokens": {"type": "number"}, "cache_creation_input_token_cost_above_272k_tokens_flex": {"type": "number"}, + "cache_creation_input_token_cost_above_272k_tokens_ultrafast": {"type": "number"}, "cache_creation_input_token_cost_above_272k_tokens_priority": {"type": "number"}, "cache_creation_input_token_cost_above_200k_tokens_batches": {"type": "number"}, "cache_creation_input_token_cost_above_272k_tokens_batches": {"type": "number"}, "cache_creation_input_token_cost_batches": {"type": "number"}, "cache_creation_input_token_cost_flex": {"type": "number"}, "cache_creation_input_token_cost_priority": {"type": "number"}, + "cache_creation_input_token_cost_ultrafast": {"type": "number"}, "cache_read_input_token_cost": {"type": "number"}, "cache_read_input_token_cost_above_32k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_128k_tokens": {"type": "number"}, @@ -782,7 +784,9 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_read_input_token_cost_above_256k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_272k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_272k_tokens_flex": {"type": "number"}, + "cache_read_input_token_cost_above_272k_tokens_ultrafast": {"type": "number"}, "cache_read_input_token_cost_above_512k_tokens": {"type": "number"}, + "input_cost_per_token_above_272k_tokens_ultrafast": {"type": "number"}, "cache_read_input_token_cost_batches": {"type": "number"}, "cache_read_input_token_cost_above_272k_tokens_batches": {"type": "number"}, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": {"type": "number"}, @@ -809,11 +813,13 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_read_input_token_cost_flex": {"type": "number"}, "cache_read_input_token_cost_priority": {"type": "number"}, "cache_read_input_token_cost_balanced": {"type": "number"}, + "cache_read_input_token_cost_ultrafast": {"type": "number"}, "cache_read_input_token_cost_above_200k_tokens_priority": {"type": "number"}, "cache_read_input_token_cost_above_272k_tokens_priority": {"type": "number"}, "input_cost_per_token_flex": {"type": "number"}, "input_cost_per_token_priority": {"type": "number"}, "input_cost_per_token_balanced": {"type": "number"}, + "input_cost_per_token_ultrafast": {"type": "number"}, "input_cost_per_token_above_200k_tokens_priority": {"type": "number"}, "input_cost_per_token_above_272k_tokens_priority": {"type": "number"}, "input_cost_per_token_above_272k_tokens_batches": {"type": "number"}, @@ -822,8 +828,10 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_token_flex": {"type": "number"}, "output_cost_per_token_priority": {"type": "number"}, "output_cost_per_token_balanced": {"type": "number"}, + "output_cost_per_token_ultrafast": {"type": "number"}, "output_cost_per_token_above_200k_tokens_priority": {"type": "number"}, "output_cost_per_token_above_272k_tokens_priority": {"type": "number"}, + "output_cost_per_token_above_272k_tokens_ultrafast": {"type": "number"}, "output_cost_per_token_above_272k_tokens_batches": {"type": "number"}, "output_cost_per_token_above_272k_tokens_flex": {"type": "number"}, "regional_endpoint_uplift_multiplier": {"type": "number"}, diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index a2d944fcf39..ab4f6c12431 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -161,6 +161,7 @@ OPTION_NAMES: Final = ( "logger_fn", "verbose", "no-log", + "log_client_error_tracebacks", "max_agentic_loops", "guardrails", "prompt_id", diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 0b71b51dc3f..03e29c02021 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -24,7 +24,7 @@ "dayjs": "1.11.19", "jwt-decode": "4.0.0", "lucide-react": "0.513.0", - "moment": "2.30.1", + "moment": "2.31.0", "next": "16.3.3", "next-themes": "^0.4.6", "nuqs": "^2.9.4", @@ -4965,9 +4965,9 @@ } }, "node_modules/brace-expansion": { - "version": "5.0.9", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.9.tgz", - "integrity": "sha512-ScQ4IuvIEF1TMlP7Zt+vjJ//9zlPb2SDcxWxM3bk8s6t6GGdJ7KO1dCcTidOPJKePW30LE/2cT7wCyPho9/Wxg==", + "version": "5.0.12", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.12.tgz", + "integrity": "sha512-YovQ3rzhaLMIrDjNDMkNS01tea93qhEhG5xy8f6+R0l+dw3Ki+5sCoIoI942iuLZTHWogWktgwVDhU09iNEimQ==", "dev": true, "license": "MIT", "dependencies": { @@ -9646,9 +9646,9 @@ } }, "node_modules/moment": { - "version": "2.30.1", - "resolved": "https://registry.npmjs.org/moment/-/moment-2.30.1.tgz", - "integrity": "sha512-uEmtNhbDOrWPFS+hdjFCBfy9f2YoyzRpwcl+DqpC6taX21FzsTLQVbMV/W7PzNSX6x/bhC1zA3c2UQ5NzH6how==", + "version": "2.31.0", + "resolved": "https://registry.npmjs.org/moment/-/moment-2.31.0.tgz", + "integrity": "sha512-0acOTfMiWOheYS4eoWb80yYMb/JLvVv9SHbs2PehaDzfUG0Bw855SKyk0IKTnPGa5+U2bmi3W68l1+sGLX/pvw==", "license": "MIT", "engines": { "node": "*" diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 233a0e63881..3bf32d37faf 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -40,7 +40,7 @@ "dayjs": "1.11.19", "jwt-decode": "4.0.0", "lucide-react": "0.513.0", - "moment": "2.30.1", + "moment": "2.31.0", "next": "16.3.3", "next-themes": "^0.4.6", "nuqs": "^2.9.4", @@ -98,7 +98,7 @@ "overrides": { "prismjs": "1.30.0", "js-yaml": "4.3.2", - "brace-expansion": "5.0.9", + "brace-expansion": "5.0.12", "glob": "13.0.0", "minimatch": "10.2.4", "ws": "8.21.0", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx index bef938cd31c..17bb8bbfec4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx @@ -33,6 +33,12 @@ describe("AgentsTable", () => { } }); + it("right-aligns the Spend (USD) column", () => { + render(); + expect(screen.getByRole("columnheader", { name: "Spend (USD)" })).toHaveClass("text-right"); + expect(screen.getByRole("columnheader", { name: "Agent Name" })).not.toHaveClass("text-right"); + }); + it("renders the agent's model and opens the detail view when the ID cell is clicked", async () => { const user = userEvent.setup(); const onAgentClick = vi.fn(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTableColumns.tsx index a8fe3973a42..9ec1eb097d2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTableColumns.tsx @@ -90,7 +90,7 @@ export const getAgentsTableColumns = ({ { id: "spend", accessorKey: "spend", - meta: { title: "Spend (USD)" }, + meta: { title: "Spend (USD)", numeric: true }, header: ({ column }) => , size: 130, enableSorting: true, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx index 6606a4e6aaf..8d1ee100a2a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx @@ -50,6 +50,8 @@ describe("ProviderDiscountTable", () => { expect(screen.getByRole("columnheader", { name: "Provider" })).toBeInTheDocument(); expect(screen.getByRole("columnheader", { name: "Discount Percentage" })).toBeInTheDocument(); expect(screen.getByRole("columnheader", { name: "Actions" })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Discount Percentage" })).toHaveClass("text-right"); + expect(screen.getByRole("columnheader", { name: "Provider" })).not.toHaveClass("text-right"); }); it("should display provider display names in the table", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.tsx index fcc4c2af935..3d8be33fc4a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.tsx @@ -80,10 +80,11 @@ const ProviderDiscountTable: React.FC = ({ }, { header: "Discount Percentage", + numeric: true, cell: (row) => { const { displayName } = getProviderLogoAndName(row.provider); return ( -
+
{editingProvider === row.provider ? ( <> { expect(screen.getByRole("columnheader", { name: "Provider" })).toBeInTheDocument(); expect(screen.getByRole("columnheader", { name: "Margin" })).toBeInTheDocument(); expect(screen.getByRole("columnheader", { name: "Actions" })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Margin" })).toHaveClass("text-right"); + expect(screen.getByRole("columnheader", { name: "Provider" })).not.toHaveClass("text-right"); }); it("should display the provider display name", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx index 04823ac4aa0..5352695ef0a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.tsx @@ -123,10 +123,11 @@ const ProviderMarginTable: React.FC = ({ }, { header: "Margin", + numeric: true, cell: (row) => { const displayName = marginRowDisplayName(row.provider); return ( -
+
{editingProvider === row.provider ? ( <>
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts index 9210e25e1a8..597c5f7b2da 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts @@ -4,7 +4,7 @@ import { fetchMCPServers } from "@/components/networking"; import { MCPServer } from "@/components/mcp_tools/types"; import useAuthorized from "../useAuthorized"; -const mcpServersKeys = createQueryKeys("mcpServers"); +export const mcpServersKeys = createQueryKeys("mcpServers"); export const useMCPServers = (teamId?: string | null) => { const { accessToken } = useAuthorized(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts index 05025adc5e6..7d1d035b4d4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts @@ -1,4 +1,11 @@ -import { keepPreviousData, useInfiniteQuery, useQuery, useQueryClient, UseQueryResult } from "@tanstack/react-query"; +import { + keepPreviousData, + QueryClient, + useInfiniteQuery, + useQuery, + useQueryClient, + UseQueryResult, +} from "@tanstack/react-query"; import { Team } from "@/components/key_team_helpers/key_list"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { fetchTeams } from "@/app/(dashboard)/networking"; @@ -110,7 +117,7 @@ export const useTeamsTable = ( }); }; -const teamKeys = createQueryKeys("teams"); +export const teamKeys = createQueryKeys("teams"); export const useTeams = (): UseQueryResult => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ @@ -179,6 +186,11 @@ export const useTeam = (teamId?: string) => { const infiniteTeamKeys = createQueryKeys("infiniteTeams"); +export const invalidateTeamQueries = (queryClient: QueryClient) => + Promise.all( + [teamsTableKeys, teamKeys, infiniteTeamKeys].map((keys) => queryClient.invalidateQueries({ queryKey: keys.all })), + ); + export const useInfiniteTeams = (pageSize: number = 50, search?: string, organizationId?: string | null) => { const { accessToken, userId, userRole } = useAuthorized(); const isAdmin = userRole === "Admin" || userRole === "Admin Viewer"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.test.tsx index 2f7f989c099..7e1eda9143f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.test.tsx @@ -7,13 +7,21 @@ import * as networking from "@/components/networking"; import { setSecureItem } from "@/utils/secureStorage"; import { EDIT_OAUTH_UI_STATE_KEY } from "./mcp_server_edit"; import type { MCPServer } from "@/components/mcp_tools/types"; +import { mcpServersKeys } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; vi.mock(".", () => ({ MCPToolsViewer: () =>
tools viewer
, })); vi.mock("./mcp_server_edit", () => ({ - default: () =>
edit form
, + default: ({ mcpServer, onSuccess }: { mcpServer: MCPServer; onSuccess: (server: MCPServer) => void }) => ( +
+ edit form + +
+ ), EDIT_OAUTH_UI_STATE_KEY: "litellm-mcp-oauth-edit-state", })); @@ -33,9 +41,15 @@ const baseServer = { auth_type: "api_key", } as MCPServer; -const renderView = (overrides: Partial = {}, props: Record = {}) => +const newQueryClient = () => new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + +const renderView = ( + overrides: Partial = {}, + props: Record = {}, + queryClient: QueryClient = newQueryClient(), +) => render( - + { expect(await screen.findByText("edit form")).toBeInTheDocument(); }); + it("drops the cached server list and tool catalog once the edit form saves", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: Infinity } } }); + const serversKey = mcpServersKeys.list(); + const toolsKey = ["mcpTools", "srv-1", {}, null]; + const otherToolsKey = ["mcpTools", "srv-2", {}, null]; + queryClient.setQueryData(serversKey, [baseServer]); + queryClient.setQueryData(toolsKey, { tools: [] }); + queryClient.setQueryData(otherToolsKey, { tools: [] }); + const onBack = vi.fn(); + renderView({}, { onBack }, queryClient); + + await userEvent.click(screen.getByRole("tab", { name: "Settings" })); + await userEvent.click(await screen.findByRole("button", { name: "Edit Settings" })); + await userEvent.click(await screen.findByRole("button", { name: "save edit" })); + + expect(queryClient.getQueryState(serversKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(toolsKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(otherToolsKey)?.isInvalidated).toBe(false); + expect(onBack).toHaveBeenCalledTimes(1); + }); + it("opens straight into the edit form when isEditing is set", async () => { renderView({}, { isEditing: true }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx index c97596ce0f6..278a98a2fa6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx @@ -1,4 +1,6 @@ import React, { useState } from "react"; +import { useQueryClient } from "@tanstack/react-query"; +import { mcpServersKeys } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { ArrowLeft, Eye, EyeOff } from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -64,6 +66,7 @@ export const MCPServerView: React.FC = ({ }) => { // Open the editing Settings tab on first render when returning from the edit OAuth // redirect, so the "token fetched" feedback shows where the user left off (Settings=2). + const queryClient = useQueryClient(); const canEdit = isProxyAdmin && !isViewOnly && !mcpServer.is_config; const returningFromEditOAuth = isReturningFromEditOAuth(canEdit, mcpServer.server_id); const [editing, setEditing] = useState(isEditing || returningFromEditOAuth); @@ -75,6 +78,8 @@ export const MCPServerView: React.FC = ({ const canRevokeUserCredentials = userRole !== null && isProxyAdminRole(userRole) && !isViewOnly; const handleSuccess = (updated: MCPServer) => { + void queryClient.invalidateQueries({ queryKey: mcpServersKeys.all }); + void queryClient.invalidateQueries({ queryKey: ["mcpTools", updated.server_id] }); setEditing(false); onBack(); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.test.tsx index cc0169d745c..be6130b0288 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.test.tsx @@ -174,6 +174,8 @@ describe("AllModelsTable", () => { const { rerender } = render(); expect(screen.getByText("$30")).toBeInTheDocument(); expect(screen.getByText("$60")).toBeInTheDocument(); + expect(screen.getByRole("cell", { name: /\$30/ })).toHaveClass("text-right"); + expect(screen.getByRole("columnheader", { name: /costs/i })).toHaveClass("text-right"); rerender(); expect(screen.queryByText(/^\$/)).not.toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx index cbc31747688..9581d3db198 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx @@ -437,7 +437,7 @@ export const getModelsTableColumns = ({ { id: COSTS_COLUMN_ID, accessorFn: (row) => row.input_cost, - meta: { title: "Costs" }, + meta: { title: "Costs", numeric: true }, header: ({ column }) => , enableSorting: true, size: 130, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx index 889a17bc88d..15b2a30b50c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx @@ -19,7 +19,15 @@ import { } from "@/components/ui/combobox"; import { Meter, MeterIndicator, MeterTrack } from "@/components/shared/Meter"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; -import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { + NUMERIC_CELL_CLASS, + Table, + TableBody, + TableCell, + TableHead, + TableHeader, + TableRow, +} from "@/components/ui/table"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { AreaChart, BarChart, DonutChart } from "@/components/shared/charts"; @@ -651,14 +659,14 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use Provider - Spend + Spend {spendByProvider.map((provider) => ( {provider.provider} - + @@ -840,8 +848,8 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use Customer - Spend - Total Events + Spend + Total Events @@ -849,10 +857,10 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use {topUsers?.map((user: any, index: number) => ( {user.end_user} - + - {user.total_count} + {user.total_count} ))} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx index 9d163fe2c08..839eb406200 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx @@ -82,6 +82,16 @@ describe("OrganizationsTable", () => { } }); + it("right-aligns the money and count columns only", () => { + renderWithProviders(); + for (const header of ["Spend (USD)", "Budget (USD)", "Members"]) { + expect(screen.getByRole("columnheader", { name: header })).toHaveClass("text-right"); + } + for (const header of ["Organization Name", "TPM / RPM Limits"]) { + expect(screen.getByRole("columnheader", { name: header })).not.toHaveClass("text-right"); + } + }); + it("opens the detail view when the organization ID cell is clicked", async () => { const user = userEvent.setup(); const onOrganizationClick = vi.fn(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx index 0fea6c6606e..5f170a32941 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTableColumns.tsx @@ -129,7 +129,7 @@ export const getOrganizationsTableColumns = ({ { id: "spend", accessorKey: "spend", - meta: { title: "Spend (USD)" }, + meta: { title: "Spend (USD)", numeric: true }, header: ({ column }) => , size: 120, enableSorting: true, @@ -137,7 +137,7 @@ export const getOrganizationsTableColumns = ({ }, { id: "max_budget", - meta: { title: "Budget (USD)" }, + meta: { title: "Budget (USD)", numeric: true }, header: "Budget (USD)", size: 120, enableSorting: false, @@ -163,7 +163,7 @@ export const getOrganizationsTableColumns = ({ }, { id: "members", - meta: { title: "Members" }, + meta: { title: "Members", numeric: true }, header: "Members", size: 100, enableSorting: false, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx index 459e3fd8c92..67a648dacd1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/BulkEditUsers.tsx @@ -9,7 +9,16 @@ import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Checkbox } from "@/components/ui/checkbox"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { Separator } from "@/components/ui/separator"; -import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { + NUMERIC_CELL_CLASS, + Table, + TableBody, + TableCell, + TableHead, + TableHeader, + TableRow, +} from "@/components/ui/table"; +import { cn } from "@/lib/cva.config"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; interface BulkEditUserModalProps { @@ -250,7 +259,7 @@ const BulkEditUserModal: React.FC = ({ User ID Email Current Role - Budget + Budget @@ -263,7 +272,7 @@ const BulkEditUserModal: React.FC = ({ {possibleUIRoles?.[user.user_role]?.ui_label || user.user_role} - + diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx index 3a2fea66b0c..678b811921c 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.test.tsx @@ -50,6 +50,9 @@ describe("getModelHubTableColumns", () => { expect(screen.getByText("128.0K / 16.4K")).toBeInTheDocument(); expect(screen.getByText("$2.50")).toBeInTheDocument(); expect(screen.getByText("$10.00")).toBeInTheDocument(); + expect(screen.getByRole("cell", { name: "128.0K / 16.4K" })).toHaveClass("text-right"); + expect(screen.getByRole("cell", { name: /\$2\.50/ })).toHaveClass("text-right"); + expect(screen.getByRole("cell", { name: "gpt-4o" })).not.toHaveClass("text-right"); }); it("shows capability badges only for supported features", () => { diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx index 9f74771f3b1..fc71a9340ac 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTableColumns.tsx @@ -143,7 +143,7 @@ export const getModelHubTableColumns = ({ onModelClick }: ModelHubTableColumnsDe { id: "max_input_tokens", accessorKey: "max_input_tokens", - meta: { title: "Tokens", className: "hidden lg:table-cell" }, + meta: { title: "Tokens", className: "hidden lg:table-cell", numeric: true }, header: ({ column }) => , size: 110, enableSorting: true, @@ -165,7 +165,7 @@ export const getModelHubTableColumns = ({ onModelClick }: ModelHubTableColumnsDe { id: "input_cost_per_token", accessorKey: "input_cost_per_token", - meta: { title: "Cost/1M", skeleton: "twoLine" }, + meta: { title: "Cost/1M", skeleton: "twoLine", numeric: true }, header: ({ column }) => , size: 110, enableSorting: true, diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx index 9699ea2b7d1..ad89b62d611 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx @@ -161,6 +161,12 @@ describe("sort contract – only backend-sortable columns are sortable", () => { }); }); + it("right-aligns Spend / Budget but not Team", () => { + renderTable(); + expect(screen.getByRole("columnheader", { name: "Spend / Budget" })).toHaveClass("text-right"); + expect(screen.getByRole("columnheader", { name: "Team" })).not.toHaveClass("text-right"); + }); + it("does not make Spend / Budget sortable (the backend rejects sort_by=spend)", () => { renderTable(); expect(screen.queryByText("Spend / Budget").closest("button")).toBeNull(); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx index 84369a58307..05bfccf6417 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx @@ -210,7 +210,7 @@ export const getTeamTableColumns = ({ { id: "spend", accessorKey: "spend", - meta: { title: "Spend / Budget", skeleton: "meter" }, + meta: { title: "Spend / Budget", skeleton: "meter", numeric: true }, header: "Spend / Budget", size: 200, enableSorting: false, @@ -234,7 +234,7 @@ export const getTeamTableColumns = ({ }, { id: "members", - meta: { title: "Members" }, + meta: { title: "Members", numeric: true }, header: "Members", size: 110, enableSorting: false, @@ -242,7 +242,7 @@ export const getTeamTableColumns = ({ }, { id: "models", - meta: { title: "Models" }, + meta: { title: "Models", numeric: true }, header: "Models", size: 100, enableSorting: false, diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 7ef7f1fcb09..423723a5938 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -207,6 +207,12 @@ it("should render VirtualKeysTable component", () => { expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); }); +it("right-aligns the Spend / Budget column", async () => { + renderWithProviders(); + expect(await screen.findByRole("columnheader", { name: /^Spend/ })).toHaveClass("text-right"); + expect(screen.getByRole("columnheader", { name: /^Key$/ })).not.toHaveClass("text-right"); +}); + it("shows the Budget Reset column by default", async () => { renderWithProviders(); await waitFor(() => { diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index 477d7b0ecb4..d69e90b1882 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -267,7 +267,7 @@ export const getKeyTableColumns = ({ { id: "spend", accessorKey: "spend", - meta: { title: "Spend / Budget", skeleton: "meter" }, + meta: { title: "Spend / Budget", skeleton: "meter", numeric: true }, header: ({ table }) => , size: 180, enableSorting: true, diff --git a/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx b/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx index 8669f067d6d..b5029e98f9c 100644 --- a/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx +++ b/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx @@ -1,7 +1,15 @@ import React, { useState, useEffect } from "react"; import { Button, buttonVariants } from "@/components/ui/button"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; -import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { + NUMERIC_CELL_CLASS, + Table, + TableBody, + TableCell, + TableHead, + TableHeader, + TableRow, +} from "@/components/ui/table"; import { Download, FileText, FileWarning, Trash2, TriangleAlert, Upload } from "lucide-react"; import { userCreateCall, invitationCreateCall, getProxyUISettings } from "./networking"; import Papa from "papaparse"; @@ -798,7 +806,7 @@ const BulkCreateUsersButton: React.FC = ({ Email Role Teams - Budget + Budget Status @@ -809,7 +817,7 @@ const BulkCreateUsersButton: React.FC = ({ {record.user_email} {record.user_role} {record.teams} - {record.max_budget} + {record.max_budget} {renderStatusCell(record)} ))} diff --git a/ui/litellm-dashboard/src/components/chat/KeysPanel.tsx b/ui/litellm-dashboard/src/components/chat/KeysPanel.tsx index 8b8bd6c8189..fe37102eaba 100644 --- a/ui/litellm-dashboard/src/components/chat/KeysPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/KeysPanel.tsx @@ -10,7 +10,16 @@ import { Label } from "@/components/ui/label"; import { Badge } from "@/components/ui/badge"; import { Skeleton } from "@/components/ui/skeleton"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; -import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { + NUMERIC_CELL_CLASS, + Table, + TableBody, + TableCell, + TableHead, + TableHeader, + TableRow, +} from "@/components/ui/table"; +import { cn } from "@/lib/cva.config"; import { toast } from "@/lib/toast"; import { keyListCall, regenerateKeyCall } from "../networking"; import { KeyResponse } from "../key_team_helpers/key_list"; @@ -176,7 +185,9 @@ const KeysPanel: React.FC = ({ accessToken, userId, premiumUser }) => { Key - Spend + + Spend + Expires Created {premiumUser && ( @@ -220,7 +231,9 @@ const KeysPanel: React.FC = ({ accessToken, userId, premiumUser }) => { Key - Spend + + Spend + Expires Created {premiumUser && ( @@ -237,7 +250,7 @@ const KeysPanel: React.FC = ({ accessToken, userId, premiumUser }) => { {maskKey(record.key_name)} {record.key_alias &&
{record.key_alias}
} - + ${record.spend?.toFixed(2) ?? "0.00"} {record.max_budget != null && record.max_budget > 0 && ( / ${record.max_budget.toFixed(2)} diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx index a16f11eab60..418f27c8188 100644 --- a/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx @@ -213,3 +213,20 @@ describe("MemberTable actions", () => { expect(screen.getByText("No members found")).toBeInTheDocument(); }); }); + +describe("MemberTable numeric columns", () => { + it("right-aligns the header and cells of a numeric extra column only", () => { + renderTable({ + members: [MEMBERS[0]], + extraColumns: [ + { title: "Spend (USD)", key: "spend", numeric: true, render: () => $1.50 }, + { title: "Joined", key: "joined", render: () => Aug 1 }, + ], + }); + + expect(screen.getByRole("columnheader", { name: "Spend (USD)" })).toHaveClass("text-right", "tabular-nums"); + expect(screen.getByRole("cell", { name: "$1.50" })).toHaveClass("text-right", "tabular-nums"); + expect(screen.getByRole("columnheader", { name: "Joined" })).not.toHaveClass("text-right"); + expect(screen.getByRole("cell", { name: "Aug 1" })).not.toHaveClass("text-right"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx index 6efc82e2d35..3cbe59bd839 100644 --- a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx +++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx @@ -24,6 +24,7 @@ export interface MemberTableColumn { key: string; render: (member: Member) => React.ReactNode; sortValue?: (member: Member) => MemberTableSortValue; + numeric?: boolean; } export interface MemberTableProps { @@ -87,6 +88,7 @@ const extraColumnDef = (column: MemberTableColumn): ColumnDef => { header: () => {column.title}, enableSorting: false, enableGlobalFilter: false, + meta: { numeric: column.numeric }, cell: ({ row }) => column.render(row.original), }; } @@ -97,6 +99,7 @@ const extraColumnDef = (column: MemberTableColumn): ColumnDef => { sortDescFirst: false, sortUndefined: "last", enableGlobalFilter: false, + meta: { numeric: column.numeric }, cell: ({ row }) => column.render(row.original), }; }; diff --git a/ui/litellm-dashboard/src/components/common_components/simple_table.test.tsx b/ui/litellm-dashboard/src/components/common_components/simple_table.test.tsx new file mode 100644 index 00000000000..885d0ed51ac --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/simple_table.test.tsx @@ -0,0 +1,25 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; + +import { SimpleTable, type SimpleTableColumn } from "./simple_table"; + +interface Row { + name: string; + spend: number; +} + +const columns: SimpleTableColumn[] = [ + { header: "Name", accessor: "name" }, + { header: "Spend", accessor: "spend", numeric: true }, +]; + +describe("SimpleTable numeric columns", () => { + it("right-aligns the header and cells of a numeric column only", () => { + render(); + + expect(screen.getByRole("columnheader", { name: "Spend" })).toHaveClass("text-right", "tabular-nums"); + expect(screen.getByRole("cell", { name: "42" })).toHaveClass("text-right", "tabular-nums"); + expect(screen.getByRole("columnheader", { name: "Name" })).not.toHaveClass("text-right"); + expect(screen.getByRole("cell", { name: "Alice" })).not.toHaveClass("text-right"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/simple_table.tsx b/ui/litellm-dashboard/src/components/common_components/simple_table.tsx index 6a30a2e0273..a4b84d28801 100644 --- a/ui/litellm-dashboard/src/components/common_components/simple_table.tsx +++ b/ui/litellm-dashboard/src/components/common_components/simple_table.tsx @@ -1,11 +1,20 @@ import React from "react"; -import { Table, TableHeader, TableRow, TableHead, TableBody, TableCell } from "@/components/ui/table"; +import { + NUMERIC_CELL_CLASS, + Table, + TableHeader, + TableRow, + TableHead, + TableBody, + TableCell, +} from "@/components/ui/table"; export interface SimpleTableColumn { header: string; accessor?: keyof T; cell?: (row: T) => React.ReactNode; width?: string; + numeric?: boolean; } interface SimpleTableProps { @@ -34,7 +43,11 @@ export function SimpleTable({ {columns.map((column, index) => ( - + {column.header} ))} @@ -51,7 +64,7 @@ export function SimpleTable({ data.map((row, rowIndex) => ( {columns.map((column, colIndex) => ( - + {column.cell ? column.cell(row) : String(row[column.accessor as keyof T] ?? "")} ))} diff --git a/ui/litellm-dashboard/src/components/organization/organization_view.tsx b/ui/litellm-dashboard/src/components/organization/organization_view.tsx index c800d12ad62..f325e92d2b5 100644 --- a/ui/litellm-dashboard/src/components/organization/organization_view.tsx +++ b/ui/litellm-dashboard/src/components/organization/organization_view.tsx @@ -134,6 +134,7 @@ const OrganizationInfoView: React.FC = ({ { title: "Spend (USD)", key: "spend", + numeric: true, sortValue: (record: Member) => orgMemberFor(record)?.spend ?? null, render: (record: Member) => , }, diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index a7e4befa6ac..df88287b369 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -141,8 +141,33 @@ const expansionColumns: ColumnDef[] = [ }, ]; +const numericColumns: ColumnDef[] = [ + { + accessorKey: "name", + header: "Name", + cell: ({ row }) => {row.original.name}, + }, + { + id: "spend", + header: ({ column }) => , + meta: { numeric: true }, + cell: () => $1.50, + }, +]; + const CHARLIE_ALICE_BOB: Person[] = [person("c", "Charlie"), person("a", "Alice"), person("b", "Bob")]; +describe("DataTable numeric columns", () => { + it("right-aligns the header and cells of a numeric column only", () => { + render(); + + expect(screen.getByRole("columnheader", { name: "Spend" })).toHaveClass("text-right", "tabular-nums"); + expect(screen.getByRole("cell", { name: "$1.50" })).toHaveClass("text-right", "tabular-nums"); + expect(screen.getByRole("columnheader", { name: "Name" })).not.toHaveClass("text-right"); + expect(screen.getByRole("cell", { name: "Alice" })).not.toHaveClass("text-right"); + }); +}); + describe("DataTable sorting", () => { it("client mode reorders rows when the sort header is clicked", async () => { const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index 340f8d4f44f..feb1615b4e1 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -31,6 +31,7 @@ import { Fragment, useEffect, useState } from "react"; import { Skeleton } from "@/components/ui/skeleton"; import { + NUMERIC_CELL_CLASS, Table as TableRoot, TableBody, TableCell, @@ -193,7 +194,7 @@ function DataTableHeadCell({ header, size, stickyHeader, enableColumnResi className={cn( "relative text-muted-foreground", size === "compact" ? "h-8 px-2 py-1 text-xs" : "", - meta?.numeric ? "text-right" : "", + meta?.numeric ? NUMERIC_CELL_CLASS : "", meta?.className, meta?.headerClassName, sticky.className, @@ -238,7 +239,7 @@ function DataTableBodyCell({ cell, size, stickyHeader, enableColumnResizi className={cn( "overflow-hidden text-ellipsis", size === "compact" ? "px-2 py-1 text-xs" : "", - meta?.numeric ? "text-right tabular-nums" : "", + meta?.numeric ? NUMERIC_CELL_CLASS : "", meta?.className, sticky.className, )} diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx index 2d84d78d56c..456dc91b13f 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx @@ -77,7 +77,8 @@ vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAllProxyModels: vi.fn(), })); -vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", async (importOriginal) => ({ + ...(await importOriginal()), useTeam: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index a693ee971d4..4650f4b6987 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -83,7 +83,8 @@ vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAllProxyModels: vi.fn(), })); -vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", async (importOriginal) => ({ + ...(await importOriginal()), useTeam: vi.fn(), })); @@ -233,7 +234,7 @@ vi.mock("../key_team_helpers/filter_helpers", () => ({ import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; -import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { teamKeys, teamsTableKeys, useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; @@ -1146,6 +1147,26 @@ describe("TeamInfoView", () => { }); }); + it("invalidates the cached team list and team detail queries after saving team settings", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["gpt-4"] })); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + const tableKey = teamsTableKeys.list({ page: 1, limit: 10 }); + const detailKey = teamKeys.detail("123"); + testQueryClient.setQueryData(tableKey, { teams: [], total: 0 }); + testQueryClient.setQueryData(detailKey, { team_id: "123" }); + + renderWithProviders(); + + await user.click(await screen.findByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await screen.findByLabelText("Team Name"); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(testQueryClient.getQueryState(tableKey)?.isInvalidated).toBe(true)); + expect(testQueryClient.getQueryState(detailKey)?.isInvalidated).toBe(true); + }); + const openSettingsEditorForTeam = async ( user: ReturnType, teamOverrides: Record, diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 3845f94593d..5dfcf1d8e35 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -2,6 +2,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import type { components } from "@/lib/http/schema"; import useCan from "@/app/(dashboard)/hooks/useCan"; import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { invalidateTeamQueries } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useQueryClient } from "@tanstack/react-query"; import UserSearchModal from "@/components/common_components/user_search_modal"; import { @@ -915,7 +916,8 @@ const TeamInfoView: React.FC = ({ const persistTeamUpdate = async (token: string, updateData: Record) => { await teamUpdateCall(token, updateData); - queryClient.invalidateQueries({ queryKey: organizationKeys.all }); + void queryClient.invalidateQueries({ queryKey: organizationKeys.all }); + void invalidateTeamQueries(queryClient); setIsEditing(false); }; diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index bd7398bc6c8..1d50a9a4670 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -292,6 +292,9 @@ describe("TeamMembersComponent", () => { expect(screen.getByText("$100.50")).toBeInTheDocument(); expect(screen.getByText("$1,538.26")).toBeInTheDocument(); + expect(screen.getByRole("cell", { name: "$100.50" })).toHaveClass("text-right"); + expect(screen.getByRole("columnheader", { name: /^Team Member Budget \(USD\)/ })).toHaveClass("text-right"); + expect(screen.getByRole("columnheader", { name: "User Email" })).not.toHaveClass("text-right"); expect(screen.getByText(/100 RPM/)).toBeInTheDocument(); expect(screen.getByText(/10000 TPM/)).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index e8576ce6dfe..a24d1b1e7cd 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -186,6 +186,7 @@ export default function TeamMemberTab({ ), key: "spend", + numeric: true, sortValue: (record: Member) => getUserCurrentCycleSpend(record.user_id), render: (record: Member) => , }, @@ -199,6 +200,7 @@ export default function TeamMemberTab({ ), key: "total_spend", + numeric: true, sortValue: (record: Member) => getUserTotalSpend(record.user_id), render: (record: Member) => , }, @@ -212,11 +214,12 @@ export default function TeamMemberTab({ ), key: "budget", + numeric: true, sortValue: (record: Member) => getUserBudget(record.user_id), render: (record: Member) => { const source = getUserBudgetSource(record.user_id); return ( - + {source !== "none" && ( diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx index 08408fb4ff4..7561264ea65 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx @@ -131,6 +131,14 @@ describe("TeamVirtualKeysTable", () => { }); }); + it("right-aligns the Spend (USD) and Budget (USD) columns", async () => { + renderWithProviders(); + + expect(await screen.findByRole("columnheader", { name: "Spend (USD)" })).toHaveClass("text-right"); + expect(screen.getByRole("columnheader", { name: "Budget (USD)" })).toHaveClass("text-right"); + expect(screen.getByRole("columnheader", { name: "Key ID" })).not.toHaveClass("text-right"); + }); + it("should display keys in table when data is loaded", async () => { mockUseKeys.mockReturnValue({ data: { diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index 5b1b71e060e..4e4fa3f4bb2 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -285,7 +285,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi { id: "spend", accessorKey: "spend", - meta: { title: "Spend (USD)" }, + meta: { title: "Spend (USD)", numeric: true }, header: ({ column }) => , size: 100, enableSorting: true, @@ -294,7 +294,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi { id: "max_budget", accessorKey: "max_budget", - meta: { title: "Budget (USD)" }, + meta: { title: "Budget (USD)", numeric: true }, header: ({ column }) => , size: 110, enableSorting: true, diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 522af5a85ad..09042dea930 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -7,10 +7,11 @@ vi.mock("@/lib/toast", () => ({ })); // ---- Hoisted shared mocks (safe to use inside vi.mock factories) ---- -const { keyUpdateCallMock, keyDeleteCallMock, mockUseAuthorized } = vi.hoisted(() => { +const { keyUpdateCallMock, keyDeleteCallMock, invalidateQueriesMock, mockUseAuthorized } = vi.hoisted(() => { return { keyUpdateCallMock: vi.fn().mockResolvedValue({}), keyDeleteCallMock: vi.fn().mockResolvedValue({}), + invalidateQueriesMock: vi.fn().mockResolvedValue(undefined), mockUseAuthorized: vi.fn(), }; }); @@ -170,7 +171,7 @@ vi.mock("@tanstack/react-query", async (importOriginal) => { const actual = await importOriginal(); return { ...actual, - useQueryClient: () => ({ invalidateQueries: vi.fn() }), + useQueryClient: () => ({ invalidateQueries: invalidateQueriesMock }), }; }); @@ -366,6 +367,24 @@ describe("KeyInfoView handleKeyUpdate mcp_toolsets", () => { }); }); +describe("KeyInfoView handleKeyUpdate cache sync", () => { + it("should invalidate every cached key query so the list and detail views re-read the saved key", async () => { + keyUpdateCallMock.mockResolvedValueOnce({ + object_permission: { mcp_servers: ["srv-1"], mcp_tool_permissions: { "srv-1": ["read_wiki"] } }, + }); + renderView(true); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { token: "tok_123", metadata: {} }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Key updated successfully")); + expect(invalidateQueriesMock).toHaveBeenCalledWith({ queryKey: ["keys"] }); + }); +}); + describe("KeyInfoView handleKeyUpdate skills", () => { it("should forward the skills the edit form supplies into object_permission and drop the form key", async () => { renderView(true); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 7eb09926caf..cfb5e9fa1f8 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -381,8 +381,8 @@ export default function KeyInfoView({ const newKeyValues = await keyUpdateCall(accessToken, formValues); - // Update local state setCurrentKeyData((prevData) => (prevData ? { ...prevData, ...newKeyValues } : undefined)); + void queryClient.invalidateQueries({ queryKey: keyKeys.all }); if (onKeyDataUpdate) { onKeyDataUpdate(newKeyValues); diff --git a/ui/litellm-dashboard/src/components/ui/table.tsx b/ui/litellm-dashboard/src/components/ui/table.tsx index 6271a9e89ac..1c4c1a981de 100644 --- a/ui/litellm-dashboard/src/components/ui/table.tsx +++ b/ui/litellm-dashboard/src/components/ui/table.tsx @@ -4,6 +4,8 @@ import * as React from "react"; import { cn } from "@/lib/cva.config"; +const NUMERIC_CELL_CLASS = "text-right tabular-nums"; + const Table = React.forwardRef>( ({ className, ...props }, ref) => (
@@ -96,4 +98,4 @@ const TableCaption = React.forwardRef { expect(within(row).getByText("gpt-4o")).toBeInTheDocument(); }); + it("shows the x-litellm-call-id with its own copy button when it differs from the request id", () => { + renderHeader(logEntry({ request_id: "chatcmpl-h", litellm_call_id: "call-h" }), false); + + expect(screen.getByText("chatcmpl-h")).toBeInTheDocument(); + expect(screen.getByText("call-h")).toBeInTheDocument(); + expect(screen.getByText("x-litellm-call-id:")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Copy Request ID" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Copy x-litellm-call-id" })).toBeInTheDocument(); + }); + + it("omits the x-litellm-call-id line and button when the ids match", () => { + renderHeader(logEntry({ request_id: "same-h", litellm_call_id: "same-h" }), false); + + expect(screen.queryByText("x-litellm-call-id:")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Copy x-litellm-call-id" })).not.toBeInTheDocument(); + }); + it("falls back to the request id row when the log names no model", () => { renderHeader(logEntry({ model: "", custom_llm_provider: "" }), true); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx index 65b5801602c..a4afdb68fae 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx @@ -50,6 +50,8 @@ export function DrawerHeader({ const providerInfo = provider ? getProviderLogoAndName(provider) : null; const showToggleWithProvider = isSidebarCollapsed && Boolean(providerInfo || log.model); const showToggleWithRequestId = isSidebarCollapsed && !showToggleWithProvider; + const callId: string | null = + log.litellm_call_id && log.litellm_call_id !== log.request_id ? log.litellm_call_id : null; return (
{showToggleWithRequestId && } - +
+ + {callId && ( +
+ + x-litellm-call-id: + + +
+ )} +
@@ -140,15 +155,22 @@ function ModelProviderSection({ ); } -/** - * Request ID display with copy functionality - */ -function RequestIdSection({ requestId }: { requestId: string }) { +function CopyableId({ + value, + label, + fontSize, + muted, +}: { + value: string; + label: string; + fontSize: number; + muted?: boolean; +}) { const [copied, setCopied] = useState(false); const handleCopy = async () => { try { - await navigator.clipboard.writeText(requestId); + await navigator.clipboard.writeText(value); setCopied(true); setTimeout(() => setCopied(false), 1200); } catch { @@ -157,38 +179,36 @@ function RequestIdSection({ requestId }: { requestId: string }) { }; return ( -
- - - - } + + + + } + > + {value} + - - {requestId} - - -
+ {copied ? : } + + + {value} + + ); } diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index 4d451a9c7b9..68590b6de2d 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -1,4 +1,4 @@ -import { render, screen } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; @@ -7,6 +7,13 @@ import { DataTable } from "@/components/shared/DataTable"; import type { LogEntry } from "./columns"; import { getRequestLogsTableColumns } from "./RequestLogsTableColumns"; +const { copyToClipboardMock } = vi.hoisted(() => ({ copyToClipboardMock: vi.fn() })); + +vi.mock("@/utils/dataUtils", async (importOriginal) => ({ + ...(await importOriginal()), + copyToClipboard: copyToClipboardMock, +})); + const logEntry = (overrides: Partial): LogEntry => ({ request_id: "req-1", api_key: "key-1", @@ -267,13 +274,64 @@ describe("batch rows", () => { }); it("leaves ordinary request ids untouched", () => { - renderRows([logEntry({ request_id: "chatcmpl-42" })]); + renderRows([logEntry({ request_id: "chatcmpl-42", litellm_call_id: "chatcmpl-42" })]); expect(screen.getByText("chatcmpl-42")).toBeInTheDocument(); expect(screen.queryByText("batch cost")).not.toBeInTheDocument(); }); }); +describe("Request ID column", () => { + it("shows only the request id in the cell and the x-litellm-call-id in its tooltip when they differ", async () => { + const user = userEvent.setup(); + renderRows([logEntry({ request_id: "chatcmpl-9", litellm_call_id: "call-uuid-9" })]); + + expect(screen.getByText("chatcmpl-9")).toBeInTheDocument(); + expect(screen.queryByText("call-uuid-9")).not.toBeInTheDocument(); + + await user.hover(screen.getByText("chatcmpl-9")); + expect(await screen.findByText("x-litellm-call-id: call-uuid-9")).toBeInTheDocument(); + }); + + it("copies the x-litellm-call-id from the tooltip without opening the row", async () => { + const user = userEvent.setup(); + const onRowClick = vi.fn(); + render( + row.request_id} + size="compact" + onRowClick={onRowClick} + />, + ); + + await user.hover(screen.getByText("chatcmpl-9")); + fireEvent.click(await screen.findByRole("button", { name: "Copy x-litellm-call-id" })); + + expect(copyToClipboardMock).toHaveBeenCalledWith("call-uuid-9"); + expect(onRowClick).not.toHaveBeenCalled(); + }); + + it("keeps the plain id tooltip when request id and call id are the same", async () => { + const user = userEvent.setup(); + renderRows([logEntry({ request_id: "same-id-7", litellm_call_id: "same-id-7" })]); + + await user.hover(screen.getByText("same-id-7")); + await waitFor(() => expect(screen.getAllByText("same-id-7")).toHaveLength(2)); + expect(screen.queryByText(/x-litellm-call-id/)).not.toBeInTheDocument(); + }); + + it("keeps the plain id tooltip when the row carries no call id", async () => { + const user = userEvent.setup(); + renderRows([logEntry({ request_id: "chatcmpl-no-call", litellm_call_id: null })]); + + await user.hover(screen.getByText("chatcmpl-no-call")); + await waitFor(() => expect(screen.getAllByText("chatcmpl-no-call")).toHaveLength(2)); + expect(screen.queryByText(/x-litellm-call-id/)).not.toBeInTheDocument(); + }); +}); + describe("Model column", () => { it("lists every model used across a conversation, not only the representative call's model", () => { const conversationCall: Partial = { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index 5705f41f3de..df7f55d7d76 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -1,10 +1,11 @@ "use client"; import type { ColumnDef } from "@tanstack/react-table"; +import { Copy } from "lucide-react"; import { DataTableSortHeader } from "@/components/shared/DataTable"; import { CellTooltip, DateCell, IdCell, MoneyCell, StatusBadge } from "@/components/shared/table_cells"; -import { getSpendString } from "@/utils/dataUtils"; +import { copyToClipboard, getSpendString } from "@/utils/dataUtils"; import { getProviderLogoAndName } from "../provider_info_helpers"; import { getBatchIdFromRequestId, getBatchRequestCounts, isBatchCallType } from "./batchLogUtils"; @@ -30,6 +31,28 @@ const readMcpLogoUrl = (metadata: Record | undefined): string | return typeof url === "string" && url !== "" ? url : undefined; }; +function RequestIdWithCallIdTooltip({ requestId, callId }: { requestId: string; callId: string }) { + return ( + + {requestId} + + x-litellm-call-id: {callId} + + + + ); +} + const getLogoUrl = (row: LogEntry, provider: string): string => readMcpLogoUrl(row.metadata) ?? (provider ? getProviderLogoAndName(provider).logo : ""); @@ -160,7 +183,14 @@ export const getRequestLogsTableColumns = ({
); } - return ; + const callId = log.litellm_call_id && log.litellm_call_id !== log.request_id ? log.litellm_call_id : null; + return ( + : undefined} + /> + ); }, }, { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ecfabf33e87..cc0435c3d5e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24285,6 +24285,45 @@ export interface components { /** Uri */ uri?: string; }; + /** AgentIdentityBinding */ + AgentIdentityBinding: { + /** + * Active + * @default true + */ + active: boolean; + /** Agent Id */ + agent_id: string; + /** Client Id */ + client_id: string; + /** Issuer */ + issuer: string; + /** Last Authenticated At */ + last_authenticated_at?: string | null; + /** + * Provider + * @constant + */ + provider: "microsoft_entra"; + /** + * Required Roles + * @default [] + */ + required_roles: string[]; + /** + * Required Scopes + * @default [ + * "user_impersonation" + * ] + */ + required_scopes: string[]; + /** Revision */ + revision: string; + /** Service Principal Id */ + service_principal_id?: string | null; + /** Tenant Id */ + tenant_id: string; + }; /** * AgentInterface * @description Declares a combination of a target URL and a transport protocol. @@ -24440,8 +24479,30 @@ export interface components { created_at?: string | null; /** Created By */ created_by?: string | null; + /** + * Enabled + * @default true + */ + enabled: boolean; + /** + * Execution Mode + * @default autonomous + * @enum {string} + */ + execution_mode: "autonomous" | "delegated" | "both"; /** Extra Headers */ extra_headers?: string[] | null; + identity?: components["schemas"]["AgentIdentityBinding"] | null; + /** + * Identity Managed + * @default false + */ + identity_managed: boolean; + /** + * Jwt Auth Configured + * @default false + */ + jwt_auth_configured: boolean; /** Keys */ keys?: components["schemas"]["AgentKeySummary"][] | null; kill_switch?: components["schemas"]["AgentKillSwitchConfig"] | null; @@ -32969,6 +33030,8 @@ export interface components { cache_creation_input_token_cost_above_272k_tokens_flex?: number | null; /** Cache Creation Input Token Cost Above 272K Tokens Priority */ cache_creation_input_token_cost_above_272k_tokens_priority?: number | null; + /** Cache Creation Input Token Cost Above 272K Tokens Ultrafast */ + cache_creation_input_token_cost_above_272k_tokens_ultrafast?: number | null; /** Cache Creation Input Token Cost Batches */ cache_creation_input_token_cost_batches?: number | null; /** Cache Creation Input Token Cost Flex */ @@ -32997,6 +33060,8 @@ export interface components { cache_read_input_token_cost_above_272k_tokens_flex?: number | null; /** Cache Read Input Token Cost Above 272K Tokens Priority */ cache_read_input_token_cost_above_272k_tokens_priority?: number | null; + /** Cache Read Input Token Cost Above 272K Tokens Ultrafast */ + cache_read_input_token_cost_above_272k_tokens_ultrafast?: number | null; /** Cache Read Input Token Cost Above 512K Tokens */ cache_read_input_token_cost_above_512k_tokens?: number | null; /** Cache Read Input Token Cost Balanced */ @@ -33081,6 +33146,8 @@ export interface components { input_cost_per_token_above_272k_tokens_flex?: number | null; /** Input Cost Per Token Above 272K Tokens Priority */ input_cost_per_token_above_272k_tokens_priority?: number | null; + /** Input Cost Per Token Above 272K Tokens Ultrafast */ + input_cost_per_token_above_272k_tokens_ultrafast?: number | null; /** Input Cost Per Token Above 512K Tokens */ input_cost_per_token_above_512k_tokens?: number | null; /** Input Cost Per Token Balanced */ @@ -33208,6 +33275,8 @@ export interface components { output_cost_per_token_above_272k_tokens_flex?: number | null; /** Output Cost Per Token Above 272K Tokens Priority */ output_cost_per_token_above_272k_tokens_priority?: number | null; + /** Output Cost Per Token Above 272K Tokens Ultrafast */ + output_cost_per_token_above_272k_tokens_ultrafast?: number | null; /** Output Cost Per Token Above 512K Tokens */ output_cost_per_token_above_512k_tokens?: number | null; /** Output Cost Per Token Balanced */ @@ -46800,6 +46869,8 @@ export interface components { cache_creation_input_token_cost_above_272k_tokens_flex?: number | null; /** Cache Creation Input Token Cost Above 272K Tokens Priority */ cache_creation_input_token_cost_above_272k_tokens_priority?: number | null; + /** Cache Creation Input Token Cost Above 272K Tokens Ultrafast */ + cache_creation_input_token_cost_above_272k_tokens_ultrafast?: number | null; /** Cache Creation Input Token Cost Batches */ cache_creation_input_token_cost_batches?: number | null; /** Cache Creation Input Token Cost Flex */ @@ -46828,6 +46899,8 @@ export interface components { cache_read_input_token_cost_above_272k_tokens_flex?: number | null; /** Cache Read Input Token Cost Above 272K Tokens Priority */ cache_read_input_token_cost_above_272k_tokens_priority?: number | null; + /** Cache Read Input Token Cost Above 272K Tokens Ultrafast */ + cache_read_input_token_cost_above_272k_tokens_ultrafast?: number | null; /** Cache Read Input Token Cost Above 512K Tokens */ cache_read_input_token_cost_above_512k_tokens?: number | null; /** Cache Read Input Token Cost Balanced */ @@ -46912,6 +46985,8 @@ export interface components { input_cost_per_token_above_272k_tokens_flex?: number | null; /** Input Cost Per Token Above 272K Tokens Priority */ input_cost_per_token_above_272k_tokens_priority?: number | null; + /** Input Cost Per Token Above 272K Tokens Ultrafast */ + input_cost_per_token_above_272k_tokens_ultrafast?: number | null; /** Input Cost Per Token Above 512K Tokens */ input_cost_per_token_above_512k_tokens?: number | null; /** Input Cost Per Token Balanced */ @@ -47039,6 +47114,8 @@ export interface components { output_cost_per_token_above_272k_tokens_flex?: number | null; /** Output Cost Per Token Above 272K Tokens Priority */ output_cost_per_token_above_272k_tokens_priority?: number | null; + /** Output Cost Per Token Above 272K Tokens Ultrafast */ + output_cost_per_token_above_272k_tokens_ultrafast?: number | null; /** Output Cost Per Token Above 512K Tokens */ output_cost_per_token_above_512k_tokens?: number | null; /** Output Cost Per Token Balanced */ @@ -66054,7 +66131,7 @@ export interface operations { group_by_session?: boolean; /** @description Keyset cursor '||' from a previous group_by_session page. UI route only, honored when sorting by startTime */ session_cursor?: string | null; - /** @description Match a log whose request_id, api_key (hash), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ + /** @description Match a log whose request_id, litellm_call_id, api_key (hash), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ search?: string | null; }; header?: never; @@ -66174,7 +66251,7 @@ export interface operations { group_by_session?: boolean; /** @description Keyset cursor '||' from a previous group_by_session page. UI route only, honored when sorting by startTime */ session_cursor?: string | null; - /** @description Match a log whose request_id, api_key (hash), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ + /** @description Match a log whose request_id, litellm_call_id, api_key (hash), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ search?: string | null; }; header?: never; diff --git a/uv.lock b/uv.lock index 527f53bd372..4b9dbaba39b 100644 --- a/uv.lock +++ b/uv.lock @@ -4500,7 +4500,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.104.0" +version = "1.105.0" source = { editable = "." } dependencies = [ { name = "aiohttp" }, @@ -4962,12 +4962,12 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.71" +version = "0.1.72" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.102" +version = "0.4.103" source = { editable = "litellm-proxy-extras" } [[package]] @@ -7857,14 +7857,14 @@ wheels = [ [[package]] name = "pyjwt" -version = "2.13.0" +version = "2.14.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3b/81/58d0ac84e1ef3a3843791d6954d94c0b33d526c75eeb1efbce9d0a4c4077/pyjwt-2.13.0.tar.gz", hash = "sha256:41571c89ca91598c79e8ef18a2d07367d4810fbbd6f637794879baf1b7703423", size = 107515, upload-time = "2026-05-21T19:54:36.618Z" } +sdist = { url = "https://files.pythonhosted.org/packages/af/c3/8a3b59c25070cc61dc517fbdfa5dc0904670c96f605cc69759dc09166b99/pyjwt-2.14.0.tar.gz", hash = "sha256:77283c83fb56ecf566a886c757a714bc83668e38156de2cce8263302f42e0b86", size = 113177, upload-time = "2026-09-11T13:11:54.638Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a3/5e/ecf12fdb62546d64385c158514e9b2b671f7832108ef2ecd2020ce0af2d1/pyjwt-2.13.0-py3-none-any.whl", hash = "sha256:66adcc2aff09b3f1bbd95fc1e1577df8ac8723c978552fd43304c8a290ac5728", size = 31274, upload-time = "2026-05-21T19:54:35.362Z" }, + { url = "https://files.pythonhosted.org/packages/9c/97/672cb32ce0dfea44b740cb7b4f97038463b9cf7c0ead1aacf595572851d6/pyjwt-2.14.0-py3-none-any.whl", hash = "sha256:ad0cef71c756a56e74863c2919cf0985f72decbcfcb550ee2f422e7c62b5eedc", size = 32896, upload-time = "2026-09-11T13:11:53.409Z" }, ] [package.optional-dependencies]