Merge branch 'main' into feat/openguardrails-guardrail-integration

This commit is contained in:
Thomas 2026-03-21 19:50:03 +08:00 committed by GitHub
commit 3bf83084ca
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
154 changed files with 9450 additions and 2508 deletions

View file

@ -42,7 +42,7 @@ commands:
"pydantic==2.11.0" "mcp==1.25.0" "requests-mock>=1.12.1" \
"responses==0.25.7" "pytest-xdist==3.6.1" "pytest-timeout==2.2.0" \
"pytest-cov==5.0.0" "semantic_router==0.1.10" "fastapi-offline==1.7.3" \
"a2a"
"a2a" "parameterized>=0.9.0"
- setup_litellm_enterprise_pip
- save_cache:
paths:
@ -1115,7 +1115,7 @@ jobs:
for dir in "${IGNORE_DIRS[@]}"; do
IGNORE_ARGS="$IGNORE_ARGS --ignore=$dir"
done
python -m pytest -v tests/llm_translation $IGNORE_ARGS --junitxml=test-results/junit.xml --durations=20 -n 8 --timeout=120 --timeout_method=thread
python -m pytest -v tests/llm_translation $IGNORE_ARGS --junitxml=test-results/junit.xml --durations=20 -n 8 --timeout=120 --timeout_method=thread --retries 2 --retry-delay 5
no_output_timeout: 15m
# Store test results
@ -1331,7 +1331,7 @@ jobs:
command: |
pwd
ls
python -m pytest -vv tests/unified_google_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5
python -m pytest -vv tests/unified_google_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 --retries 3 --retry-delay 5
no_output_timeout: 15m
- run:
name: Rename the coverage files

View file

@ -28,9 +28,12 @@ jobs:
find . -type d -name "__pycache__" -exec rm -rf {} + || true
find . -name "*.pyc" -delete || true
- name: Check poetry.lock is up to date
run: |
poetry check --lock || (echo "❌ poetry.lock is out of sync with pyproject.toml. Run 'poetry lock' locally and commit the result." && exit 1)
- name: Install dependencies
run: |
poetry lock
poetry install --with dev
- name: Check Black formatting

View file

@ -14,12 +14,12 @@ repos:
types: [python]
files: (litellm/|litellm_proxy_extras/|enterprise/).*\.py
exclude: ^litellm/__init__.py$
# - id: black
# name: black
# entry: poetry run black
# language: system
# types: [python]
# files: (litellm/|litellm_proxy_extras/|enterprise/).*\.py
- id: black
name: black
entry: poetry run black
language: system
types: [python]
files: (litellm/|litellm_proxy_extras/).*\.py
- repo: https://github.com/pycqa/flake8
rev: 7.0.0 # The version of flake8 to use
hooks:

View file

@ -163,6 +163,9 @@ run_grype_scans() {
"CVE-2026-25639" # axios - full fix requires 1.x major version bump; pinned to >=0.30.2 to clear other axios CVEs, upgrade to 1.x in follow-up
"CVE-2026-2297" # Python 3.13 SourcelessFileLoader audit hook bypass - no fix available in base image
"GHSA-qffp-2rhf-9h96" # tar hardlink path traversal - from nodejs_wheel bundled npm, not used in application runtime code
"CVE-2026-2673" # OpenSSL 3.6.1 TLS 1.3 key exchange group negotiation issue - no fix available yet
"CVE-2026-3644" # Python 3.13 vulnerability - no fix available in base image
"CVE-2026-4224" # Python 3.13 Expat parser stack overflow in ElementDeclHandler - no fix available in base image
)
# Build JSON array of allowlisted CVE IDs for jq

View file

@ -0,0 +1,78 @@
---
slug: guardrail-logging-secret-exposure-incident
title: "Incident Report: Guardrail logging exposed secret headers in spend logs and traces"
date: 2026-03-18T10:00:00
authors:
- litellm
tags: [incident-report, security, guardrails]
hide_table_of_contents: false
---
**Date:** March 18, 2026
**Duration:** Unknown
**Severity:** High
**Status:** Resolved
## Summary
When a custom guardrail returned the full LiteLLM request/data dictionary, the guardrail response logged by LiteLLM could include `secret_fields.raw_headers`, including plaintext `Authorization` headers containing API keys or other credentials.
This information could then propagate to logging and observability surfaces that consume guardrail metadata, including:
- **Spend logs in the LiteLLM UI:** visible to admins with access to spend-log data
- **OpenTelemetry traces:** visible to anyone with access to the relevant telemetry backend
LLM calls, proxy routing, and provider execution were not blocked by this bug. The impact was exposure of sensitive request headers in observability and logging paths.
{/* truncate */}
---
## Background
LiteLLM keeps internal request data (including request headers) for use during the call. That data is not meant to be written to logs or telemetry.
When custom guardrails run, their outcomes are logged so they can appear in spend logs, OpenTelemetry traces, and other observability backends. If a guardrail returned the full request payload instead of a minimal result, that internal request data could be included in what was logged. Before the fix, the guardrail logging path did not strip that data before sending it to those systems.
```mermaid
flowchart TD
inboundRequest["1. Incoming proxy request"] --> storeSecrets["2. Store internal request data"]
storeSecrets --> guardrailRuns["3. Custom guardrail runs"]
guardrailRuns --> fullDataReturn["4. Guardrail returns full request payload"]
fullDataReturn --> loggingBuild["5. Build guardrail log payload"]
loggingBuild --> spendLogs["6a. Persist to spend logs / UI"]
loggingBuild --> otelTraces["6b. Attach to OTEL guardrail spans"]
```
---
## Root Cause
The root cause was incomplete sanitization in the guardrail logging path. When building the payload that gets sent to spend logs and traces, LiteLLM prepared guardrail responses for logging but did not strip internal request data (such as headers) from them. If a guardrail returned a response that included that data, it was passed through to the logging and observability systems unchanged.
---
## Impact
This issue required all of the following:
1. A custom guardrail returned the full LiteLLM request/data dictionary, or another response object containing `secret_fields`.
2. LiteLLM logged that guardrail response through the standard guardrail logging path.
3. An operator, admin, or telemetry consumer had access to the resulting logs or traces.
When those conditions were met, sensitive values could become visible through:
- **Spend logs / UI responses:** guardrail metadata could be included in spend-log payloads rendered in the admin UI.
- **OpenTelemetry traces:** `guardrail_response` could be written as a span attribute on guardrail spans.
- **Other downstream observability backends:** any integration consuming the same guardrail metadata could receive the leaked values.
This was a logging and telemetry exposure bug. It did not let callers bypass auth, access other tenants directly, or change model behavior, but it could expose plaintext credentials to people with access to those observability systems.
---
## Guidance For Users
- Upgrade to LiteLLM 1.82.3+.
- If you operated custom guardrails that return the full request/data dict, review whether spend logs or telemetry traces were retained during the affected period.
- Rotate any credentials that may have appeared in `Authorization` or other forwarded request headers in those systems.
- Apply least-privilege access controls to spend-log views and telemetry backends that may contain request-derived metadata.

View file

@ -401,8 +401,10 @@ router_settings:
| AUTH_STRATEGY | Strategy used for authentication (e.g., OAuth, API key)
| AUTO_REDIRECT_UI_LOGIN_TO_SSO | Flag to enable automatic redirect of UI login page to SSO when SSO is configured. Default is **false**
| AUDIO_SPEECH_CHUNK_SIZE | Chunk size for audio speech processing. Default is 1024
| ANTHROPIC_API_KEY | API key for Anthropic service
| ANTHROPIC_API_KEY | API key for Anthropic service. Uses `x-api-key` header for authentication.
| ANTHROPIC_AUTH_TOKEN | Alternative auth token for Anthropic service. Uses `Authorization: Bearer` header instead of `x-api-key`. Used as fallback when `ANTHROPIC_API_KEY` is not set.
| ANTHROPIC_API_BASE | Base URL for Anthropic API. Default is https://api.anthropic.com
| ANTHROPIC_BASE_URL | Alternative to `ANTHROPIC_API_BASE` for setting the Anthropic API base URL. Used as fallback when `ANTHROPIC_API_BASE` is not set.
| ANTHROPIC_TOKEN_COUNTING_BETA_VERSION | Beta version header for Anthropic token counting API. Default is `token-counting-2024-11-01`
| AWS_ACCESS_KEY_ID | Access Key ID for AWS services
| AWS_BATCH_ROLE_ARN | ARN of the AWS IAM role for batch operations
@ -902,6 +904,7 @@ router_settings:
| OTEL_SERVICE_NAME | Service name identifier for OpenTelemetry
| OTEL_TRACER_NAME | Tracer name for OpenTelemetry tracing
| OTEL_LOGS_EXPORTER | Exporter type for OpenTelemetry logs (e.g., console)
| OTEL_IGNORE_CONTEXT_PROPAGATION | When true, ignore parent span context propagation in OpenTelemetry callbacks
| PAGERDUTY_API_KEY | API key for PagerDuty Alerting
| PANW_PRISMA_AIRS_API_KEY | API key for PANW Prisma AIRS service
| PANW_PRISMA_AIRS_API_BASE | Base URL for PANW Prisma AIRS service

View file

@ -602,6 +602,22 @@ Since you shouldn't use 12.5, round down to **10** to leave a safety buffer. Thi
- Total maximum connections: 8 workers × 10 connections = 80 connections
- This stays safely under your database's 100 connection limit
## LiteLLM License Key (Enterprise)
To enable [LiteLLM Enterprise features](https://docs.litellm.ai/docs/proxy/enterprise), set your license key as an environment variable:
```bash
export LITELLM_LICENSE="eyJ..."
```
The license key is a JWT token provided when you purchase a LiteLLM Enterprise license. Once set, LiteLLM will automatically detect and activate enterprise features.
You can also add it to your `.env` file:
```env
LITELLM_LICENSE="eyJ..."
```
## Extras

View file

@ -47,6 +47,10 @@ pip install litellm==1.82.3
- **FLUX Kontext image editing**`flux-kontext-pro` and `flux-kontext-max` added to Black Forest Labs, alongside `flux-pro-1.0-fill` and `flux-pro-1.0-expand` for inpainting and outpainting
- **116 new models, 132 deprecated models cleaned up** — Major model map refresh including Mistral Magistral, Dashscope Qwen3 VL, xAI Grok via Azure AI, ZAI GLM-5, Serper Search; removal of OpenAI GPT-3.5/GPT-4 legacy variants, Gemini 1.5, and Vertex AI PaLM2
- **SageMaker Nova provider** — [New `sagemaker_nova` provider for Amazon Nova models on SageMaker](../../docs/providers/aws_sagemaker) - [PR #21542](https://github.com/BerriAI/litellm/pull/21542)
- **Hashicorp Vault secret manager** — Config override backend powered by Hashicorp Vault, with full UI for managing vault-sourced credentials - [PR #22939](https://github.com/BerriAI/litellm/pull/22939), [PR #23036](https://github.com/BerriAI/litellm/pull/23036)
- **Responses API WebSocket streaming** — Real-time WebSocket streaming for the Responses API, including support across all providers - [PR #22559](https://github.com/BerriAI/litellm/pull/22559), [PR #22771](https://github.com/BerriAI/litellm/pull/22771)
- **Org Admin RBAC expansion** — Org Admins can now access team management endpoints, view and invite internal users, and manage team membership without requiring a global admin role - [PR #23085](https://github.com/BerriAI/litellm/pull/23085), [PR #23080](https://github.com/BerriAI/litellm/pull/23080)
- **Guardrail mode defaults and tag-based modes** — Set a default guardrail mode list globally, and specify a list of modes in tag-based guardrail configs - [PR #22676](https://github.com/BerriAI/litellm/pull/22676), [PR #23020](https://github.com/BerriAI/litellm/pull/23020)
- **Secret redaction in logs** — API keys, tokens, and credentials automatically scrubbed from all proxy log output. Enabled by default; opt out with `LITELLM_DISABLE_REDACT_SECRETS=true` - [PR #23668](https://github.com/BerriAI/litellm/pull/23668)
- **Streaming stability fix** — Critical fix for `RuntimeError: Cannot send a request, as the client has been closed.` crashes after ~1 hour in production - [PR #22926](https://github.com/BerriAI/litellm/pull/22926)
@ -54,7 +58,7 @@ pip install litellm==1.82.3
## New Providers and Endpoints
### New Providers (5 new providers)
### New Providers (7 new providers)
| Provider | Supported LiteLLM Endpoints | Description |
| -------- | --------------------------- | ----------- |
@ -63,6 +67,8 @@ pip install litellm==1.82.3
| [Black Forest Labs](../../docs/providers/black_forest_labs) (`black_forest_labs/`) | `/images/generations`, `/images/edits` | FLUX image generation and editing — Kontext Pro/Max, Pro 1.0 Fill/Expand |
| [Serper](../../docs/providers/serper) (`serper/`) | `/search` | Web search via Serper API |
| [SageMaker Nova](../../docs/providers/aws_sagemaker) (`sagemaker_nova/`) | `/chat/completions` | Amazon Nova models via SageMaker endpoint |
| [Google Search API](../../docs/providers/google_search) (`google_search/`) | `/search` | Google Search API integration - [PR #22752](https://github.com/BerriAI/litellm/pull/22752) |
| [Bedrock Mantle](../../docs/providers/bedrock) (`bedrock_mantle/`) | `/chat/completions` | Amazon Bedrock via Mantle — alternative auth and routing path for Bedrock models - [PR #22866](https://github.com/BerriAI/litellm/pull/22866) |
---
@ -238,22 +244,92 @@ pip install litellm==1.82.3
- **[Responses API](../../docs/response_api)**
- Handle `response.failed`, `response.incomplete`, and `response.cancelled` terminal event types in background streaming — previously only `response.completed` was handled - [PR #23492](https://github.com/BerriAI/litellm/pull/23492)
- WebSocket streaming support for Responses API — real-time streaming via WebSocket for all providers - [PR #22559](https://github.com/BerriAI/litellm/pull/22559), [PR #22771](https://github.com/BerriAI/litellm/pull/22771)
- WebRTC support for real-time audio/video communication - [PR #23446](https://github.com/BerriAI/litellm/pull/23446)
- Responses API support for OpenAI-compatible JSON providers (`openai_like`) - [PR #21398](https://github.com/BerriAI/litellm/pull/21398)
- Route `gpt-5.4+` calls using both tools and reasoning to the Responses API automatically - [PR #23577](https://github.com/BerriAI/litellm/pull/23577)
#### Bug Fixes
- **[Anthropic Files API](../../docs/providers/anthropic)**
- Full Anthropic Files API support — upload, retrieve, list, and delete files; use file references in messages - [PR #16594](https://github.com/BerriAI/litellm/pull/16594)
- **[Mistral](../../docs/providers/mistral)**
- Voxtral audio transcription support — `mistral/voxtral-mini-*` and `mistral/voxtral-*` for audio transcription via Mistral - [PR #22801](https://github.com/BerriAI/litellm/pull/22801)
- **[OpenAI](../../docs/providers/openai)**
- `litellm.acount_tokens()` public API — async token counting with full OpenAI provider support - [PR #22809](https://github.com/BerriAI/litellm/pull/22809)
- Normalize `reasoning_effort` dict to string for chat completion API - [PR #22981](https://github.com/BerriAI/litellm/pull/22981)
- **[OpenRouter](../../docs/providers/openrouter)**
- Image edit support for OpenRouter models - [PR #22403](https://github.com/BerriAI/litellm/pull/22403)
- **[Google Vertex AI](../../docs/providers/vertex)**
- VIDEO modality token usage tracking in `completion_tokens_details` - [PR #22550](https://github.com/BerriAI/litellm/pull/22550)
- **Images API**
- `input_fidelity` parameter for image edit API - [PR #23201](https://github.com/BerriAI/litellm/pull/23201)
- **General**
- Per-request `enable_json_schema_validation` flag for thread-safe JSON schema validation - [PR #21233](https://github.com/BerriAI/litellm/pull/21233)
- Model cost aliases expansion — define aliases in the cost map that inherit pricing from a parent model - [PR #23314](https://github.com/BerriAI/litellm/pull/23314), [PR #23457](https://github.com/BerriAI/litellm/pull/23457)
- Wildcards model support for the Files API - [PR #22740](https://github.com/BerriAI/litellm/pull/22740)
#### Bugs
- **[Anthropic](../../docs/providers/anthropic)**
- Preserve native tool format (web_search, bash, tool_search, etc.) when guardrails convert tools for the Anthropic Messages API - [PR #23526](https://github.com/BerriAI/litellm/pull/23526)
- Enforce `type: "object"` on tool input schemas in `_map_tool_helper` — fixes tool call failures for strict-schema providers - [PR #23103](https://github.com/BerriAI/litellm/pull/23103)
- Deduplicate `tool_result` messages by `tool_call_id` — prevents duplicate tool result errors in multi-turn conversations - [PR #23104](https://github.com/BerriAI/litellm/pull/23104)
- Map `reasoning_effort` to `output_config` for Claude 4.6 models - [PR #22220](https://github.com/BerriAI/litellm/pull/22220)
- **[Google Gemini](../../docs/providers/gemini)**
- Correct streaming `finish_reason` for tool calls — was incorrectly returning `null` instead of `tool_calls` - [PR #21577](https://github.com/BerriAI/litellm/pull/21577)
- Preserve `$ref` in JSON Schema for Gemini 2.0+ — schema references were being stripped, breaking structured output - [PR #21597](https://github.com/BerriAI/litellm/pull/21597)
- Handle `minimal` `reasoning_effort` param for Gemini 3.1 models - [PR #22920](https://github.com/BerriAI/litellm/pull/22920)
- **[Google Vertex AI](../../docs/providers/vertex)**
- Pass through native Gemini `imageConfig` params for image generation - [PR #21585](https://github.com/BerriAI/litellm/pull/21585)
- Prevent content truncation when `finish_reason` races ahead of content chunks in streaming - [PR #22692](https://github.com/BerriAI/litellm/pull/22692)
- Strip LiteLLM-internal keys from `extra_body` before merging to Gemini request body - [PR #23131](https://github.com/BerriAI/litellm/pull/23131)
- Drop unsupported `output_config` parameter from all Vertex AI requests - [PR #22884](https://github.com/BerriAI/litellm/pull/22884)
- Skip schema transforms for Gemini 2.0+ tool parameters — avoids breaking native Gemini schema handling - [PR #23265](https://github.com/BerriAI/litellm/pull/23265)
- **[OpenRouter](../../docs/providers/openrouter)**
- Pattern-based fix for native model double-stripping when provider prefix matches model name - [PR #22320](https://github.com/BerriAI/litellm/pull/22320)
- Use provider-reported usage in streaming responses when `stream_options` is not set - [PR #21592](https://github.com/BerriAI/litellm/pull/21592)
- **[AWS Bedrock](../../docs/providers/bedrock)**
- Extract region and model ID from `bedrock/{region}/{model}` path format - [PR #22546](https://github.com/BerriAI/litellm/pull/22546)
- Strip `scope` from `cache_control` for Anthropic messages on Bedrock and Azure AI - [PR #22867](https://github.com/BerriAI/litellm/pull/22867)
- Populate `completion_tokens_details` in Responses API responses - [PR #23243](https://github.com/BerriAI/litellm/pull/23243)
- **[Azure AI](../../docs/providers/azure_ai)**
- Resolve `api_base` from environment variable in Document Intelligence OCR - [PR #21581](https://github.com/BerriAI/litellm/pull/21581)
- **[Moonshot / Kimi](../../docs/providers/openai_compatible)**
- Auto-fill `reasoning_content` for Moonshot Kimi reasoning models - [PR #23580](https://github.com/BerriAI/litellm/pull/23580)
- Preserve `image_url` blocks in multimodal messages for Moonshot - [PR #21595](https://github.com/BerriAI/litellm/pull/21595)
- **[HuggingFace](../../docs/providers/huggingface)**
- Forward `extra_headers` to HuggingFace embedding API - [PR #23525](https://github.com/BerriAI/litellm/pull/23525)
- **Token Counting / Cost**
- Fix `count_tokens` to include system prompts and tools in token counting API requests - [PR #22301](https://github.com/BerriAI/litellm/pull/22301)
- Pass all custom pricing fields to `register_model` in `completion()` and `embedding()` - [PR #22552](https://github.com/BerriAI/litellm/pull/22552)
- **Tools / Function Calling**
- Gracefully repair truncated JSON in tool call arguments — prevents crashes on malformed tool responses - [PR #22503](https://github.com/BerriAI/litellm/pull/22503)
- Fix `output_item.done` for function calls not emitting `finish_reason` in streaming - [PR #22553](https://github.com/BerriAI/litellm/pull/22553)
- Preserve thinking block order with multiple web searches - [PR #23093](https://github.com/BerriAI/litellm/pull/23093)
- **General**
- Normalize `content_filtered` finish reason across providers - [PR #23564](https://github.com/BerriAI/litellm/pull/23564)
- Unify `finish_reason` mapping to OpenAI-compatible values across all providers - [PR #22138](https://github.com/BerriAI/litellm/pull/22138)
- Fix custom cost tracking on deployments for `/v1/messages` and `/v1/responses` - [PR #23647](https://github.com/BerriAI/litellm/pull/23647)
- Fix per-request custom pricing when `router_model_id` has no pricing data — now falls back to model name
- Fix batch list showing stale `validating` status after completion - [PR #22982](https://github.com/BerriAI/litellm/pull/22982)
- Fix batch retrieve returning raw `output_file_id` when `model_id` is missing - [PR #23194](https://github.com/BerriAI/litellm/pull/23194)
- Encode batch IDs when `x-litellm-model` header is used - [PR #22653](https://github.com/BerriAI/litellm/pull/22653)
- Map `reasoning` to `reasoning_content` in streaming Delta for gpt-oss providers - [PR #22803](https://github.com/BerriAI/litellm/pull/22803)
---
@ -264,10 +340,31 @@ pip install litellm==1.82.3
- **Virtual Keys**
- Add Organization dropdown to Create/Edit Key form — `organization_id` is now a first-class field in Key Ownership - [PR #23595](https://github.com/BerriAI/litellm/pull/23595)
- Allow setting `organization_id` on `/key/update` — keys can be assigned or moved to a different organization after creation - [PR #23557](https://github.com/BerriAI/litellm/pull/23557)
- Manual Spend Reset for virtual keys from the UI — admins can reset key spend to zero on demand - [PR #22715](https://github.com/BerriAI/litellm/pull/22715)
- BYOK (Bring Your Own Key) — client-side provider API key takes precedence over proxy key for Anthropic `/v1/messages` - [PR #22964](https://github.com/BerriAI/litellm/pull/22964)
- UI login session duration configurable via `LITELLM_UI_SESSION_DURATION` environment variable - [PR #22182](https://github.com/BerriAI/litellm/pull/22182)
- Auto-redirect UI login to SSO via `auto_redirect_ui_login_to_sso: true` in config.yaml - [PR #23367](https://github.com/BerriAI/litellm/pull/23367)
- **Access Control (RBAC)**
- Org Admins can now access team management endpoints — `/team/new`, `/team/update`, `/team/delete`, `/team/member_add`, `/team/member_delete` - [PR #23085](https://github.com/BerriAI/litellm/pull/23085), [PR #23095](https://github.com/BerriAI/litellm/pull/23095)
- Org Admins can view and invite internal users — full user management without requiring global admin role - [PR #23080](https://github.com/BerriAI/litellm/pull/23080)
- Allow Admin Viewers to access Audit Logs — view-only admin role now includes audit log access - [PR #23419](https://github.com/BerriAI/litellm/pull/23419)
- RBAC for Vector Stores and Agents — key/team-level access control for vector store and agent resources - [PR #22858](https://github.com/BerriAI/litellm/pull/22858)
- User filter scope (`scope_user_search_to_org`) is now opt-in — previously default-on, causing unintended restriction - [PR #23057](https://github.com/BerriAI/litellm/pull/23057)
- **Vector Stores**
- Vector Store management endpoints — retrieve, list, update, and delete vector stores via `/v1/vector_stores/*` - [PR #23435](https://github.com/BerriAI/litellm/pull/23435)
- **Teams**
- Batch expiry setting for teams — configure a default expiry duration for all team keys - [PR #22705](https://github.com/BerriAI/litellm/pull/22705)
- Team Admin can reset key spend - [PR #22725](https://github.com/BerriAI/litellm/pull/22725)
- **Internal Users**
- Add/Remove Team Membership directly from the Internal Users info page — includes searchable dropdown and role selector; no longer requires navigating to each team - [PR #23638](https://github.com/BerriAI/litellm/pull/23638)
- **Models**
- Attach knowledge base to model via UI - [PR #22656](https://github.com/BerriAI/litellm/pull/22656)
- **Default Team Settings**
- Modernize page to antd (consistent with rest of app) - [PR #23614](https://github.com/BerriAI/litellm/pull/23614)
- Fix: default team params (budget, duration, tpm, rpm, permissions) now correctly applied on `/team/new` - [PR #23614](https://github.com/BerriAI/litellm/pull/23614)
@ -290,6 +387,13 @@ pip install litellm==1.82.3
- Fix Public Model Hub not showing config-defined models after save - [PR #23501](https://github.com/BerriAI/litellm/pull/23501)
- Fix fallback popup model dropdown z-index issue - [PR #23516](https://github.com/BerriAI/litellm/pull/23516)
- Fix double-counting bug in org/team key limit checks on `/key/update`
- Fix invite link allowing multiple password resets for the same link - [PR #22462](https://github.com/BerriAI/litellm/pull/22462)
- Fix key expiry default duration not being applied when `duration` is not set - [PR #22956](https://github.com/BerriAI/litellm/pull/22956)
- Fix all proxy models not including model access groups in key creation - [PR #23236](https://github.com/BerriAI/litellm/pull/23236)
- Fix admin viewers unable to see all organizations - [PR #22940](https://github.com/BerriAI/litellm/pull/22940)
- Fix Audit Logs UI: added server-side pagination, filtering, and drawer view - [PR #22476](https://github.com/BerriAI/litellm/pull/22476)
- Fix virtual keys in teams view not applying the team filter correctly - [PR #23065](https://github.com/BerriAI/litellm/pull/23065)
- Fix team expiry enforcement validation - [PR #22728](https://github.com/BerriAI/litellm/pull/22728)
---
@ -297,6 +401,13 @@ pip install litellm==1.82.3
### Logging
- **[Helicone](../../docs/observability/helicone_integration)**
- Add Gemini and Vertex AI support to HeliconeLogger — routes Gemini and Vertex AI requests through the correct Helicone provider URL - [PR #19288](https://github.com/BerriAI/litellm/pull/19288)
- Fix correct provider URL for Vertex AI Gemini models - [PR #22603](https://github.com/BerriAI/litellm/pull/22603)
- **[Langfuse](../../docs/proxy/logging#langfuse)**
- Fix failure path kwargs inconsistency causing dropped traces on failed requests - [PR #22390](https://github.com/BerriAI/litellm/pull/22390)
- **[Vantage](https://vantage.sh)**
- Add Vantage integration for FOCUS 1.2 CSV export — export LiteLLM proxy spend data as FinOps Open Cost & Usage Specification reports, with time-windowed filenames to prevent overwrites - [PR #23333](https://github.com/BerriAI/litellm/pull/23333)
@ -305,7 +416,10 @@ pip install litellm==1.82.3
### Guardrails
No major guardrail changes in this release.
- **Guardrail mode default list** — Configure a default list of guardrail modes applied globally when no per-request mode is specified - [PR #22676](https://github.com/BerriAI/litellm/pull/22676)
- **Tag-based guardrail mode lists** — Specify a list of modes in tag-based guardrail configs instead of a single mode - [PR #23020](https://github.com/BerriAI/litellm/pull/23020)
- **Fix presidio PII token leak** — Edge case where Anthropic handle in Presidio caused PII data exposure in token response - [PR #22627](https://github.com/BerriAI/litellm/pull/22627)
- **Fix OTEL orphaned guardrail traces** — Span redundancy and missing response IDs in OpenTelemetry guardrail traces - [PR #23001](https://github.com/BerriAI/litellm/pull/23001)
### Prompt Management
@ -313,7 +427,32 @@ No major prompt management changes in this release.
### Secret Managers
No major secret manager changes in this release.
- **[Hashicorp Vault](../../docs/secret_managers)** — Full Hashicorp Vault integration as a config override backend — secrets defined in Vault are fetched at startup and override `config.yaml` values. UI support for managing vault-sourced credentials included - [PR #22939](https://github.com/BerriAI/litellm/pull/22939), [PR #23036](https://github.com/BerriAI/litellm/pull/23036)
---
## MCP Gateway
#### Features
- **Token authentication for MCP servers** — configure `auth_type: "bearer"` per MCP server to require token-based auth on tool calls - [PR #23260](https://github.com/BerriAI/litellm/pull/23260)
- **Team-scoped MCP server filtering** — keys created under a team only see MCP servers available to that team - [PR #23323](https://github.com/BerriAI/litellm/pull/23323)
- **Per-server health recheck in UI** — trigger a health check for individual MCP servers without reloading all servers - [PR #23328](https://github.com/BerriAI/litellm/pull/23328)
#### Bugs
- Fix MCP server URL and tools management issues causing tool discovery to fail - [PR #22751](https://github.com/BerriAI/litellm/pull/22751)
- Fix MCP server health checks triggering on server deletion - [PR #23063](https://github.com/BerriAI/litellm/pull/23063)
---
## Spend Tracking, Budgets and Rate Limiting
- **Fix budget-linked keys never having spend reset** — Keys linked to budget objects were not having their spend reset on the configured reset interval - [PR #20688](https://github.com/BerriAI/litellm/pull/20688)
- **Flex pricing support** — Add `flex_pricing` field to cost map for providers that offer dynamic pricing tiers - [PR #22992](https://github.com/BerriAI/litellm/pull/22992)
- **Fix spend log cleanup** — Resolved lock tracking, integer retention, and skip-log-level issues in spend log cleanup job - [PR #22687](https://github.com/BerriAI/litellm/pull/22687)
- **Fix WebSearch spend log deduplication** — WebSearch interception was failing with thinking enabled; fixed along with spend log dedup - [PR #22679](https://github.com/BerriAI/litellm/pull/22679)
- **Fix TypeError when request has no API key** — Spend tracking was throwing unhandled exception when API key was absent from request - [PR #23363](https://github.com/BerriAI/litellm/pull/23363)
---
@ -323,6 +462,10 @@ No major secret manager changes in this release.
- **Fix OOM / Prisma connection loss** on large installs — unbounded managed-object poll was exhausting Prisma connections after ~6070 minutes on instances with 336K+ queued response rows - [PR #23472](https://github.com/BerriAI/litellm/pull/23472)
- **Centralize logging kwarg updates** — root cause fix migrating all logging updates to a single function, eliminating kwarg inconsistencies across logging paths - [PR #23659](https://github.com/BerriAI/litellm/pull/23659)
- **Fix tiktoken cache for non-root offline containers** — tiktoken cache now works correctly in offline environments running as non-root users - [PR #23498](https://github.com/BerriAI/litellm/pull/23498)
- **Block proxy startup when Redis transaction buffer has no Redis** — prevents silent data loss when `use_redis_transaction_buffer: true` is set without a Redis connection - [PR #23019](https://github.com/BerriAI/litellm/pull/23019)
- **Fix `InFlightRequestsMiddleware` crash** — undefined kwargs in middleware were causing request failures - [PR #22523](https://github.com/BerriAI/litellm/pull/22523)
- **Fix `BaseModelResponseIterator` crash on non-string stream chunks** — streaming was crashing when providers returned non-string chunk data - [PR #23497](https://github.com/BerriAI/litellm/pull/23497)
- **Fix `SERVER_ROOT_PATH` prefix handling** — strip prefix before checking mapped pass-through routes to prevent double-prefix issues - [PR #23414](https://github.com/BerriAI/litellm/pull/23414)
- **Add CodSpeed continuous performance benchmarks** — automated performance regression tracking on CI - [PR #23676](https://github.com/BerriAI/litellm/pull/23676)
---
@ -342,6 +485,16 @@ No major secret manager changes in this release.
---
## Documentation Updates
- Add Anthropic `/v1/messages``/responses` parameter mapping reference - [PR #22893](https://github.com/BerriAI/litellm/pull/22893)
- Update Okta SSO docs and custom SSO handler example - [PR #22786](https://github.com/BerriAI/litellm/pull/22786)
- Add `LITELLM_MAX_BUDGET_PER_SESSION_TTL` to environment variables reference - [PR #23186](https://github.com/BerriAI/litellm/pull/23186)
- Add DB query performance guidelines to `CLAUDE.md` - [PR #23196](https://github.com/BerriAI/litellm/pull/23196)
- Add Gemini Vertex AI PayGo/priority cost tracking docs - [PR #22948](https://github.com/BerriAI/litellm/pull/22948)
---
## New Contributors
* @ryanh-ai made their first contribution in [PR #21542](https://github.com/BerriAI/litellm/pull/21542)
@ -359,14 +512,17 @@ No major secret manager changes in this release.
## Diff Summary
## 03/16/2026
* New Providers: 5
* New Providers: 7
* New Models / Updated Models: 116 new, 132 removed
* LLM API Endpoints: 5
* Management Endpoints / UI: 11
* AI Integrations: 2
* Performance / Reliability: 5
* LLM API Endpoints: 37
* Management Endpoints / UI: 31
* AI Integrations: 8
* MCP Gateway: 5
* Spend Tracking, Budgets and Rate Limiting: 5
* Performance / Loadbalancing / Reliability improvements: 9
* Security: 3
* Database / Proxy Operations: 2
* Documentation Updates: 5
---

View file

@ -48,6 +48,20 @@ const sidebars = {
slug: "/guardrail_providers"
},
items: [
{
type: "category",
label: "Contributing to Guardrails",
items: [
"adding_provider/generic_guardrail_api",
"adding_provider/simple_guardrail_tutorial",
"adding_provider/adding_guardrail_support",
]
},
{
type: "doc",
id: "proxy/guardrails/team_based_guardrails",
label: "Team Bring-Your-Own Guardrails",
},
...[
"proxy/guardrails/qualifire",
"proxy/guardrails/aim_security",

View file

@ -524,6 +524,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
optional_params.api_base
or litellm.api_base
or get_secret_str("ANTHROPIC_API_BASE")
or get_secret_str("ANTHROPIC_BASE_URL")
)
api_key = (
optional_params.api_key

View file

@ -1444,6 +1444,7 @@ SENTRY_DENYLIST = [
"credential",
"OPENAI_API_KEY",
"ANTHROPIC_API_KEY",
"ANTHROPIC_AUTH_TOKEN",
"AZURE_API_KEY",
"COHERE_API_KEY",
"REPLICATE_API_KEY",

View file

@ -757,7 +757,7 @@ def _map_traffic_type_to_service_tier(traffic_type: Optional[str]) -> Optional[s
"""
if traffic_type is None:
return None
service_tier = _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER.get(traffic_type.upper())
service_tier = _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER.get(str(traffic_type).upper())
return service_tier

View file

@ -291,7 +291,7 @@ class DataDogLogger(
dd_payload = DatadogPayload(
ddsource=get_datadog_source(),
ddtags=get_datadog_tags(),
ddtags=",".join(get_datadog_tags()),
hostname=get_datadog_hostname(),
message=safe_dumps(message_payload),
service=get_datadog_service(),
@ -442,7 +442,9 @@ class DataDogLogger(
verbose_logger.debug("Datadog: Logger - Logging payload = %s", json_payload)
dd_payload = DatadogPayload(
ddsource=get_datadog_source(),
ddtags=get_datadog_tags(standard_logging_object=standard_logging_object),
ddtags=",".join(
get_datadog_tags(standard_logging_object=standard_logging_object)
),
hostname=get_datadog_hostname(),
message=json_payload,
service=get_datadog_service(),
@ -545,7 +547,7 @@ class DataDogLogger(
_dd_message_str = safe_dumps(_payload_dict)
_dd_payload = DatadogPayload(
ddsource=get_datadog_source(),
ddtags=get_datadog_tags(),
ddtags=",".join(get_datadog_tags()),
hostname=get_datadog_hostname(),
message=_dd_message_str,
service=get_datadog_service(),
@ -587,7 +589,7 @@ class DataDogLogger(
_dd_message_str = safe_dumps(_payload_dict)
_dd_payload = DatadogPayload(
ddsource=get_datadog_source(),
ddtags=get_datadog_tags(),
ddtags=",".join(get_datadog_tags()),
hostname=get_datadog_hostname(),
message=_dd_message_str,
service=get_datadog_service(),
@ -678,7 +680,7 @@ class DataDogLogger(
dd_payload = DatadogPayload(
ddsource=get_datadog_source(),
ddtags=get_datadog_tags(),
ddtags=",".join(get_datadog_tags()),
hostname=get_datadog_hostname(),
message=json_payload,
service=get_datadog_service(),

View file

@ -38,8 +38,13 @@ def get_datadog_pod_name() -> str:
def get_datadog_tags(
standard_logging_object: Optional[StandardLoggingPayload] = None,
) -> str:
"""Build Datadog tags string used by multiple integrations."""
) -> List[str]:
"""Build Datadog tags as a list of individual tag strings.
Returns a list of "key:value" strings suitable for Datadog LLM Observability
(which expects tags as an array). For Datadog Logs API (ddtags), join with
comma: ",".join(get_datadog_tags(...)).
"""
base_tags = {
"env": get_datadog_env(),
@ -66,4 +71,4 @@ def get_datadog_tags(
if team_tag:
tags.append(f"team:{team_tag}")
return ",".join(tags)
return tags

View file

@ -203,7 +203,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
type="span",
attributes=DDSpanAttributes(
ml_app=get_datadog_service(),
tags=[get_datadog_tags()],
tags=get_datadog_tags(),
spans=self.log_queue,
),
),
@ -315,7 +315,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
duration=int((end_time - start_time).total_seconds() * 1e9),
metrics=metrics,
status="error" if error_info else "ok",
tags=[get_datadog_tags(standard_logging_object=standard_logging_payload)],
tags=get_datadog_tags(standard_logging_object=standard_logging_payload),
)
apm_trace_id = self._get_apm_trace_id()

View file

@ -5,7 +5,6 @@ import os
import random
import traceback
import types
from litellm._uuid import uuid
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
@ -14,10 +13,11 @@ from pydantic import BaseModel # type: ignore
import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.integrations.langsmith_mock_client import (
should_use_langsmith_mock,
create_mock_langsmith_client,
should_use_langsmith_mock,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -110,6 +110,60 @@ class LangsmithLogger(CustomBatchLogger):
LANGSMITH_TENANT_ID=_credentials_tenant_id,
)
def _extract_metadata_fields(
self, metadata: dict, credentials: LangsmithCredentialsObject
):
return {
"project_name": metadata.get(
"project_name", credentials["LANGSMITH_PROJECT"]
),
"run_name": metadata.get("run_name", self.langsmith_default_run_name),
"run_id": metadata.get("id", metadata.get("run_id", None)),
"parent_run_id": metadata.get("parent_run_id", None),
"trace_id": metadata.get("trace_id", None),
"session_id": metadata.get("session_id", None),
"dotted_order": metadata.get("dotted_order", None),
}
def _build_extra_metadata(self, metadata: Dict):
extra_metadata = dict(metadata)
requester_metadata = extra_metadata.get("requester_metadata")
if requester_metadata and isinstance(requester_metadata, dict):
for key in ("session_id", "thread_id", "conversation_id"):
if key in requester_metadata and key not in extra_metadata:
extra_metadata[key] = requester_metadata[key]
return extra_metadata
def _build_outputs_with_usage(
self, payload: StandardLoggingPayload
) -> Dict[str, Any]:
response = payload["response"]
outputs: Dict[str, Any]
if isinstance(response, dict):
outputs = {**response}
else:
outputs = {"output": response}
outputs["usage_metadata"] = {
"input_tokens": payload.get("prompt_tokens", 0),
"output_tokens": payload.get("completion_tokens", 0),
"total_tokens": payload.get("total_tokens", 0),
"total_cost": payload.get("response_cost", 0),
}
return outputs
def _ensure_required_ids(self, data: dict, run_id: Optional[str]):
if "id" not in data or data["id"] is None:
run_id = str(uuid.uuid4())
data["id"] = run_id
if "trace_id" not in data or data["trace_id"] is None:
if run_id is not None and isinstance(run_id, str):
data["trace_id"] = run_id
if "dotted_order" not in data or data["dotted_order"] is None:
if run_id is not None and isinstance(run_id, str):
data["dotted_order"] = self.make_dot_order(run_id=run_id)
def _prepare_log_data(
self,
kwargs,
@ -121,44 +175,28 @@ class LangsmithLogger(CustomBatchLogger):
try:
_litellm_params = kwargs.get("litellm_params", {}) or {}
metadata = _litellm_params.get("metadata", {}) or {}
project_name = metadata.get(
"project_name", credentials["LANGSMITH_PROJECT"]
)
run_name = metadata.get("run_name", self.langsmith_default_run_name)
run_id = metadata.get("id", metadata.get("run_id", None))
parent_run_id = metadata.get("parent_run_id", None)
trace_id = metadata.get("trace_id", None)
session_id = metadata.get("session_id", None)
dotted_order = metadata.get("dotted_order", None)
fields = self._extract_metadata_fields(metadata, credentials)
verbose_logger.debug(
f"Langsmith Logging - project_name: {project_name}, run_name {run_name}"
f"Langsmith Logging - project_name: {fields['project_name']}, run_name {fields['run_name']}"
)
# Ensure everything in the payload is converted to str
payload: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object", None
)
if payload is None:
raise Exception("Error logging request payload. Payload=none.")
metadata = payload[
"metadata"
] # ensure logged metadata is json serializable
extra_metadata = dict(metadata)
requester_metadata = extra_metadata.get("requester_metadata")
if requester_metadata and isinstance(requester_metadata, dict):
for key in ("session_id", "thread_id", "conversation_id"):
if key in requester_metadata and key not in extra_metadata:
extra_metadata[key] = requester_metadata[key]
metadata = payload["metadata"]
extra_metadata = self._build_extra_metadata(dict(metadata))
outputs = self._build_outputs_with_usage(payload)
data = {
"name": run_name,
"run_type": "llm", # this should always be llm, since litellm always logs llm calls. Langsmith allow us to log "chain"
"name": fields["run_name"],
"run_type": "llm",
"inputs": payload,
"outputs": payload["response"],
"session_name": project_name,
"outputs": outputs,
"session_name": fields["project_name"],
"start_time": payload["startTime"],
"end_time": payload["endTime"],
"tags": payload["request_tags"],
@ -168,46 +206,19 @@ class LangsmithLogger(CustomBatchLogger):
if payload["error_str"] is not None and payload["status"] == "failure":
data["error"] = payload["error_str"]
if run_id:
data["id"] = run_id
if parent_run_id:
data["parent_run_id"] = parent_run_id
if trace_id:
data["trace_id"] = trace_id
if session_id:
data["session_id"] = session_id
if dotted_order:
data["dotted_order"] = dotted_order
run_id: Optional[str] = data.get("id") # type: ignore
if "id" not in data or data["id"] is None:
"""
for /batch langsmith requires id, trace_id and dotted_order passed as params
"""
run_id = str(uuid.uuid4())
data["id"] = run_id
if (
"trace_id" not in data
or data["trace_id"] is None
and (run_id is not None and isinstance(run_id, str))
for key in (
"id",
"parent_run_id",
"trace_id",
"session_id",
"dotted_order",
):
data["trace_id"] = run_id
if (
"dotted_order" not in data
or data["dotted_order"] is None
and (run_id is not None and isinstance(run_id, str))
):
data["dotted_order"] = self.make_dot_order(run_id=run_id) # type: ignore
field_key = "run_id" if key == "id" else key
if fields[field_key]:
data[key] = fields[field_key]
self._ensure_required_ids(data, fields["run_id"])
verbose_logger.debug("Langsmith Logging data on langsmith: %s", data)
return data
except Exception:
raise

View file

@ -84,6 +84,8 @@ from litellm.types.llms.openai import (
OpenAIModerationResponse,
ResponseAPIUsage,
ResponseCompletedEvent,
ResponseFailedEvent,
ResponseIncompleteEvent,
ResponsesAPIResponse,
)
from litellm.types.mcp import MCPPostCallResponseObject
@ -516,6 +518,23 @@ class Logging(LiteLLMLoggingBaseClass):
),
)
def get_router_model_id(self) -> Optional[str]:
"""Extract the router deployment model_id from litellm_params.
Checks both litellm_metadata and metadata for model_info.id.
Used by cost calculators to look up custom pricing registered
under the deployment's model_info.id in litellm.model_cost.
"""
if not hasattr(self, "litellm_params"):
return None
for key in ("litellm_metadata", "metadata"):
meta = self.litellm_params.get(key, {}) or {}
info = meta.get("model_info", {}) or {}
model_id = info.get("id")
if model_id is not None:
return model_id
return None
def update_environment_variables(
self,
litellm_params: Dict,
@ -1455,6 +1474,12 @@ class Logging(LiteLLMLoggingBaseClass):
): # use model_id if not already set
router_model_id = hidden_params["model_id"]
# Fallback: extract router_model_id from litellm_params when not available
# from the result object. ResponsesAPIResponse objects (used by /v1/responses
# streaming) don't carry _hidden_params["model_id"] like ModelResponse does.
if router_model_id is None:
router_model_id = self.get_router_model_id()
## RESPONSE COST ##
custom_pricing = use_custom_pricing_for_model(
litellm_params=(
@ -2958,8 +2983,7 @@ class Logging(LiteLLMLoggingBaseClass):
if (
isinstance(callback, CustomLogger)
and is_sync_request
and self.call_type
!= CallTypes.pass_through.value
and self.call_type != CallTypes.pass_through.value
): # custom logger class
callback.log_failure_event(
start_time=start_time,
@ -3307,7 +3331,10 @@ class Logging(LiteLLMLoggingBaseClass):
return result
elif isinstance(result, TextCompletionResponse):
return result
elif isinstance(result, ResponseCompletedEvent):
elif isinstance(
result,
(ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent),
):
## return unified Usage object
if isinstance(result.response.usage, ResponseAPIUsage):
transformed_usage = (
@ -3328,7 +3355,6 @@ class Logging(LiteLLMLoggingBaseClass):
return result.response
else:
return None
return None
def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelResponse:
"""

View file

@ -160,6 +160,7 @@ class CustomStreamWrapper:
self.chunks: List = (
[]
) # keep track of the returned chunks - used for calculating the input/output tokens for stream options
self._repeated_messages_count = 1
self.is_function_call = self.check_is_function_call(logging_obj=logging_obj)
self.created: Optional[int] = None
self._last_returned_hidden_params: Optional[dict] = None
@ -241,7 +242,7 @@ class CustomStreamWrapper:
except Exception as e:
raise e
def safety_checker(self) -> None:
def raise_on_model_repetition(self) -> None:
"""
Fixes - https://github.com/BerriAI/litellm/issues/5158
@ -249,28 +250,35 @@ class CustomStreamWrapper:
Raises - InternalServerError, if LLM enters infinite loop while streaming
"""
if len(self.chunks) >= litellm.REPEATED_STREAMING_CHUNK_LIMIT:
# Get the last n chunks
last_chunks = self.chunks[-litellm.REPEATED_STREAMING_CHUNK_LIMIT :]
if len(self.chunks) < 2:
return
# Extract the relevant content from the chunks
last_contents = [chunk.choices[0].delta.content for chunk in last_chunks]
last_content = self.chunks[-1].choices[0].delta.content
# Check if all extracted contents are identical
if all(content == last_contents[0] for content in last_contents):
if (
last_contents[0] is not None
and isinstance(last_contents[0], str)
and len(last_contents[0]) > 2
): # ignore empty content - https://github.com/BerriAI/litellm/issues/5158#issuecomment-2287156946
# All last n chunks are identical
raise litellm.InternalServerError(
message="The model is repeating the same chunk = {}.".format(
last_contents[0]
),
model="",
llm_provider="",
)
if (
last_content is None
or not isinstance(last_content, str)
or len(last_content) <= 2
): # ignore empty content - https://github.com/BerriAI/litellm/issues/5158#issuecomment-2287156946
self._repeated_messages_count = 1
return
second_to_last_content = self.chunks[-2].choices[0].delta.content
if last_content == second_to_last_content:
self._repeated_messages_count += 1
else:
self._repeated_messages_count = 1
if self._repeated_messages_count >= litellm.REPEATED_STREAMING_CHUNK_LIMIT:
# All last n chunks are identical
raise litellm.InternalServerError(
message="The model is repeating the same chunk = {}.".format(
last_content
),
model="",
llm_provider="",
)
def check_special_tokens(self, chunk: str, finish_reason: Optional[str]):
"""
@ -924,7 +932,7 @@ class CustomStreamWrapper:
if (
is_chunk_non_empty
): # cannot set content of an OpenAI Object to be an empty string
self.safety_checker()
self.raise_on_model_repetition()
hold, model_response_str = self.check_special_tokens(
chunk=completion_obj["content"],
finish_reason=model_response.choices[0].finish_reason,
@ -1893,15 +1901,19 @@ class CustomStreamWrapper:
"usage",
getattr(complete_streaming_response, "usage"),
)
try:
_cache_copy = complete_streaming_response.model_copy(deep=True)
_log_copy = complete_streaming_response.model_copy(deep=True)
except RuntimeError:
_cache_copy = complete_streaming_response.model_copy()
_log_copy = complete_streaming_response.model_copy()
self.cache_streaming_response(
processed_chunk=complete_streaming_response.model_copy(
deep=True
),
processed_chunk=_cache_copy,
cache_hit=cache_hit,
)
executor.submit(
self.logging_obj.success_handler,
complete_streaming_response.model_copy(deep=True),
_log_copy,
None,
None,
cache_hit,
@ -2113,11 +2125,13 @@ class CustomStreamWrapper:
"usage",
getattr(complete_streaming_response, "usage"),
)
try:
_copy = complete_streaming_response.model_copy(deep=True)
except RuntimeError:
_copy = complete_streaming_response.model_copy()
asyncio.create_task(
self.async_cache_streaming_response(
processed_chunk=complete_streaming_response.model_copy(
deep=True
),
processed_chunk=_copy,
cache_hit=cache_hit,
)
)

View file

@ -42,9 +42,8 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
api_base: Optional[str] = None,
) -> dict:
"""Validate and prepare environment-specific headers and parameters."""
# Resolve api_key from environment if not provided
api_key = api_key or self.anthropic_model_info.get_api_key()
if api_key is None:
auth_header = self.anthropic_model_info.get_auth_header(api_key)
if auth_header is None:
raise ValueError(
"Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params"
)
@ -52,8 +51,8 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
"accept": "application/json",
"anthropic-version": "2023-06-01",
"content-type": "application/json",
"x-api-key": api_key,
}
_headers.update(auth_header)
# Add beta header for message batches
if "anthropic-beta" not in headers:
headers["anthropic-beta"] = "message-batches-2024-09-24"

View file

@ -48,6 +48,10 @@ from litellm.types.llms.openai import (
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
)
from litellm.types.responses.main import (
OutputCodeInterpreterCall,
build_code_interpreter_log_outputs,
)
from litellm.types.utils import (
Delta,
GenericStreamingChunk,
@ -538,6 +542,12 @@ class ModelResponseIterator:
# Accumulate compaction blocks for multi-turn reconstruction
self.compaction_blocks: List[Dict[str, Any]] = []
# Track server tool use inputs and results for code_interpreter_results
self._server_tool_inputs: Dict[str, Any] = {}
self.tool_results: List[Dict[str, Any]] = []
self._current_server_tool_id: Optional[str] = None
self._container_id: Optional[str] = None
def check_empty_tool_call_args(self) -> bool:
"""
Check if the tool call block so far has been an empty string
@ -682,6 +692,39 @@ class ModelResponseIterator:
return content_block_start
def _build_code_interpreter_results(self) -> list:
"""Convert accumulated tool_results to OutputCodeInterpreterCall objects.
Called during streaming to produce provider-neutral code_interpreter_results
alongside the raw tool_results, so the Responses API layer doesn't need
Anthropic-specific knowledge.
Returns the full cumulative list each time (not incremental), matching
how web_search_results works. stream_chunk_builder uses "last value
wins" for list-valued provider_specific_fields keys, so the last
emission must contain every result.
"""
results = []
for tr in self.tool_results:
if tr.get("type") != "bash_code_execution_tool_result":
continue
call_id = tr.get("tool_use_id", "")
content = tr.get("content", {})
log_outputs = build_code_interpreter_log_outputs(content)
tool_input = self._server_tool_inputs.get(call_id, {})
code = tool_input.get("command", "") if isinstance(tool_input, dict) else ""
results.append(
OutputCodeInterpreterCall(
type="code_interpreter_call",
id=call_id,
code=code,
container_id=self._container_id,
status="completed",
outputs=log_outputs,
)
)
return results
def chunk_parser(self, chunk: dict) -> ModelResponseStream: # noqa: PLR0915
try:
type_chunk = chunk.get("type", "") or ""
@ -748,6 +791,23 @@ class ModelResponseIterator:
),
index=self.tool_index,
)
# Track server tool use inputs for code_interpreter_results.
# The initial input in content_block_start is typically {}
# for streaming; the full input arrives via input_json_delta
# and is assembled at content_block_stop.
if (
content_block_start["content_block"]["type"]
== "server_tool_use"
):
self._current_server_tool_id = content_block_start[
"content_block"
]["id"]
tool_input = content_block_start["content_block"].get(
"input", {}
)
self._server_tool_inputs[
self._current_server_tool_id
] = tool_input
# Include caller information if present (for programmatic tool calling)
if "caller" in content_block_start["content_block"]:
caller_data = content_block_start["content_block"]["caller"]
@ -808,10 +868,12 @@ class ModelResponseIterator:
elif content_type != "tool_search_tool_result":
# Handle other tool results (code execution, etc.)
# Skip tool_search_tool_result as it's internal metadata
if not hasattr(self, "tool_results"):
self.tool_results = []
self.tool_results.append(content_block_start["content_block"])
provider_specific_fields["tool_results"] = self.tool_results
# Convert to provider-neutral code_interpreter_results
provider_specific_fields[
"code_interpreter_results"
] = self._build_code_interpreter_results()
elif type_chunk == "content_block_stop":
ContentBlockStop(**chunk) # type: ignore
@ -828,6 +890,26 @@ class ModelResponseIterator:
),
index=self.tool_index,
)
# Update server_tool_inputs with fully assembled input
# from input_json_delta chunks (content_block_start has {})
if (
self.current_content_block_type == "server_tool_use"
and self._current_server_tool_id
):
args = ""
for block in self.content_blocks:
if block["delta"]["type"] == "input_json_delta":
partial_json = block["delta"].get("partial_json")
if isinstance(partial_json, str):
args += partial_json
if args:
try:
self._server_tool_inputs[
self._current_server_tool_id
] = json.loads(args)
except (json.JSONDecodeError, TypeError):
pass
self._current_server_tool_id = None
# Reset response_format tool tracking when block stops
self.is_response_format_tool = False
# Reset current content block type
@ -840,6 +922,17 @@ class ModelResponseIterator:
finish_reason, usage, container = self._handle_message_delta(chunk)
if container:
provider_specific_fields["container"] = container
# Store container_id and re-emit code_interpreter_results
# so stream_chunk_builder's last-value-wins picks up the
# version with container_id populated.
container_id = (
container.get("id") if isinstance(container, dict) else None
)
if container_id and self.tool_results:
self._container_id = container_id
provider_specific_fields[
"code_interpreter_results"
] = self._build_code_interpreter_results()
elif type_chunk == "message_start":
"""
Anthropic

View file

@ -50,6 +50,10 @@ from litellm.types.llms.openai import (
OpenAIMcpServerTool,
OpenAIWebSearchOptions,
)
from litellm.types.responses.main import (
OutputCodeInterpreterCall,
build_code_interpreter_log_outputs,
)
from litellm.types.utils import (
CacheCreationTokenDetails,
CompletionTokensDetailsWrapper,
@ -1522,7 +1526,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
tool_results = []
tool_results.append(content)
elif content.get("thinking", None) is not None:
elif content.get("type") == "thinking":
if thinking_blocks is None:
thinking_blocks = []
thinking_blocks.append(cast(ChatCompletionThinkingBlock, content))
@ -1682,6 +1686,96 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
return usage
def _build_code_by_id_map(
self, tool_calls: List[ChatCompletionToolCallChunk]
) -> Dict[str, str]:
code_by_id: Dict[str, str] = {}
for tc in tool_calls:
try:
args = json.loads(tc.get("function", {}).get("arguments", "{}"))
call_id = tc.get("id")
command = args.get("command", "")
if isinstance(call_id, str):
code_by_id[call_id] = command if isinstance(command, str) else ""
except Exception:
pass
return code_by_id
def _build_code_interpreter_results(
self,
tool_results: List[Any],
code_by_id: Dict[str, str],
container_id: Optional[str],
) -> List[OutputCodeInterpreterCall]:
code_interpreter_results = []
for tr in tool_results:
if tr.get("type") != "bash_code_execution_tool_result":
continue
call_id = tr.get("tool_use_id", "")
content = tr.get("content", {})
log_outputs = build_code_interpreter_log_outputs(content)
code_interpreter_results.append(
OutputCodeInterpreterCall(
type="code_interpreter_call",
id=call_id,
code=code_by_id.get(call_id, ""),
container_id=container_id,
status="completed",
outputs=log_outputs,
)
)
return code_interpreter_results
def _build_provider_specific_fields(
self,
completion_response: dict,
citations: Optional[List[Any]],
thinking_blocks: Optional[
List[
Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]
]
],
web_search_results: Optional[List[Any]],
tool_results: Optional[List[Any]],
compaction_blocks: Optional[List[Any]],
tool_calls: List[ChatCompletionToolCallChunk],
) -> Dict[str, Any]:
provider_specific_fields: Dict[str, Any] = {
"citations": citations,
"thinking_blocks": thinking_blocks,
}
context_management = completion_response.get("context_management")
if context_management is not None:
provider_specific_fields["context_management"] = context_management
if web_search_results is not None:
provider_specific_fields["web_search_results"] = web_search_results
if tool_results is not None:
provider_specific_fields["tool_results"] = tool_results
container_id = (
completion_response.get("container", {}).get("id")
if isinstance(completion_response.get("container"), dict)
else None
)
code_by_id = self._build_code_by_id_map(tool_calls)
code_interpreter_results = self._build_code_interpreter_results(
tool_results, code_by_id, container_id
)
provider_specific_fields[
"code_interpreter_results"
] = code_interpreter_results
container = completion_response.get("container")
if container is not None:
provider_specific_fields["container"] = container
if compaction_blocks is not None:
provider_specific_fields["compaction_blocks"] = compaction_blocks
return provider_specific_fields
def transform_parsed_response(
self,
completion_response: dict,
@ -1702,98 +1796,73 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
status_code=raw_response.status_code,
headers=response_headers,
)
else:
text_content = ""
citations: Optional[List[Any]] = None
thinking_blocks: Optional[
List[
Union[
ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock
]
]
] = None
reasoning_content: Optional[str] = None
tool_calls: List[ChatCompletionToolCallChunk] = []
(
text_content,
citations,
thinking_blocks,
reasoning_content,
tool_calls,
web_search_results,
tool_results,
compaction_blocks,
) = self.extract_response_content(completion_response=completion_response)
(
text_content,
citations,
thinking_blocks,
reasoning_content,
tool_calls,
web_search_results,
tool_results,
compaction_blocks,
) = self.extract_response_content(completion_response=completion_response)
if (
prefix_prompt is not None
and not text_content.startswith(prefix_prompt)
and not litellm.disable_add_prefix_to_prompt
):
text_content = prefix_prompt + text_content
if (
prefix_prompt is not None
and not text_content.startswith(prefix_prompt)
and not litellm.disable_add_prefix_to_prompt
):
text_content = prefix_prompt + text_content
context_management: Optional[Dict] = completion_response.get(
"context_management"
)
provider_specific_fields = self._build_provider_specific_fields(
completion_response,
citations,
thinking_blocks,
web_search_results,
tool_results,
compaction_blocks,
tool_calls,
)
container: Optional[Dict] = completion_response.get("container")
_message = litellm.Message(
tool_calls=tool_calls,
content=text_content or None,
provider_specific_fields=provider_specific_fields,
thinking_blocks=thinking_blocks,
reasoning_content=reasoning_content,
)
_message.provider_specific_fields = provider_specific_fields
provider_specific_fields: Dict[str, Any] = {
"citations": citations,
"thinking_blocks": thinking_blocks,
}
if context_management is not None:
provider_specific_fields["context_management"] = context_management
if web_search_results is not None:
provider_specific_fields["web_search_results"] = web_search_results
if tool_results is not None:
provider_specific_fields["tool_results"] = tool_results
if container is not None:
provider_specific_fields["container"] = container
if compaction_blocks is not None:
provider_specific_fields["compaction_blocks"] = compaction_blocks
json_mode_message = self._transform_response_for_json_mode(
json_mode=json_mode,
tool_calls=tool_calls,
)
if json_mode_message is not None:
completion_response["stop_reason"] = "stop"
_message = json_mode_message
_message = litellm.Message(
tool_calls=tool_calls,
content=text_content or None,
provider_specific_fields=provider_specific_fields,
thinking_blocks=thinking_blocks,
reasoning_content=reasoning_content,
)
_message.provider_specific_fields = provider_specific_fields
model_response.choices[0].message = _message
model_response._hidden_params["original_response"] = completion_response[
"content"
]
model_response.choices[0].finish_reason = cast(
OpenAIChatCompletionFinishReason,
map_finish_reason(completion_response["stop_reason"]),
)
## HANDLE JSON MODE - anthropic returns single function call
json_mode_message = self._transform_response_for_json_mode(
json_mode=json_mode,
tool_calls=tool_calls,
)
if json_mode_message is not None:
completion_response["stop_reason"] = "stop"
_message = json_mode_message
model_response.choices[0].message = _message # type: ignore
model_response._hidden_params["original_response"] = completion_response[
"content"
] # allow user to access raw anthropic tool calling response
model_response.choices[0].finish_reason = cast(
OpenAIChatCompletionFinishReason,
map_finish_reason(completion_response["stop_reason"]),
)
## CALCULATING USAGE
usage = self.calculate_usage(
usage_object=completion_response["usage"],
reasoning_content=reasoning_content,
completion_response=completion_response,
speed=speed,
)
setattr(model_response, "usage", usage) # type: ignore
setattr(model_response, "usage", usage)
model_response.created = int(time.time())
model_response.model = completion_response["model"]
_hidden_params["provider_specific_fields"] = provider_specific_fields
model_response._hidden_params = _hidden_params
return model_response

View file

@ -359,9 +359,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
Returns:
List of beta header strings
"""
from litellm.types.llms.anthropic import (
ANTHROPIC_EFFORT_BETA_HEADER,
)
from litellm.types.llms.anthropic import ANTHROPIC_EFFORT_BETA_HEADER
betas = []
@ -390,7 +388,8 @@ class AnthropicModelInfo(BaseLLMModelInfo):
def get_anthropic_headers(
self,
api_key: str,
api_key: Optional[str] = None,
auth_token: Optional[str] = None,
anthropic_version: Optional[str] = None,
computer_tool_used: Optional[str] = None,
prompt_caching_set: bool = False,
@ -451,7 +450,9 @@ class AnthropicModelInfo(BaseLLMModelInfo):
headers["authorization"] = f"Bearer {api_key}"
headers["anthropic-dangerous-direct-browser-access"] = "true"
betas.add(ANTHROPIC_OAUTH_BETA_HEADER)
else:
elif auth_token and not api_key:
headers["authorization"] = f"Bearer {auth_token}"
elif api_key:
headers["x-api-key"] = api_key
if user_anthropic_beta_headers is not None:
@ -485,9 +486,14 @@ class AnthropicModelInfo(BaseLLMModelInfo):
headers, api_key = optionally_handle_anthropic_oauth(
headers=headers, api_key=api_key
)
api_key = AnthropicModelInfo.get_api_key(api_key)
# Resolve auth_token from ANTHROPIC_AUTH_TOKEN if api_key is not set
auth_token: Optional[str] = None
if api_key is None:
auth_token = AnthropicModelInfo.get_auth_token()
if api_key is None and auth_token is None:
raise litellm.AuthenticationError(
message="Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` in your environment vars",
message="Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` or `ANTHROPIC_AUTH_TOKEN` in your environment vars",
llm_provider="anthropic",
model=model,
)
@ -519,6 +525,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
prompt_caching_set=prompt_caching_set,
pdf_used=pdf_used,
api_key=api_key,
auth_token=auth_token,
file_id_used=file_id_used,
web_search_tool_used=web_search_tool_used,
is_vertex_request=optional_params.get("is_vertex_request", False),
@ -543,6 +550,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
return (
api_base
or get_secret_str("ANTHROPIC_API_BASE")
or get_secret_str("ANTHROPIC_BASE_URL")
or "https://api.anthropic.com"
)
@ -552,6 +560,35 @@ class AnthropicModelInfo(BaseLLMModelInfo):
return api_key or get_secret_str("ANTHROPIC_API_KEY")
@staticmethod
def get_auth_token(auth_token: Optional[str] = None) -> Optional[str]:
"""Get auth token from ANTHROPIC_AUTH_TOKEN env var.
Unlike api_key (which uses X-Api-Key header), auth_token uses
Authorization: Bearer header, matching the official Anthropic SDK behavior.
"""
from litellm.secret_managers.main import get_secret_str
return auth_token or get_secret_str("ANTHROPIC_AUTH_TOKEN")
@staticmethod
def get_auth_header(api_key: Optional[str] = None) -> Optional[dict]:
"""Resolve Anthropic credentials and return the appropriate auth header dict.
Checks ANTHROPIC_API_KEY first (-> x-api-key), then
ANTHROPIC_AUTH_TOKEN (-> Authorization: Bearer).
Returns None if neither is available.
"""
resolved_key = AnthropicModelInfo.get_api_key(api_key)
if resolved_key is not None:
if is_anthropic_oauth_key(resolved_key):
return {"authorization": f"Bearer {resolved_key}"}
return {"x-api-key": resolved_key}
auth_token = AnthropicModelInfo.get_auth_token()
if auth_token is not None:
return {"authorization": f"Bearer {auth_token}"}
return None
@staticmethod
def get_base_model(model: Optional[str] = None) -> Optional[str]:
return model.replace("anthropic/", "") if model else None
@ -560,14 +597,16 @@ class AnthropicModelInfo(BaseLLMModelInfo):
self, api_key: Optional[str] = None, api_base: Optional[str] = None
) -> List[str]:
api_base = AnthropicModelInfo.get_api_base(api_base)
api_key = AnthropicModelInfo.get_api_key(api_key)
if api_base is None or api_key is None:
auth_header = AnthropicModelInfo.get_auth_header(api_key)
if api_base is None or auth_header is None:
raise ValueError(
"ANTHROPIC_API_BASE or ANTHROPIC_API_KEY is not set. Please set the environment variable, to query Anthropic's `/models` endpoint."
"ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN is not set. Please set the environment variable, to query Anthropic's `/models` endpoint."
)
headers = {"anthropic-version": "2023-06-01"}
headers.update(auth_header)
response = litellm.module_level_client.get(
url=f"{api_base}/v1/models",
headers={"x-api-key": api_key, "anthropic-version": "2023-06-01"},
headers=headers,
)
try:

View file

@ -23,7 +23,6 @@ from ...common_utils import (
optionally_handle_anthropic_oauth,
)
DEFAULT_ANTHROPIC_API_BASE = "https://api.anthropic.com"
DEFAULT_ANTHROPIC_API_VERSION = "2023-06-01"
@ -127,7 +126,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
api_base = api_base or DEFAULT_ANTHROPIC_API_BASE
api_base = (
AnthropicModelInfo.get_api_base(api_base) or "https://api.anthropic.com"
)
if not api_base.endswith("/v1/messages"):
api_base = f"{api_base}/v1/messages"
return api_base
@ -142,17 +143,15 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> Tuple[dict, Optional[str]]:
import os
# Check for Anthropic OAuth token in Authorization header
headers, api_key = optionally_handle_anthropic_oauth(
headers=headers, api_key=api_key
)
if api_key is None:
api_key = os.getenv("ANTHROPIC_API_KEY")
if "x-api-key" not in headers and "authorization" not in headers and api_key:
headers["x-api-key"] = api_key
if "x-api-key" not in headers and "authorization" not in headers:
auth_header = AnthropicModelInfo.get_auth_header(api_key)
if auth_header is not None:
headers.update(auth_header)
if "anthropic-version" not in headers:
headers["anthropic-version"] = DEFAULT_ANTHROPIC_API_VERSION
if "content-type" not in headers:

View file

@ -8,10 +8,8 @@ import httpx
import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
)
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.openai import (
FileContentRequest,
HttpxBinaryResponseContent,
@ -85,9 +83,9 @@ class AnthropicFilesHandler:
# Get Anthropic API credentials
api_base = self.anthropic_model_info.get_api_base(api_base)
api_key = api_key or self.anthropic_model_info.get_api_key()
auth_header = self.anthropic_model_info.get_auth_header(api_key)
if not api_key:
if auth_header is None:
raise ValueError("Missing Anthropic API Key")
# Construct the Anthropic batch results URL
@ -97,8 +95,8 @@ class AnthropicFilesHandler:
headers = {
"accept": "application/json",
"anthropic-version": "2023-06-01",
"x-api-key": api_key,
}
headers.update(auth_header)
# Make the request to Anthropic
async_client = get_async_httpx_client(llm_provider=LlmProviders.ANTHROPIC)

View file

@ -94,14 +94,14 @@ class AnthropicFilesConfig(BaseFilesConfig):
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
api_key = AnthropicModelInfo.get_api_key(api_key)
if not api_key:
auth_header = AnthropicModelInfo.get_auth_header(api_key)
if auth_header is None:
raise ValueError(
"Anthropic API key is required. Set ANTHROPIC_API_KEY environment variable or pass api_key parameter."
"Anthropic API key is required. Set ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN environment variable or pass api_key parameter."
)
headers.update(
{
"x-api-key": api_key,
**auth_header,
"anthropic-version": "2023-06-01",
"anthropic-beta": ANTHROPIC_FILES_BETA_HEADER,
}

View file

@ -35,17 +35,18 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig):
"""Add Anthropic-specific headers"""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
# Get API key
# Get API key from litellm_params if available
api_key = None
if litellm_params:
if litellm_params is not None:
api_key = litellm_params.api_key
api_key = AnthropicModelInfo.get_api_key(api_key)
if not api_key:
raise ValueError("ANTHROPIC_API_KEY is required for Skills API")
auth_header = AnthropicModelInfo.get_auth_header(api_key)
if auth_header is None:
raise ValueError(
"ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN is required for Skills API"
)
# Add required headers
headers["x-api-key"] = api_key
headers.update(auth_header)
headers["anthropic-version"] = "2023-06-01"
# Add beta header for skills API

View file

@ -4462,6 +4462,78 @@
"supports_vision": true,
"supports_web_search": true
},
"azure/gpt-5.4-mini": {
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_token": 7.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 4.5e-06,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true,
"supports_web_search": true,
"supports_none_reasoning_effort": false,
"supports_xhigh_reasoning_effort": false
},
"azure/gpt-5.4-nano": {
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.25e-06,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true,
"supports_web_search": true,
"supports_none_reasoning_effort": false,
"supports_xhigh_reasoning_effort": false
},
"azure/gpt-image-1": {
"cache_read_input_image_token_cost": 2.5e-06,
"cache_read_input_token_cost": 1.25e-06,
@ -37032,5 +37104,157 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true
},
"volcengine/doubao-seed-2-0-pro-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 4.6e-07,
"output_cost_per_token": 2.3e-06,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 7e-07,
"output_cost_per_token": 3.5e-06,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 1.4e-06,
"output_cost_per_token": 7e-06,
"range": [
128000.0,
256000.0
]
}
]
},
"volcengine/doubao-seed-2-0-lite-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 8.7e-08,
"output_cost_per_token": 5.2e-07,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 1.3e-07,
"output_cost_per_token": 7.8e-07,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 2.6e-07,
"output_cost_per_token": 1.6e-06,
"range": [
128000.0,
256000.0
]
}
]
},
"volcengine/doubao-seed-2-0-mini-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 2.9e-08,
"output_cost_per_token": 2.9e-07,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 5.8e-08,
"output_cost_per_token": 5.8e-07,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 1.2e-07,
"output_cost_per_token": 1.2e-06,
"range": [
128000.0,
256000.0
]
}
]
},
"volcengine/doubao-seed-2-0-code-preview-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 4.6e-07,
"output_cost_per_token": 2.3e-06,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 7e-07,
"output_cost_per_token": 3.5e-06,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 1.4e-06,
"output_cost_per_token": 7e-06,
"range": [
128000.0,
256000.0
]
}
]
}
}

View file

@ -208,26 +208,52 @@ async def exchange_token_with_server(
client_id: str,
client_secret: Optional[str],
code_verifier: Optional[str],
refresh_token: Optional[str] = None,
scope: Optional[str] = None,
):
if grant_type != "authorization_code":
if grant_type not in ("authorization_code", "refresh_token"):
raise HTTPException(status_code=400, detail="Unsupported grant_type")
if mcp_server.token_url is None:
raise HTTPException(status_code=400, detail="MCP server token url is not set")
proxy_base_url = get_request_base_url(request)
token_data = {
"grant_type": "authorization_code",
"client_id": mcp_server.client_id if mcp_server.client_id else client_id,
"client_secret": mcp_server.client_secret
if mcp_server.client_secret
else client_secret,
"code": code,
"redirect_uri": f"{proxy_base_url}/callback",
}
resolved_client_id = mcp_server.client_id if mcp_server.client_id else client_id
resolved_client_secret = (
mcp_server.client_secret if mcp_server.client_secret else client_secret
)
if code_verifier:
token_data["code_verifier"] = code_verifier
if grant_type == "refresh_token":
if not refresh_token:
raise HTTPException(
status_code=400,
detail="refresh_token is required for refresh_token grant",
)
token_data: dict = {
"grant_type": "refresh_token",
"refresh_token": refresh_token,
"client_id": resolved_client_id,
}
if resolved_client_secret is not None:
token_data["client_secret"] = resolved_client_secret
if scope:
token_data["scope"] = scope
else:
if not code:
raise HTTPException(
status_code=400,
detail="code is required for authorization_code grant",
)
proxy_base_url = get_request_base_url(request)
token_data = {
"grant_type": "authorization_code",
"client_id": resolved_client_id,
"code": code,
"redirect_uri": f"{proxy_base_url}/callback",
}
if resolved_client_secret is not None:
token_data["client_secret"] = resolved_client_secret
if code_verifier:
token_data["code_verifier"] = code_verifier
async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
response = await async_client.post(
@ -375,6 +401,8 @@ async def token_endpoint(
client_id: str = Form(...),
client_secret: Optional[str] = Form(None),
code_verifier: str = Form(None),
refresh_token: Optional[str] = Form(None),
scope: Optional[str] = Form(None),
mcp_server_name: Optional[str] = None,
):
"""
@ -408,6 +436,8 @@ async def token_endpoint(
client_id=client_id,
client_secret=client_secret,
code_verifier=code_verifier,
refresh_token=refresh_token,
scope=scope,
)

View file

@ -1,41 +1,32 @@
model_list:
- model_name: gpt-3.5-turbo
litellm_params:
model: openai/gpt-3.5-turbo
api_key: os.environ/OPENAI_API_KEY
- model_name: gpt-4o
litellm_params:
model: openai/gpt-4o
api_key: os.environ/OPENAI_API_KEY
- model_name: claude-sonnet-4-5-20250929
litellm_params:
model: anthropic/claude-sonnet-4-5-20250929
- model_name: gpt-4.1-mini
# OpenAI model for /v1/chat/completions test — 200x custom pricing
- model_name: "gpt-4.1-mini"
litellm_params:
model: openai/gpt-4.1-mini
- model_name: gpt-5-mini
api_key: os.environ/OPENAI_API_KEY
model_info:
id: gpt-4.1-mini-custom-pricing
input_cost_per_token: 0.00004 # 100x standard ($0.40/1M = $0.0000004)
output_cost_per_token: 0.00016 # 100x standard ($1.60/1M = $0.0000016)
# OpenAI model for /v1/responses test — 100x custom pricing
- model_name: "gpt-5"
litellm_params:
model: openai/gpt-5-mini
- model_name: custom_litellm_model
model: openai/gpt-5
api_key: os.environ/OPENAI_API_KEY
model_info:
id: gpt-5-custom-pricing
mode: "chat"
input_cost_per_token: 125 # 100x standard ($1.25/1M = $0.00000125)
output_cost_per_token: 10 # 100x standard ($10.00/1M = $0.00001)
# Anthropic model for /v1/messages test — 100x custom pricing
- model_name: "claude-sonnet-4-20250514"
litellm_params:
model: litellm_agent/claude-sonnet-4-5-20250929
litellm_system_prompt: "Be a helpful assistant."
guardrails:
- guardrail_name: "tool_policy"
litellm_params:
guardrail: tool_policy
mode: [pre_call, post_call]
default_on: true
mcp_servers:
my_http_server:
url: "http://0.0.0.0:8001/mcp"
transport: "http"
description: "My custom MCP server"
available_on_public_internet: true
general_settings:
store_model_in_db: true
store_prompts_in_spend_logs: true
model: anthropic/claude-sonnet-4-20250514
api_key: os.environ/ANTHROPIC_API_KEY
model_info:
id: claude-sonnet-4-custom-pricing
input_cost_per_token: 0.0003 # 100x standard ($0.000003)
output_cost_per_token: 0.0015 # 100x standard ($0.000015)

View file

@ -1672,6 +1672,7 @@ class NewTeamRequest(TeamBase):
int
] = None # allow user to set TPM limit for all team members
team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m"
team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo"
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
enforced_batch_output_expires_after: Optional[dict] = None
enforced_file_expires_after: Optional[dict] = None
@ -2955,7 +2956,9 @@ class LiteLLM_ErrorLogs(LiteLLMPydanticObjectBase):
endTime: Union[str, datetime, None]
AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked", "rotated"]
AUDIT_ACTIONS = Literal[
"created", "updated", "deleted", "blocked", "unblocked", "rotated"
]
class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase):

View file

@ -29,6 +29,7 @@ from litellm.constants import (
DEFAULT_MAX_RECURSE_DEPTH,
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.proxy._types import (
RBAC_ROLES,
@ -407,18 +408,21 @@ async def common_checks( # noqa: PLR0915
# 2. If team can call model
if _model and team_object:
if not await can_team_access_model(
model=_model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=valid_token.team_model_aliases if valid_token else None,
):
raise ProxyException(
message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}",
type=ProxyErrorTypes.team_model_access_denied,
param="model",
code=status.HTTP_401_UNAUTHORIZED,
)
with tracer.trace("litellm.proxy.auth.common_checks.can_team_access_model"):
if not await can_team_access_model(
model=_model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=valid_token.team_model_aliases
if valid_token
else None,
):
raise ProxyException(
message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}",
type=ProxyErrorTypes.team_model_access_denied,
param="model",
code=status.HTTP_401_UNAUTHORIZED,
)
# Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent
if valid_token is not None and valid_token.agent_id:
@ -443,54 +447,62 @@ async def common_checks( # noqa: PLR0915
## 2.1 If user can call model (if personal key)
if _model and team_object is None and user_object is not None:
await can_user_call_model(
model=_model,
llm_router=llm_router,
user_object=user_object,
)
with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"):
await can_user_call_model(
model=_model,
llm_router=llm_router,
user_object=user_object,
)
# 1.1 - 2.2 - 3.0.2 - 3.0.3: Project checks (blocked, model access, budget)
await _run_project_checks(
project_object=project_object,
_model=_model,
llm_router=llm_router,
skip_budget_checks=skip_budget_checks,
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace("litellm.proxy.auth.common_checks.run_project_checks"):
await _run_project_checks(
project_object=project_object,
_model=_model,
llm_router=llm_router,
skip_budget_checks=skip_budget_checks,
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
# If this is a free model, skip all budget checks
if not skip_budget_checks:
# 3. If team is in budget
await _team_max_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
with tracer.trace("litellm.proxy.auth.common_checks.team_max_budget_check"):
await _team_max_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
# 3.0.5. If team is over soft budget (alert only, doesn't block)
await _team_soft_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
with tracer.trace("litellm.proxy.auth.common_checks.team_soft_budget_check"):
await _team_soft_budget_check(
team_object=team_object,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
# 3.1. If organization is in budget
await _organization_max_budget_check(
valid_token=valid_token,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace(
"litellm.proxy.auth.common_checks.organization_max_budget_check"
):
await _organization_max_budget_check(
valid_token=valid_token,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await _tag_max_budget_check(
request_body=request_body,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
with tracer.trace("litellm.proxy.auth.common_checks.tag_max_budget_check"):
await _tag_max_budget_check(
request_body=request_body,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
# 4. If user is in budget
## 4.1 check personal budget, if personal key
@ -508,14 +520,15 @@ async def common_checks( # noqa: PLR0915
)
## 4.2 check team member budget, if team key
await _check_team_member_budget(
team_object=team_object,
user_object=user_object,
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace("litellm.proxy.auth.common_checks.check_team_member_budget"):
await _check_team_member_budget(
team_object=team_object,
user_object=user_object,
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
# 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget
if (
@ -554,19 +567,21 @@ async def common_checks( # noqa: PLR0915
)
# 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store
await vector_store_access_check(
request_body=request_body,
team_object=team_object,
valid_token=valid_token,
)
with tracer.trace("litellm.proxy.auth.common_checks.vector_store_access_check"):
await vector_store_access_check(
request_body=request_body,
team_object=team_object,
valid_token=valid_token,
)
# 12. [OPTIONAL] Tool allowlist - key/team allowed_tools (no DB in hot path)
await check_tools_allowlist(
request_body=request_body,
valid_token=valid_token,
team_object=team_object,
route=route,
)
with tracer.trace("litellm.proxy.auth.common_checks.check_tools_allowlist"):
await check_tools_allowlist(
request_body=request_body,
valid_token=valid_token,
team_object=team_object,
route=route,
)
return True

View file

@ -548,13 +548,12 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
custom_auth_api_key: bool = False
try:
# get the request body
await pre_db_read_auth_checks(
request_data=request_data,
request=request,
route=route,
)
with tracer.trace("litellm.proxy.auth.pre_db_read_auth_checks"):
await pre_db_read_auth_checks(
request_data=request_data,
request=request,
route=route,
)
pass_through_endpoints: Optional[List[dict]] = general_settings.get(
"pass_through_endpoints", None
)
@ -588,9 +587,10 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
### USER-DEFINED AUTH FUNCTION ###
if enterprise_custom_auth is not None:
response = await enterprise_custom_auth(
request=request, api_key=api_key, user_custom_auth=user_custom_auth
)
with tracer.trace("litellm.proxy.auth.enterprise_custom_auth"):
response = await enterprise_custom_auth(
request=request, api_key=api_key, user_custom_auth=user_custom_auth
)
if response is not None and isinstance(response, UserAPIKeyAuth):
validated = UserAPIKeyAuth.model_validate(response)
validated = await _run_post_custom_auth_checks(
@ -706,18 +706,19 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
# Fall through to virtual key checks
if do_standard_jwt_auth:
result = await JWTAuthManager.auth_builder(
request_data=request_data,
general_settings=general_settings,
api_key=api_key,
jwt_handler=jwt_handler,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
parent_otel_span=parent_otel_span,
request_headers=_safe_get_request_headers(request),
)
with tracer.trace("litellm.proxy.auth.jwt_auth_builder"):
result = await JWTAuthManager.auth_builder(
request_data=request_data,
general_settings=general_settings,
api_key=api_key,
jwt_handler=jwt_handler,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
parent_otel_span=parent_otel_span,
request_headers=_safe_get_request_headers(request),
)
is_proxy_admin = result["is_proxy_admin"]
team_id = result["team_id"]
@ -909,15 +910,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
try:
end_user_params["end_user_id"] = end_user_id
# get end-user object
_end_user_object = await get_end_user_object(
end_user_id=end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
)
with tracer.trace("litellm.proxy.auth.get_end_user_object"):
_end_user_object = await get_end_user_object(
end_user_id=end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
)
if _end_user_object is not None:
end_user_params[
"allowed_model_region"
@ -960,14 +961,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if valid_token is None:
## Check CACHE
try:
valid_token = await get_key_object(
hashed_token=hash_token(api_key),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
check_cache_only=True,
)
with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"):
valid_token = await get_key_object(
hashed_token=hash_token(api_key),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
check_cache_only=True,
)
except Exception:
verbose_logger.debug("api key not found in cache.")
valid_token = None
@ -1139,13 +1141,14 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
api_key = hash_token(token=api_key)
try:
valid_token = await get_key_object(
hashed_token=api_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace("litellm.proxy.auth.get_key_object_from_db"):
valid_token = await get_key_object(
hashed_token=api_key,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
except ProxyException as e:
if e.code == 401 or e.code == "401":
e.message = "Authentication Error, Invalid proxy server token passed. Received API Key = {}, Key Hash (Token) ={}. Unable to find token in cache or `LiteLLM_VerificationTokenTable`".format(
@ -1233,14 +1236,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
# Check 2. If user_id for this token is in budget - done in common_checks()
if valid_token.user_id is not None:
try:
user_obj = await get_user_object(
user_id=valid_token.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace("litellm.proxy.auth.get_user_object"):
user_obj = await get_user_object(
user_id=valid_token.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
except Exception as e:
verbose_logger.debug(
"litellm.proxy.auth.user_api_key_auth.py::user_api_key_auth() - Unable to get user from db/cache. Setting user_obj to None. Exception received - {}".format(
@ -1329,71 +1333,73 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
if not skip_budget_checks:
# Check 4. Token Spend is under budget
if RouteChecks.is_llm_api_route(route=route):
await _virtual_key_max_budget_check(
with tracer.trace("litellm.proxy.auth.budget_checks"):
# Check 4. Token Spend is under budget
if RouteChecks.is_llm_api_route(route=route):
await _virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 5. Max Budget Alert Check
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 5. Max Budget Alert Check
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 6. Soft Budget Check
await _virtual_key_soft_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 5. Token Model Spend is under Model budget
max_budget_per_model = valid_token.model_max_budget
current_model = request_data.get("model", None)
if (
max_budget_per_model is not None
and isinstance(max_budget_per_model, dict)
and len(max_budget_per_model) > 0
and prisma_client is not None
and current_model is not None
and valid_token.token is not None
):
## GET THE SPEND FOR THIS MODEL
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=current_model,
# Check 6. Soft Budget Check
await _virtual_key_soft_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 5b. End-user model max budget
end_user_mmb = valid_token.end_user_model_max_budget
if (
end_user_mmb is not None
and isinstance(end_user_mmb, dict)
and len(end_user_mmb) > 0
and current_model is not None
and valid_token.end_user_id is not None
):
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=current_model,
)
# Check 5. Token Model Spend is under Model budget
max_budget_per_model = valid_token.model_max_budget
current_model = request_data.get("model", None)
if (
max_budget_per_model is not None
and isinstance(max_budget_per_model, dict)
and len(max_budget_per_model) > 0
and prisma_client is not None
and current_model is not None
and valid_token.token is not None
):
## GET THE SPEND FOR THIS MODEL
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=current_model,
)
# Check 5b. End-user model max budget
end_user_mmb = valid_token.end_user_model_max_budget
if (
end_user_mmb is not None
and isinstance(end_user_mmb, dict)
and len(end_user_mmb) > 0
and current_model is not None
and valid_token.end_user_id is not None
):
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=current_model,
)
# Check 6: Additional Common Checks across jwt + key auth
if valid_token.team_id is not None:
try:
_team_obj = await get_team_object(
team_id=valid_token.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
with tracer.trace("litellm.proxy.auth.get_team_object"):
_team_obj = await get_team_object(
team_id=valid_token.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
except HTTPException:
_team_obj = LiteLLM_TeamTableCachedObj(
team_id=valid_token.team_id,
@ -1431,11 +1437,14 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
litellm.max_budget > 0 and prisma_client is not None
): # user set proxy max budget
cache_key = "{}:spend".format(litellm_proxy_admin_name)
global_proxy_spend = await _fetch_global_spend_with_event_coordination(
cache_key=cache_key,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"):
global_proxy_spend = (
await _fetch_global_spend_with_event_coordination(
cache_key=cache_key,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
)
if global_proxy_spend is not None:
call_info = CallInfo(
@ -1452,21 +1461,22 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
user_info=call_info,
)
)
_ = await common_checks(
request=request,
request_body=request_data,
team_object=_team_obj,
user_object=user_obj,
end_user_object=_end_user_object,
general_settings=general_settings,
global_proxy_spend=global_proxy_spend,
route=route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
skip_budget_checks=skip_budget_checks,
project_object=_project_obj,
)
with tracer.trace("litellm.proxy.auth.common_checks"):
_ = await common_checks(
request=request,
request_body=request_data,
team_object=_team_obj,
user_object=user_obj,
end_user_object=_end_user_object,
general_settings=general_settings,
global_proxy_spend=global_proxy_spend,
route=route,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
skip_budget_checks=skip_budget_checks,
project_object=_project_obj,
)
# Token passed all checks
if valid_token is None:
raise HTTPException(401, detail="Invalid API key")

View file

@ -1260,7 +1260,9 @@ class ProxyBaseLLMRequestProcessing:
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=(
_litellm_logging_obj.litellm_call_id if _litellm_logging_obj else None
_litellm_logging_obj.litellm_call_id
if _litellm_logging_obj
else self.data.get("litellm_call_id")
),
model_id=model_id,
version=version,

View file

@ -27,10 +27,17 @@ async def get_ui_config():
admin_ui_disabled = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true"
sso_configured = _has_user_setup_sso()
from litellm.proxy.proxy_server import proxy_config
is_control_plane = len(proxy_config.worker_registry) > 0
return UiDiscoveryEndpoints(
server_root_path=get_server_root_path(),
proxy_base_url=get_proxy_base_url(),
auto_redirect_to_sso=sso_configured and auto_redirect_ui_login_to_sso,
admin_ui_disabled=admin_ui_disabled,
sso_configured=sso_configured,
is_control_plane=is_control_plane,
workers=proxy_config.worker_registry if is_control_plane else [],
)

View file

@ -41,10 +41,8 @@ from litellm.proxy._experimental.mcp_server.db import (
from litellm.proxy._types import *
from litellm.proxy._types import LiteLLM_VerificationToken
from litellm.proxy.auth.auth_checks import (
_cache_key_object,
_delete_cache_key_object,
can_team_access_model,
get_key_object,
get_org_object,
get_project_object,
get_team_object,
@ -1656,7 +1654,7 @@ async def _get_and_validate_existing_key(
LiteLLM_VerificationToken: The existing key row
Raises:
HTTPException: If key is not found
ProxyException: 404 if key is not found
"""
if prisma_client is None:
raise HTTPException(
@ -1664,16 +1662,18 @@ async def _get_and_validate_existing_key(
detail={"error": "Database not connected"},
)
existing_key_row = await prisma_client.get_data(
token=token,
table_name="key",
query_type="find_unique",
hashed_token = _hash_token_if_needed(token=token)
existing_key_row = await prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": hashed_token}
)
if existing_key_row is None:
raise HTTPException(
status_code=404,
detail={"error": f"Key not found: {token}"},
raise ProxyException(
message="Key not found.",
type=ProxyErrorTypes.not_found_error,
param="key",
code=status.HTTP_404_NOT_FOUND,
)
return existing_key_row
@ -2111,19 +2111,11 @@ async def update_key_fn(
key = data_json.pop("key")
# get the row from db
if prisma_client is None:
raise Exception("Not connected to DB!")
existing_key_row = await prisma_client.get_data(
token=data.key, table_name="key", query_type="find_unique"
existing_key_row = await _get_and_validate_existing_key(
token=data.key,
prisma_client=prisma_client,
)
if existing_key_row is None:
raise HTTPException(
status_code=404,
detail={"error": f"Team not found, passed team_id={data.team_id}"},
)
await _validate_update_key_data(
data=data,
existing_key_row=existing_key_row,
@ -2158,6 +2150,8 @@ async def update_key_fn(
)
_data = {**non_default_values, "token": key}
if prisma_client is None:
raise Exception("Not connected to DB!")
response = await prisma_client.update_data(token=key, data=_data)
# Delete - key from cache, since it's been updated!
@ -2330,6 +2324,8 @@ async def bulk_update_keys(
error_message = error_detail.get("error", str(e))
else:
error_message = str(error_detail)
elif isinstance(e, ProxyException):
error_message = e.message
else:
error_message = str(e)
@ -4945,18 +4941,19 @@ async def block_key(
route="/key/block",
)
if litellm.store_audit_logs is True:
# make an audit log for key update
record = await prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": hashed_token}
# Check if the key exists before trying to block it
existing_record = await prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": hashed_token}
)
if existing_record is None:
raise ProxyException(
message="Key not found.",
type=ProxyErrorTypes.not_found_error,
param="key",
code=status.HTTP_404_NOT_FOUND,
)
if record is None:
raise ProxyException(
message=f"Key {data.key} not found",
type=ProxyErrorTypes.bad_request_error,
param="key",
code=status.HTTP_404_NOT_FOUND,
)
if litellm.store_audit_logs is True:
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
@ -4970,7 +4967,7 @@ async def block_key(
object_id=hashed_token,
action="blocked",
updated_values="{}",
before_value=record.model_dump_json(),
before_value=existing_record.model_dump_json(),
)
)
)
@ -4979,24 +4976,9 @@ async def block_key(
where={"token": hashed_token}, data={"blocked": True} # type: ignore
)
## UPDATE KEY CACHE
### get cached object ###
key_object = await get_key_object(
## UPDATE KEY CACHE - invalidate so next read re-fetches from DB
await _delete_cache_key_object(
hashed_token=hashed_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
### update cached object ###
key_object.blocked = True
### store cached object ###
await _cache_key_object(
hashed_token=hashed_token,
user_api_key_obj=key_object,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
@ -5068,18 +5050,19 @@ async def unblock_key(
route="/key/unblock",
)
if litellm.store_audit_logs is True:
# make an audit log for key update
record = await prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": hashed_token}
# Check if the key exists before trying to unblock it
existing_record = await prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": hashed_token}
)
if existing_record is None:
raise ProxyException(
message="Key not found.",
type=ProxyErrorTypes.not_found_error,
param="key",
code=status.HTTP_404_NOT_FOUND,
)
if record is None:
raise ProxyException(
message=f"Key {data.key} not found",
type=ProxyErrorTypes.bad_request_error,
param="key",
code=status.HTTP_404_NOT_FOUND,
)
if litellm.store_audit_logs is True:
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
@ -5091,9 +5074,9 @@ async def unblock_key(
changed_by_api_key=user_api_key_dict.api_key,
table_name=LitellmTableNames.KEY_TABLE_NAME,
object_id=hashed_token,
action="blocked",
action="unblocked",
updated_values="{}",
before_value=record.model_dump_json(),
before_value=existing_record.model_dump_json(),
)
)
)
@ -5102,24 +5085,9 @@ async def unblock_key(
where={"token": hashed_token}, data={"blocked": False} # type: ignore
)
## UPDATE KEY CACHE
### get cached object ###
key_object = await get_key_object(
## UPDATE KEY CACHE - invalidate so next read re-fetches from DB
await _delete_cache_key_object(
hashed_token=hashed_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
)
### update cached object ###
key_object.blocked = False
### store cached object ###
await _cache_key_object(
hashed_token=hashed_token,
user_api_key_obj=key_object,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)

View file

@ -1214,17 +1214,9 @@ if MCP_AVAILABLE:
"error": "User does not have permission to create mcp servers. You can only create mcp servers if you are a PROXY_ADMIN."
},
)
elif payload.server_id is not None:
# fail if the mcp server with id already exists
mcp_server = await get_mcp_server(prisma_client, payload.server_id)
if mcp_server is not None:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."
},
)
elif (
# Block reserved special server IDs
if (
SpecialMCPServerName.all_team_servers == payload.server_id
or SpecialMCPServerName.all_proxy_servers == payload.server_id
):
@ -1235,6 +1227,17 @@ if MCP_AVAILABLE:
},
)
if payload.server_id is not None:
# fail if the mcp server with id already exists
mcp_server = await get_mcp_server(prisma_client, payload.server_id)
if mcp_server is not None:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."
},
)
# TODO: audit log for create
# Admin-created servers are always active — clear any submission lifecycle
@ -1399,6 +1402,8 @@ if MCP_AVAILABLE:
client_id: Optional[str] = Form(None),
client_secret: Optional[str] = Form(None),
code_verifier: Optional[str] = Form(None),
refresh_token: Optional[str] = Form(None),
scope: Optional[str] = Form(None),
):
mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
resolved_client_id = mcp_server.client_id or client_id or ""
@ -1422,6 +1427,8 @@ if MCP_AVAILABLE:
client_id=resolved_client_id,
client_secret=client_secret,
code_verifier=code_verifier,
refresh_token=refresh_token,
scope=scope,
)
@router.post(

View file

@ -724,6 +724,7 @@ async def new_team( # noqa: PLR0915
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
- team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
- team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
- team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
- team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
- team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
@ -934,6 +935,7 @@ async def new_team( # noqa: PLR0915
team_member_budget=data.team_member_budget,
team_member_rpm_limit=data.team_member_rpm_limit,
team_member_tpm_limit=data.team_member_tpm_limit,
team_member_budget_duration=data.team_member_budget_duration,
):
data_json = await TeamMemberBudgetHandler.create_team_member_budget_table(
data=data,
@ -942,6 +944,7 @@ async def new_team( # noqa: PLR0915
team_member_budget=data.team_member_budget,
team_member_rpm_limit=data.team_member_rpm_limit,
team_member_tpm_limit=data.team_member_tpm_limit,
team_member_budget_duration=data.team_member_budget_duration,
)
## ADD TO TEAM TABLE
@ -3334,7 +3337,9 @@ def _convert_teams_to_response_models(
use_deleted_table: bool,
) -> List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]:
"""Convert raw Prisma team rows to response models."""
team_list: List[Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]] = []
team_list: List[
Union[TeamListItem, LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]
] = []
for team in teams:
try:
team_dict = team.model_dump()

View file

@ -16,6 +16,7 @@ import os
import secrets
from copy import deepcopy
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
from urllib.parse import urlencode, urlparse
if TYPE_CHECKING:
import httpx
@ -301,6 +302,7 @@ async def google_login(
source: Optional[str] = None,
key: Optional[str] = None,
existing_key: Optional[str] = None,
return_to: Optional[str] = None,
): # noqa: PLR0915
"""
Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env
@ -394,13 +396,23 @@ async def google_login(
is True
):
verbose_proxy_logger.info(f"Redirecting to SSO login for {redirect_url}")
return await SSOAuthenticationHandler.get_sso_login_redirect(
sso_redirect = await SSOAuthenticationHandler.get_sso_login_redirect(
redirect_url=redirect_url,
microsoft_client_id=microsoft_client_id,
google_client_id=google_client_id,
generic_client_id=generic_client_id,
state=cli_state,
)
if return_to is not None and sso_redirect is not None:
SSOAuthenticationHandler._validate_return_to(return_to)
sso_redirect.set_cookie(
key="litellm_cp_return_to",
value=return_to,
max_age=600,
httponly=True,
samesite="lax",
)
return sso_redirect
elif ui_username is not None:
# No Google, Microsoft SSO
# Use UI Credentials set in .env
@ -1312,12 +1324,17 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa:
request=request, key=key_id, existing_key=existing_key, result=result
)
# Control-plane cross-origin: read return_to from cookie.
# Starlette's cookie_parser already handles RFC 2109 unquoting.
cp_return_to: Optional[str] = request.cookies.get("litellm_cp_return_to")
return await SSOAuthenticationHandler.get_redirect_response_from_openid(
result=result,
request=request,
received_response=received_response,
generic_client_id=generic_client_id,
ui_access_mode=ui_access_mode,
return_to=cp_return_to,
)
@ -1760,6 +1777,38 @@ class SSOAuthenticationHandler:
Handler for SSO Authentication across all SSO providers
"""
@staticmethod
def _validate_return_to(return_to: str) -> None:
"""
Validate that return_to matches the configured control_plane_url origin.
Raises HTTPException(400) if:
- control_plane_url is not configured in general_settings
- return_to origin does not match control_plane_url origin
"""
from litellm.proxy.proxy_server import general_settings
control_plane_url = general_settings.get("control_plane_url")
if control_plane_url is None:
raise HTTPException(
status_code=400,
detail="return_to is not allowed: control_plane_url is not configured",
)
def _origin(url: str) -> tuple:
parsed = urlparse(url)
scheme = (parsed.scheme or "").lower()
hostname = (parsed.hostname or "").lower()
default_port = 443 if scheme == "https" else 80
port = parsed.port if parsed.port is not None else default_port
return (scheme, hostname, port)
if _origin(return_to) != _origin(control_plane_url):
raise HTTPException(
status_code=400,
detail="return_to does not match the configured control_plane_url",
)
@staticmethod
async def get_sso_login_redirect(
redirect_url: str,
@ -2358,6 +2407,7 @@ class SSOAuthenticationHandler:
received_response: Optional[dict] = None,
generic_client_id: Optional[str] = None,
ui_access_mode: Optional[Dict] = None,
return_to: Optional[str] = None,
) -> RedirectResponse:
import jwt
@ -2367,6 +2417,7 @@ class SSOAuthenticationHandler:
master_key,
premium_user,
proxy_logging_obj,
redis_usage_cache,
user_api_key_cache,
user_custom_sso,
)
@ -2534,6 +2585,36 @@ class SSOAuthenticationHandler:
master_key or "",
algorithm="HS256",
)
# Control-plane cross-origin: store JWT behind a single-use opaque
# code (60s TTL) so the token never appears in browser history / logs.
# The control plane redeems it via POST /v3/login/exchange.
if return_to is not None:
SSOAuthenticationHandler._validate_return_to(return_to)
code = secrets.token_urlsafe(32)
cache_key = f"login_code:{code}"
cache_value = {"token": jwt_token, "redirect_url": return_to}
if redis_usage_cache is not None:
await redis_usage_cache.async_set_cache(
key=cache_key, value=cache_value, ttl=60
)
else:
await user_api_key_cache.async_set_cache(
key=cache_key, value=cache_value, ttl=60
)
separator = "&" if "?" in return_to else "?"
redirect_url = (
return_to + separator + urlencode({"login": "success", "code": code})
)
verbose_proxy_logger.info(
"Cross-origin SSO: redirecting to control plane with login code"
)
redirect_response = RedirectResponse(url=redirect_url, status_code=303)
redirect_response.delete_cookie("litellm_cp_return_to")
return redirect_response
if user_id is not None and isinstance(user_id, str):
litellm_dashboard_ui += "?login=success"
verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}")

View file

@ -22,6 +22,7 @@ from litellm.constants import (
ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS,
BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES,
)
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.proxy._types import *
from litellm.proxy.auth.route_checks import RouteChecks
@ -585,7 +586,11 @@ async def anthropic_proxy_route(
"""
[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)
"""
base_target_url = os.getenv("ANTHROPIC_API_BASE") or "https://api.anthropic.com"
base_target_url = (
os.getenv("ANTHROPIC_API_BASE")
or os.getenv("ANTHROPIC_BASE_URL")
or "https://api.anthropic.com"
)
encoded_endpoint = httpx.URL(endpoint).path
# Ensure endpoint starts with '/' for proper URL construction
@ -606,10 +611,11 @@ async def anthropic_proxy_route(
is_streaming_request = await is_streaming_request_fn(request)
## CREATE PASS-THROUGH
auth_header = AnthropicModelInfo.get_auth_header(anthropic_api_key or None)
endpoint_func = create_pass_through_route(
endpoint=endpoint,
target=str(updated_url),
custom_headers={"x-api-key": "{}".format(anthropic_api_key)},
custom_headers=auth_header if auth_header is not None else {},
_forward_headers=True,
is_streaming_request=is_streaming_request,
) # dynamically construct pass-through endpoint based on incoming path

View file

@ -7,6 +7,7 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
from litellm.llms.anthropic import get_anthropic_config
from litellm.llms.anthropic.chat.handler import (
ModelResponseIterator as AnthropicModelResponseIterator,
@ -124,10 +125,21 @@ class AnthropicPassthroughLoggingHandler:
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
model_for_cost = f"{custom_llm_provider}/{model}"
router_model_id = logging_obj.get_router_model_id()
custom_pricing = use_custom_pricing_for_model(
litellm_params=(
logging_obj.litellm_params
if hasattr(logging_obj, "litellm_params")
else None
)
)
response_cost = litellm.completion_cost(
completion_response=litellm_model_response,
model=model_for_cost,
custom_llm_provider=custom_llm_provider,
custom_pricing=custom_pricing,
router_model_id=router_model_id,
)
kwargs["response_cost"] = response_cost
@ -319,9 +331,7 @@ class AnthropicPassthroughLoggingHandler:
import base64
from litellm._uuid import uuid
from litellm.llms.anthropic.batches.transformation import (
AnthropicBatchesConfig,
)
from litellm.llms.anthropic.batches.transformation import AnthropicBatchesConfig
from litellm.types.utils import Choices, SpecialEnums
try:

View file

@ -12,6 +12,7 @@ import click
import httpx
from dotenv import load_dotenv
import litellm
from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY
from litellm.secret_managers.main import get_secret_bool
@ -387,7 +388,7 @@ class ProxyInitializationHelpers:
@click.option("--api_base", default=None, help="API base URL.")
@click.option(
"--api_version",
default="2024-07-01-preview",
default=litellm.AZURE_DEFAULT_API_VERSION,
help="For azure - pass in the api version.",
)
@click.option(

View file

@ -541,6 +541,7 @@ from litellm.types.llms.anthropic import (
AnthropicResponseUsageBlock,
)
from litellm.types.llms.openai import HttpxBinaryResponseContent
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
ModelGroupInfoProxy,
)
@ -1546,6 +1547,7 @@ user_custom_key_generate = None
# Sentinel: prevents PKCE-no-Redis advisory from re-logging on config hot-reload.
# Tests that need to reset it can patch 'litellm.proxy.proxy_server._pkce_no_redis_warning_emitted'.
_pkce_no_redis_warning_emitted: bool = False
_cp_no_redis_warning_emitted: bool = False
user_custom_sso = None
user_custom_ui_sso_sign_in_handler = None
use_background_health_checks = None
@ -2295,6 +2297,7 @@ class ProxyConfig:
self.config: Dict[str, Any] = {}
self._last_semantic_filter_config: Optional[Dict[str, Any]] = None
self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None
self.worker_registry: List["WorkerRegistryEntry"] = []
def is_yaml(self, config_file_path: str) -> bool:
if not os.path.isfile(config_file_path):
@ -3095,6 +3098,21 @@ class ProxyConfig:
"Set PKCE_STRICT_CACHE_MISS=true to fail fast with a 401 on cache misses "
"instead of continuing without a code_verifier."
)
### CONTROL PLANE CODE-EXCHANGE PREREQUISITE CHECK ###
cp_url = general_settings.get("control_plane_url")
if cp_url and redis_usage_cache is None:
global _cp_no_redis_warning_emitted
if not _cp_no_redis_warning_emitted:
_cp_no_redis_warning_emitted = True
verbose_proxy_logger.warning(
"control_plane_url is configured but Redis is not configured for LiteLLM caching. "
"Login codes (SSO and /v3/login) will not be shared across instances — "
"the /v3/login/exchange call may land on a different pod and fail with 401. "
"Configure Redis via the 'cache' section in your proxy config, "
"or ensure sticky sessions for single-instance deployments."
)
### STORE MODEL IN DB ### feature flag for `/model/new`
store_model_in_db = general_settings.get("store_model_in_db", False)
if store_model_in_db is None:
@ -3385,7 +3403,15 @@ class ProxyConfig:
litellm.vector_store_registry.load_vector_stores_from_config(
vector_store_registry_config
)
pass
## WORKER REGISTRY (Control Plane)
worker_registry_config = config.get("worker_registry", None)
if worker_registry_config:
self.worker_registry = [
WorkerRegistryEntry(**e) for e in worker_registry_config
]
else:
self.worker_registry = []
async def _init_policy_engine(
self,
@ -11095,6 +11121,165 @@ async def login_v2(request: Request): # noqa: PLR0915
)
@router.post(
"/v3/login", include_in_schema=False
) # control-plane login — always returns token in body for cross-origin use
async def login_v3(request: Request): # noqa: PLR0915
global premium_user, general_settings, master_key
from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object
from litellm.proxy.utils import get_custom_url
try:
if not general_settings.get("control_plane_url"):
raise ProxyException(
message="/v3/login is only available on workers with control_plane_url configured",
type=ProxyErrorTypes.not_found_error,
param="control_plane_url",
code=status.HTTP_404_NOT_FOUND,
)
body = await request.json()
username = str(body.get("username"))
password = str(body.get("password"))
login_result = await authenticate_user(
username=username,
password=password,
master_key=master_key,
prisma_client=prisma_client,
)
returned_ui_token_object = create_ui_token_object(
login_result=login_result,
general_settings=general_settings,
premium_user=premium_user,
)
import jwt
jwt_token = jwt.encode(
cast(dict, returned_ui_token_object),
cast(str, master_key),
algorithm="HS256",
)
litellm_dashboard_ui = get_custom_url(str(request.base_url))
if litellm_dashboard_ui.endswith("/"):
litellm_dashboard_ui += "ui/"
else:
litellm_dashboard_ui += "/ui/"
litellm_dashboard_ui += "?login=success"
# Store JWT behind a single-use opaque code (60s TTL)
code = secrets.token_urlsafe(32)
cache_key = f"login_code:{code}"
cache_value = {"token": jwt_token, "redirect_url": litellm_dashboard_ui}
if redis_usage_cache is not None:
await redis_usage_cache.async_set_cache(
key=cache_key, value=cache_value, ttl=60
)
else:
await user_api_key_cache.async_set_cache(
key=cache_key, value=cache_value, ttl=60
)
return JSONResponse(
content={"code": code, "expires_in": 60},
status_code=status.HTTP_200_OK,
)
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.login_v3(): Exception occurred - {}".format(
str(e)
)
)
if isinstance(e, ProxyException):
raise e
elif isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", str(e)),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
)
else:
error_msg = f"{str(e)}"
raise ProxyException(
message=error_msg,
type=ProxyErrorTypes.auth_error,
param="None",
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
@router.post(
"/v3/login/exchange", include_in_schema=False
) # exchange single-use opaque code for JWT
async def login_v3_exchange(request: Request):
try:
if not general_settings.get("control_plane_url"):
raise ProxyException(
message="/v3/login/exchange is only available on workers with control_plane_url configured",
type=ProxyErrorTypes.not_found_error,
param="control_plane_url",
code=status.HTTP_404_NOT_FOUND,
)
body = await request.json()
code = body.get("code")
if not code:
raise ProxyException(
message="Missing 'code' parameter",
type=ProxyErrorTypes.auth_error,
param="code",
code=status.HTTP_400_BAD_REQUEST,
)
cache_key = f"login_code:{code}"
if redis_usage_cache is not None:
cached_data = await redis_usage_cache.async_get_cache(key=cache_key)
else:
cached_data = await user_api_key_cache.async_get_cache(key=cache_key)
if not cached_data or not isinstance(cached_data, dict):
raise ProxyException(
message="Invalid or expired login code",
type=ProxyErrorTypes.auth_error,
param="code",
code=status.HTTP_401_UNAUTHORIZED,
)
# Single-use: delete immediately
if redis_usage_cache is not None:
await redis_usage_cache.async_delete_cache(key=cache_key)
else:
await user_api_key_cache.async_delete_cache(key=cache_key)
json_response = JSONResponse(
content={
"token": cached_data["token"],
"redirect_url": cached_data["redirect_url"],
},
status_code=status.HTTP_200_OK,
)
json_response.set_cookie(key="token", value=cached_data["token"])
return json_response
except ProxyException:
raise
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.login_v3_exchange(): Exception occurred - {}".format(
str(e)
)
)
raise ProxyException(
message=str(e),
type=ProxyErrorTypes.auth_error,
param="None",
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
@app.get("/onboarding/get_token", include_in_schema=False)
async def onboarding(invite_link: str, request: Request):
"""

View file

@ -1878,28 +1878,33 @@ class ProxyLogging:
)
input: Union[list, str, dict] = ""
normalized_call_type: Optional[str] = None
if "messages" in request_data and isinstance(
request_data["messages"], list
):
input = request_data["messages"]
litellm_logging_obj.model_call_details["messages"] = input
if litellm_logging_obj.call_type != CallTypes.pass_through.value:
litellm_logging_obj.call_type = CallTypes.acompletion.value
normalized_call_type = CallTypes.acompletion.value
elif "prompt" in request_data and isinstance(request_data["prompt"], str):
input = request_data["prompt"]
litellm_logging_obj.model_call_details["prompt"] = input
if litellm_logging_obj.call_type != CallTypes.pass_through.value:
litellm_logging_obj.call_type = CallTypes.atext_completion.value
normalized_call_type = CallTypes.atext_completion.value
elif "input" in request_data and isinstance(request_data["input"], list):
input = request_data["input"]
litellm_logging_obj.model_call_details["input"] = input
if litellm_logging_obj.call_type != CallTypes.pass_through.value:
litellm_logging_obj.call_type = CallTypes.aembedding.value
normalized_call_type = CallTypes.aembedding.value
if normalized_call_type is not None:
litellm_logging_obj.call_type = normalized_call_type
litellm_logging_obj.model_call_details[
"call_type"
] = normalized_call_type
# Pass-through endpoints are logged via the callback loop's
# async_post_call_failure_hook — skip pre_call and failure handlers.
if litellm_logging_obj.call_type == CallTypes.pass_through.value:
return
litellm_logging_obj.pre_call(
input=input,
api_key="",

View file

@ -107,6 +107,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._reasoning_done_emitted = False
self._reasoning_item_id: Optional[str] = None
self._accumulated_reasoning_content_parts: List[str] = []
self._accumulated_provider_specific_fields: Dict[str, Any] = {}
def _get_or_assign_tool_output_index(self, call_id: str) -> int:
existing = self._tool_output_index_by_call_id.get(call_id)
@ -479,16 +480,36 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
event.__dict__["sequence_number"] = self._sequence_number
return event
def create_litellm_model_response(
self,
) -> Optional[ModelResponse]:
return cast(
def _merge_provider_specific_fields(self, src: dict) -> None:
"""Merge provider_specific_fields using last-value-wins for lists.
List-valued keys (web_search_results, tool_results,
code_interpreter_results, etc.) are emitted cumulatively each
emission contains the full list so far. Using "last value wins"
matches stream_chunk_builder's semantics and avoids quadratic
growth from repeated extend calls.
"""
for key, val in src.items():
self._accumulated_provider_specific_fields[key] = val
def create_litellm_model_response(self) -> Optional[ModelResponse]:
response = cast(
Optional[ModelResponse],
stream_chunk_builder(
chunks=self.collected_chat_completion_chunks,
logging_obj=self.litellm_logging_obj,
),
)
if response is not None and self._accumulated_provider_specific_fields:
if (
not hasattr(response, "_hidden_params")
or response._hidden_params is None
):
response._hidden_params = {}
response._hidden_params.setdefault("provider_specific_fields", {}).update(
self._accumulated_provider_specific_fields
)
return response
@staticmethod
def _snapshot_chunk_for_stream_chunk_builder(
@ -853,6 +874,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if chunk is not None:
chunk = cast(ModelResponseStream, chunk)
self._ensure_output_item_for_chunk(chunk)
# Accumulate provider_specific_fields from chunk and delta
for src in (
getattr(chunk, "provider_specific_fields", None),
getattr(
chunk.choices[0].delta if chunk.choices else None,
"provider_specific_fields",
None,
),
):
if src and isinstance(src, dict):
self._merge_provider_specific_fields(src)
# Proceed to transformation
self.collected_chat_completion_chunks.append(
self._snapshot_chunk_for_stream_chunk_builder(chunk)
@ -964,6 +996,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
try:
chunk = self.litellm_custom_stream_wrapper.__next__()
self._ensure_output_item_for_chunk(chunk)
# Accumulate provider_specific_fields from chunk and delta
for src in (
getattr(chunk, "provider_specific_fields", None),
getattr(
chunk.choices[0].delta if chunk.choices else None,
"provider_specific_fields",
None,
),
):
if src and isinstance(src, dict):
self._merge_provider_specific_fields(src)
# Emit any just-queued output_item event
if self._pending_response_events:
return self._pending_response_events.pop(0)

View file

@ -42,6 +42,7 @@ from litellm.types.llms.openai import (
from litellm.types.responses.main import (
GenericResponseOutputItem,
GenericResponseOutputItemContentAnnotation,
OutputCodeInterpreterCall,
OutputFunctionToolCall,
OutputImageGenerationCall,
OutputText,
@ -1696,6 +1697,7 @@ class LiteLLMCompletionResponsesConfig:
) -> List[
Union[
GenericResponseOutputItem,
OutputCodeInterpreterCall,
OutputFunctionToolCall,
OutputImageGenerationCall,
ResponseFunctionToolCall,
@ -1704,6 +1706,7 @@ class LiteLLMCompletionResponsesConfig:
responses_output: List[
Union[
GenericResponseOutputItem,
OutputCodeInterpreterCall,
OutputFunctionToolCall,
OutputImageGenerationCall,
ResponseFunctionToolCall,
@ -1725,8 +1728,63 @@ class LiteLLMCompletionResponsesConfig:
chat_completion_response=chat_completion_response
)
)
# Convert server-side tool results (e.g. Anthropic code execution)
# into code_interpreter_call output items, replacing the corresponding
# function_call items so the output matches OpenAI's native shape.
tool_result_items = (
LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(
chat_completion_response
)
)
if tool_result_items:
result_by_id = {item.id: item for item in tool_result_items}
replaced_ids = set(result_by_id.keys())
responses_output = [
(
result_by_id[getattr(item, "call_id", None)]
if (
getattr(item, "type", None) == "function_call"
and getattr(item, "call_id", None) in replaced_ids
)
else item
)
for item in responses_output
]
return responses_output
@staticmethod
def _extract_tool_result_output_items(
chat_completion_response: ModelResponse,
) -> list:
"""Extract pre-built code_interpreter_call output items from provider_specific_fields.
Provider transformers (e.g. Anthropic) convert their native tool results
into OutputCodeInterpreterCall objects and store them in
provider_specific_fields["code_interpreter_results"]. This method
simply retrieves them no provider-specific parsing here.
"""
output_items: list = []
for choice in chat_completion_response.choices or []:
message = getattr(choice, "message", None)
if not message:
continue
psf = getattr(message, "provider_specific_fields", None)
if not psf or not isinstance(psf, dict):
continue
results = psf.get("code_interpreter_results")
if results and isinstance(results, list):
for item in results:
# In the streaming path, items are plain dicts after
# model_dump() in stream_chunk_builder. Reconstruct
# Pydantic objects so responses_output has a uniform type.
if isinstance(item, dict):
output_items.append(OutputCodeInterpreterCall(**item))
else:
output_items.append(item)
return output_items
@staticmethod
def _extract_reasoning_output_items(
chat_completion_response: ModelResponse,

View file

@ -166,11 +166,12 @@ class BaseResponsesAPIStreamingIterator:
)
setattr(item, "encrypted_content", wrapped_content)
# Store the completed response
if (
openai_responses_api_chunk
and getattr(openai_responses_api_chunk, "type", None)
== ResponsesAPIStreamEvents.RESPONSE_COMPLETED
# Store the completed response (also for incomplete/failed so logging still fires)
_chunk_type = getattr(openai_responses_api_chunk, "type", None)
if openai_responses_api_chunk and _chunk_type in (
ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
ResponsesAPIStreamEvents.RESPONSE_FAILED,
):
self.completed_response = openai_responses_api_chunk
# Add cost to usage object if include_cost_in_streaming_usage is True
@ -195,10 +196,12 @@ class BaseResponsesAPIStreamingIterator:
if cost is not None:
setattr(usage_obj, "cost", cost)
except Exception:
# If cost calculation fails, continue without cost
pass
self._handle_logging_completed_response()
if _chunk_type == ResponsesAPIStreamEvents.RESPONSE_FAILED:
self._handle_logging_failed_response()
else:
self._handle_logging_completed_response()
return openai_responses_api_chunk
@ -216,6 +219,32 @@ class BaseResponsesAPIStreamingIterator:
"""Base implementation - should be overridden by subclasses"""
pass
def _handle_logging_failed_response(self):
"""
Handle logging for RESPONSE_FAILED events by routing to failure handlers.
Unlike _handle_logging_completed_response (which calls success handlers),
this constructs an exception from the response error and routes to
async_failure_handler / failure_handler so logging integrations correctly
record the call as failed.
"""
response_obj = (
getattr(self.completed_response, "response", None)
if self.completed_response
else None
)
error_info = getattr(response_obj, "error", None) if response_obj else None
error_message = "Response failed"
if isinstance(error_info, dict):
error_message = error_info.get("message", str(error_info))
exception = litellm.APIError(
status_code=500,
message=error_message,
llm_provider=self.custom_llm_provider or "",
model=self.model or "",
)
self._handle_failure(exception)
async def _call_post_streaming_deployment_hook(self, chunk):
"""
Allow callbacks to modify streaming chunks before returning (parity with chat).

View file

@ -3874,14 +3874,23 @@ class Router:
The response from the handler function
"""
handler_name = original_function.__name__
metadata_variable_name = _get_router_metadata_variable_name(
function_name="generic_api_call"
)
try:
verbose_router_logger.debug(
f"Inside _generic_api_call() - handler: {handler_name}, model: {model}; kwargs: {kwargs}"
)
self._update_kwargs_before_fallbacks(
model=model,
kwargs=kwargs,
metadata_variable_name=metadata_variable_name,
)
deployment = self.get_available_deployment(
model=model,
messages=kwargs.get("messages", None),
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
self._update_kwargs_with_deployment(
deployment=deployment, kwargs=kwargs, function_name="generic_api_call"

View file

@ -84,6 +84,7 @@ from typing_extensions import Annotated, Dict, Required, TypedDict, override
from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject
from litellm.types.responses.main import (
GenericResponseOutputItem,
OutputCodeInterpreterCall,
OutputFunctionToolCall,
OutputImageGenerationCall,
)
@ -1242,6 +1243,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject):
List[
Union[
GenericResponseOutputItem,
OutputCodeInterpreterCall,
OutputFunctionToolCall,
OutputImageGenerationCall,
ResponseFunctionToolCall,
@ -1308,13 +1310,16 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject):
if not isinstance(serialized, list):
return serialized
return [
{
k: v
for k, v in item.items()
if v is not None or k not in ("status", "content", "encrypted_content")
}
if isinstance(item, dict) and item.get("type") == "reasoning"
else item
(
{
k: v
for k, v in item.items()
if v is not None
or k not in ("status", "content", "encrypted_content")
}
if isinstance(item, dict) and item.get("type") == "reasoning"
else item
)
for item in serialized
]

View file

@ -0,0 +1,14 @@
from pydantic import BaseModel, field_validator
class WorkerRegistryEntry(BaseModel):
worker_id: str
name: str
url: str
@field_validator("url")
@classmethod
def url_must_be_http(cls, v: str) -> str:
if not v.startswith(("http://", "https://")):
raise ValueError("Worker URL must start with http:// or https://")
return v

View file

@ -1,7 +1,9 @@
from typing import Optional
from typing import List, Optional
from pydantic import BaseModel
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
class UiDiscoveryEndpoints(BaseModel):
server_root_path: str
@ -9,3 +11,5 @@ class UiDiscoveryEndpoints(BaseModel):
auto_redirect_to_sso: bool
admin_ui_disabled: bool
sso_configured: bool
is_control_plane: bool = False
workers: List[WorkerRegistryEntry] = []

View file

@ -49,6 +49,42 @@ class OutputImageGenerationCall(BaseLiteLLMOpenAIResponseObject):
result: Optional[str] # Base64 encoded image data (without data:image prefix)
class OutputCodeInterpreterCallLog(BaseLiteLLMOpenAIResponseObject):
"""Log output from a code interpreter call"""
type: Literal["logs"]
logs: str
class OutputCodeInterpreterCall(BaseLiteLLMOpenAIResponseObject):
"""A code interpreter / code execution call output"""
type: Literal["code_interpreter_call"]
id: str
code: Optional[str]
container_id: Optional[str]
status: Literal["in_progress", "completed", "incomplete", "failed"]
outputs: Optional[List[OutputCodeInterpreterCallLog]]
def build_code_interpreter_log_outputs(
content: Any,
) -> Optional[List[OutputCodeInterpreterCallLog]]:
"""Convert Anthropic bash_code_execution stdout/stderr to log outputs.
Shared by streaming (handler.py) and non-streaming (transformation.py) paths.
"""
if not isinstance(content, dict):
return None
parts = []
if content.get("stdout"):
parts.append(content["stdout"])
if content.get("stderr"):
parts.append(f"STDERR: {content['stderr']}")
logs = "".join(parts)
return [OutputCodeInterpreterCallLog(type="logs", logs=logs)] if logs else None
class GenericResponseOutputItem(BaseLiteLLMOpenAIResponseObject):
"""
Generic response API output item

View file

@ -6160,7 +6160,10 @@ def validate_environment( # noqa: PLR0915
["AZURE_API_BASE", "AZURE_API_VERSION", "AZURE_API_KEY"]
)
elif custom_llm_provider == "anthropic":
if "ANTHROPIC_API_KEY" in os.environ:
if (
"ANTHROPIC_API_KEY" in os.environ
or "ANTHROPIC_AUTH_TOKEN" in os.environ
):
keys_in_environment = True
else:
missing_keys.append("ANTHROPIC_API_KEY")
@ -6399,7 +6402,10 @@ def validate_environment( # noqa: PLR0915
missing_keys.append("OPENAI_API_KEY")
## anthropic
elif model in litellm.anthropic_models:
if "ANTHROPIC_API_KEY" in os.environ:
if (
"ANTHROPIC_API_KEY" in os.environ
or "ANTHROPIC_AUTH_TOKEN" in os.environ
):
keys_in_environment = True
else:
missing_keys.append("ANTHROPIC_API_KEY")
@ -8593,9 +8599,7 @@ class ProviderConfigManager:
return ManusFilesConfig()
elif LlmProviders.ANTHROPIC == provider:
from litellm.llms.anthropic.files.transformation import (
AnthropicFilesConfig,
)
from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig
return AnthropicFilesConfig()
return None

View file

@ -4462,6 +4462,78 @@
"supports_vision": true,
"supports_web_search": true
},
"azure/gpt-5.4-mini": {
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_token": 7.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 4.5e-06,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true,
"supports_web_search": true,
"supports_none_reasoning_effort": false,
"supports_xhigh_reasoning_effort": false
},
"azure/gpt-5.4-nano": {
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.25e-06,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_native_streaming": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_service_tier": true,
"supports_vision": true,
"supports_web_search": true,
"supports_none_reasoning_effort": false,
"supports_xhigh_reasoning_effort": false
},
"azure/gpt-image-1": {
"cache_read_input_image_token_cost": 2.5e-06,
"cache_read_input_token_cost": 1.25e-06,
@ -37032,5 +37104,157 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true
},
"volcengine/doubao-seed-2-0-pro-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 4.6e-07,
"output_cost_per_token": 2.3e-06,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 7e-07,
"output_cost_per_token": 3.5e-06,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 1.4e-06,
"output_cost_per_token": 7e-06,
"range": [
128000.0,
256000.0
]
}
]
},
"volcengine/doubao-seed-2-0-lite-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 8.7e-08,
"output_cost_per_token": 5.2e-07,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 1.3e-07,
"output_cost_per_token": 7.8e-07,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 2.6e-07,
"output_cost_per_token": 1.6e-06,
"range": [
128000.0,
256000.0
]
}
]
},
"volcengine/doubao-seed-2-0-mini-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 2.9e-08,
"output_cost_per_token": 2.9e-07,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 5.8e-08,
"output_cost_per_token": 5.8e-07,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 1.2e-07,
"output_cost_per_token": 1.2e-06,
"range": [
128000.0,
256000.0
]
}
]
},
"volcengine/doubao-seed-2-0-code-preview-260215": {
"litellm_provider": "volcengine",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"source": "https://www.volcengine.com/docs/82379/1330310",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true,
"tiered_pricing": [
{
"input_cost_per_token": 4.6e-07,
"output_cost_per_token": 2.3e-06,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 7e-07,
"output_cost_per_token": 3.5e-06,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 1.4e-06,
"output_cost_per_token": 7e-06,
"range": [
128000.0,
256000.0
]
}
]
}
}

8
poetry.lock generated
View file

@ -3473,15 +3473,15 @@ files = [
[[package]]
name = "mcp"
version = "1.25.0"
version = "1.26.0"
description = "Model Context Protocol SDK"
optional = true
python-versions = ">=3.10"
groups = ["main"]
markers = "python_version >= \"3.10\" and extra == \"proxy\""
files = [
{file = "mcp-1.25.0-py3-none-any.whl", hash = "sha256:b37c38144a666add0862614cc79ec276e97d72aa8ca26d622818d4e278b9721a"},
{file = "mcp-1.25.0.tar.gz", hash = "sha256:56310361ebf0364e2d438e5b45f7668cbb124e158bb358333cd06e49e83a6802"},
{file = "mcp-1.26.0-py3-none-any.whl", hash = "sha256:904a21c33c25aa98ddbeb47273033c435e595bbacfdb177f4bd87f6dceebe1ca"},
{file = "mcp-1.26.0.tar.gz", hash = "sha256:db6e2ef491eecc1a0d93711a76f28dec2e05999f93afd48795da1c1137142c66"},
]
[package.dependencies]
@ -8018,4 +8018,4 @@ utils = ["numpydoc"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.9,<4.0"
content-hash = "eda34dfd8b35474beffee18893d6782c7b3d0d3d2c610f66237eb97176f43527"
content-hash = "2cf958f1a04fd5f1ab0e5cfc33bdbf441b518ed6c82d0f2546bf64cd3d2f89be"

View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm"
version = "1.82.4"
version = "1.82.5"
description = "Library to easily interface with LLM API providers"
authors = ["BerriAI"]
license = "MIT"
@ -184,7 +184,7 @@ requires = ["poetry-core", "wheel"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "1.82.4"
version = "1.82.5"
version_files = [
"pyproject.toml:^version"
]

View file

@ -30,9 +30,11 @@ from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterat
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
ResponseCompletedEvent,
ResponseFailedEvent,
ResponseIncompleteEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
OutputTextDeltaEvent
OutputTextDeltaEvent,
)
@ -429,3 +431,155 @@ class TestBaseResponsesAPIStreamingIterator:
mock_logging_obj.async_failure_handler.assert_not_called()
mock_logging_obj.failure_handler.assert_not_called()
def test_process_chunk_response_failed_calls_failure_handler(self):
"""
Test that a RESPONSE_FAILED event routes to failure handlers,
not success handlers. Failed responses represent genuine LLM-level
errors and should be logged as failures.
"""
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
mock_response = Mock()
mock_response.headers = {}
mock_response.aiter_lines = Mock()
mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
mock_logging_obj.model_call_details = {"litellm_params": {}}
mock_logging_obj.async_failure_handler = Mock()
mock_logging_obj.failure_handler = Mock()
mock_logging_obj.async_success_handler = Mock()
mock_logging_obj.success_handler = Mock()
mock_config = Mock(spec=BaseResponsesAPIConfig)
mock_responses_api_response = Mock(spec=ResponsesAPIResponse)
mock_responses_api_response.id = "resp_failed_123"
mock_responses_api_response.error = {
"type": "server_error",
"message": "The model encountered an error",
}
mock_responses_api_response.usage = None
mock_failed_event = Mock(spec=ResponseFailedEvent)
mock_failed_event.type = ResponsesAPIStreamEvents.RESPONSE_FAILED
mock_failed_event.response = mock_responses_api_response
mock_config.transform_streaming_response.return_value = mock_failed_event
iterator = ResponsesAPIStreamingIterator(
response=mock_response,
model="gpt-4",
responses_api_provider_config=mock_config,
logging_obj=mock_logging_obj,
litellm_metadata={"model_info": {"id": "model_123"}},
custom_llm_provider="openai",
)
test_chunk_data = {
"type": "response.failed",
"response": {
"id": "resp_failed_123",
"error": {
"type": "server_error",
"message": "The model encountered an error",
},
},
}
with patch.object(
ResponsesAPIRequestUtils,
"_update_responses_api_response_id_with_model_id",
return_value=mock_responses_api_response,
), patch(
"litellm.responses.streaming_iterator.run_async_function"
) as mock_run_async, patch(
"litellm.responses.streaming_iterator.executor"
) as mock_executor:
result = iterator._process_chunk(json.dumps(test_chunk_data))
assert result is not None
assert result.type == ResponsesAPIStreamEvents.RESPONSE_FAILED
assert iterator.completed_response == result
# Failure handler should have been called via _handle_failure
mock_run_async.assert_called_once()
call_kwargs = mock_run_async.call_args
assert (
call_kwargs[1]["async_function"]
== mock_logging_obj.async_failure_handler
)
mock_executor.submit.assert_called_once()
submit_args = mock_executor.submit.call_args
assert submit_args[0][0] == mock_logging_obj.failure_handler
def test_process_chunk_response_incomplete_calls_success_handler(self):
"""
Test that a RESPONSE_INCOMPLETE event routes to success handlers.
Incomplete responses (e.g. max_output_tokens reached) are still valid
responses with usage data analogous to finish_reason='length' in chat.
"""
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
mock_response = Mock()
mock_response.headers = {}
mock_response.aiter_lines = Mock()
mock_logging_obj = Mock(spec=LiteLLMLoggingObj)
mock_logging_obj.model_call_details = {"litellm_params": {}}
mock_logging_obj.async_failure_handler = Mock()
mock_logging_obj.failure_handler = Mock()
mock_logging_obj.async_success_handler = Mock()
mock_logging_obj.success_handler = Mock()
mock_config = Mock(spec=BaseResponsesAPIConfig)
mock_responses_api_response = Mock(spec=ResponsesAPIResponse)
mock_responses_api_response.id = "resp_incomplete_123"
mock_responses_api_response.incomplete_details = {
"reason": "max_output_tokens"
}
mock_responses_api_response.usage = None
mock_incomplete_event = Mock(spec=ResponseIncompleteEvent)
mock_incomplete_event.type = ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE
mock_incomplete_event.response = mock_responses_api_response
mock_config.transform_streaming_response.return_value = mock_incomplete_event
iterator = ResponsesAPIStreamingIterator(
response=mock_response,
model="gpt-4",
responses_api_provider_config=mock_config,
logging_obj=mock_logging_obj,
litellm_metadata={"model_info": {"id": "model_123"}},
custom_llm_provider="openai",
)
test_chunk_data = {
"type": "response.incomplete",
"response": {
"id": "resp_incomplete_123",
"incomplete_details": {"reason": "max_output_tokens"},
},
}
with patch.object(
ResponsesAPIRequestUtils,
"_update_responses_api_response_id_with_model_id",
return_value=mock_responses_api_response,
), patch(
"asyncio.create_task"
) as mock_create_task, patch(
"litellm.responses.streaming_iterator.executor"
) as mock_executor:
result = iterator._process_chunk(json.dumps(test_chunk_data))
assert result is not None
assert result.type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE
assert iterator.completed_response == result
# Success handler should have been called (via _handle_logging_completed_response)
mock_create_task.assert_called_once()
mock_executor.submit.assert_called_once()
# Failure handlers should NOT have been called
mock_logging_obj.async_failure_handler.assert_not_called()
mock_logging_obj.failure_handler.assert_not_called()

View file

@ -593,7 +593,7 @@ def test_datadog_static_methods():
# Test tags format with default values
assert (
"env:unknown,service:litellm-server,version:unknown,HOSTNAME:"
in get_datadog_tags()
in ",".join(get_datadog_tags())
)
# Test with custom environment variables
@ -631,7 +631,7 @@ def test_datadog_static_methods():
# Test tags format with custom values
expected_custom_tags = "env:production,service:custom-service,version:1.0.0,HOSTNAME:test-host,POD_NAME:pod-123"
print("DataDogLogger._get_datadog_tags()", get_datadog_tags())
assert get_datadog_tags() == expected_custom_tags
assert ",".join(get_datadog_tags()) == expected_custom_tags
@pytest.mark.asyncio
@ -672,11 +672,11 @@ def test_get_datadog_tags():
"""Test the _get_datadog_tags static method with various inputs"""
# Test with no standard_logging_object and default env vars
base_tags = get_datadog_tags()
assert "env:" in base_tags
assert "service:" in base_tags
assert "version:" in base_tags
assert "POD_NAME:" in base_tags
assert "HOSTNAME:" in base_tags
assert any("env:" in t for t in base_tags)
assert any("service:" in t for t in base_tags)
assert any("version:" in t for t in base_tags)
assert any("POD_NAME:" in t for t in base_tags)
assert any("HOSTNAME:" in t for t in base_tags)
# Test with custom env vars
test_env = {
@ -705,12 +705,12 @@ def test_get_datadog_tags():
# Test with empty request_tags
standard_logging_obj["request_tags"] = []
tags_empty_request = get_datadog_tags(standard_logging_obj)
assert "request_tag:" not in tags_empty_request
assert not any(t.startswith("request_tag:") for t in tags_empty_request)
# Test with None request_tags
standard_logging_obj["request_tags"] = None
tags_none_request = get_datadog_tags(standard_logging_obj)
assert "request_tag:" not in tags_none_request
assert not any(t.startswith("request_tag:") for t in tags_none_request)
@pytest.mark.asyncio

View file

@ -2278,6 +2278,75 @@ async def test_post_call_failure_hook_auth_error_llm_api_route():
mock_handle_logging.assert_called_once()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"request_data, route, expected_call_type",
[
(
{"model": "bad-model", "messages": [{"role": "user", "content": "hello"}]},
"/v1/chat/completions",
"acompletion",
),
(
{"model": "bad-model", "prompt": "hello"},
"/v1/completions",
"atext_completion",
),
(
{"model": "bad-model", "input": ["hello"]},
"/v1/embeddings",
"aembedding",
),
],
)
async def test_handle_logging_proxy_only_error_syncs_normalized_call_type(
request_data, route, expected_call_type
):
from fastapi import HTTPException
from litellm.caching.caching import DualCache
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.utils import ProxyLogging
cache = DualCache()
proxy_logging = ProxyLogging(user_api_key_cache=cache)
captured_logging_obj = {}
original_function_setup = litellm.utils.function_setup
def _capture_function_setup(*args, **kwargs):
logging_obj, data = original_function_setup(*args, **kwargs)
captured_logging_obj["logging_obj"] = logging_obj
return logging_obj, data
with patch(
"litellm.proxy.utils.litellm.utils.function_setup",
side_effect=_capture_function_setup,
), patch.object(
Logging, "async_failure_handler", new=AsyncMock(return_value=None)
), patch.object(
Logging, "failure_handler", return_value=None
), patch(
"litellm.proxy.utils.threading.Thread"
) as mock_thread:
mock_thread.return_value.start = Mock()
await proxy_logging._handle_logging_proxy_only_error(
request_data=request_data,
user_api_key_dict=UserAPIKeyAuth(
api_key="test_key",
user_id="test_user",
token="test_token",
request_route=route,
),
route=route,
original_exception=HTTPException(status_code=400, detail="bad request"),
)
logging_obj = captured_logging_obj["logging_obj"]
assert logging_obj.call_type == expected_call_type
assert logging_obj.model_call_details["call_type"] == expected_call_type
@pytest.mark.asyncio
async def test_during_call_hook_parallel_execution():
"""

View file

@ -1568,6 +1568,91 @@ def test_handle_clientside_credential_with_deployment_model_name(model_list):
print("✓ _handle_clientside_credential test passed!")
def test_sync_generic_api_call_preserves_requested_model_group_in_logs():
router = Router(
model_list=[
{
"model_name": "claude-sonnet-4-6",
"litellm_params": {
"model": "bedrock/global.anthropic.claude-sonnet-4-6",
"aws_access_key_id": "test-access-key",
"aws_secret_access_key": "test-secret-key",
"aws_region_name": "us-west-2",
},
}
]
)
try:
captured_kwargs = {}
def mock_original_function(**kwargs):
captured_kwargs.update(kwargs)
return {"status": "ok"}
response = router._generic_api_call_with_fallbacks(
model="claude-sonnet-4-6",
original_function=mock_original_function,
)
assert response == {"status": "ok"}
assert (
captured_kwargs["model"] == "bedrock/global.anthropic.claude-sonnet-4-6"
)
assert (
captured_kwargs["litellm_metadata"]["model_group"] == "claude-sonnet-4-6"
)
assert (
captured_kwargs["litellm_metadata"]["deployment"]
== "bedrock/global.anthropic.claude-sonnet-4-6"
)
finally:
router.discard()
def test_sync_generic_api_call_uses_request_kwargs_for_deployment_selection():
router = Router(
model_list=[
{
"model_name": "regional-model",
"litellm_params": {
"model": "anthropic/us-model",
"api_key": "test-api-key",
"region_name": "us",
},
},
{
"model_name": "regional-model",
"litellm_params": {
"model": "anthropic/eu-model",
"api_key": "test-api-key",
"region_name": "eu",
},
},
],
enable_pre_call_checks=True,
)
try:
captured_kwargs = {}
def mock_original_function(**kwargs):
captured_kwargs.update(kwargs)
return {"status": "ok"}
response = router._generic_api_call_with_fallbacks(
model="regional-model",
original_function=mock_original_function,
messages=[{"role": "user", "content": "Hello from Europe"}],
allowed_model_region="eu",
)
assert response == {"status": "ok"}
assert captured_kwargs["model"] == "anthropic/eu-model"
finally:
router.discard()
@pytest.mark.parametrize(
"function_name, expected_metadata_key",
[

View file

@ -44,7 +44,7 @@ class TestDatadogTagsRegression:
assert "env:test-env" in tags_legacy
assert "service:test-service" in tags_legacy
# Verify NO team tag (should not invent one)
assert "team:" not in tags_legacy
assert not any(t.startswith("team:") for t in tags_legacy)
# Case 2: New feature (team info provided)
payload_with_team = StandardLoggingPayload(

View file

@ -132,3 +132,57 @@ class TestLangsmithLoggerInit:
assert (
logger.sampling_rate >= 0.0
), f"sampling_rate should be non-negative, got {logger.sampling_rate}"
class TestLangsmithPrepareLogData:
"""Regression test for #24001: _prepare_log_data must inject
usage_metadata into outputs so LangSmith's Cost column is populated."""
@patch("asyncio.create_task")
@patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
def test_outputs_contain_usage_metadata(self, mock_create_task):
logger = LangsmithLogger(
langsmith_api_key="test-key",
langsmith_project="test-project",
)
payload = {
"id": "test-id",
"response": {"choices": [{"message": {"content": "hi"}}]},
"metadata": {},
"startTime": 1.0,
"endTime": 2.0,
"request_tags": [],
"error_str": None,
"status": "success",
"response_cost": 0.0042,
"prompt_tokens": 100,
"completion_tokens": 50,
"total_tokens": 150,
}
kwargs = {
"litellm_params": {"metadata": {}},
"standard_logging_object": payload,
}
credentials = {
"LANGSMITH_API_KEY": "test-key",
"LANGSMITH_PROJECT": "test-project",
"LANGSMITH_BASE_URL": "https://api.smith.langchain.com",
}
data = logger._prepare_log_data(
kwargs=kwargs,
response_obj=None,
start_time=1.0,
end_time=2.0,
credentials=credentials,
)
assert "usage_metadata" in data["outputs"]
um = data["outputs"]["usage_metadata"]
assert um["total_cost"] == 0.0042
assert um["input_tokens"] == 100
assert um["output_tokens"] == 50
assert um["total_tokens"] == 150

View file

@ -11,7 +11,8 @@ sys.path.insert(
import time
from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
from litellm.litellm_core_utils.litellm_logging import \
Logging as LitellmLogging
from litellm.litellm_core_utils.litellm_logging import set_callbacks
from litellm.types.utils import ModelResponse, TextCompletionResponse
@ -139,7 +140,8 @@ def test_sentry_environment():
def test_use_custom_pricing_for_model():
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
from litellm.litellm_core_utils.litellm_logging import \
use_custom_pricing_for_model
litellm_params = {
"custom_llm_provider": "azure",
@ -154,7 +156,8 @@ def test_use_custom_pricing_for_model_via_litellm_metadata():
Generic API call routes (/messages, /responses) store model_info
under litellm_metadata, not metadata. Regression test for #23185.
"""
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
from litellm.litellm_core_utils.litellm_logging import \
use_custom_pricing_for_model
litellm_params = {
"litellm_metadata": {
@ -170,7 +173,8 @@ def test_use_custom_pricing_for_model_via_litellm_metadata():
def test_use_custom_pricing_not_detected_litellm_metadata_no_pricing():
"""Should return False when litellm_metadata.model_info has no pricing keys."""
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
from litellm.litellm_core_utils.litellm_logging import \
use_custom_pricing_for_model
litellm_params = {
"litellm_metadata": {
@ -180,6 +184,198 @@ def test_use_custom_pricing_not_detected_litellm_metadata_no_pricing():
assert use_custom_pricing_for_model(litellm_params) is False
def test_response_cost_calculator_uses_router_model_id_from_litellm_metadata():
"""_response_cost_calculator should extract router_model_id from
litellm_params.litellm_metadata.model_info.id when the result object
does not carry _hidden_params (e.g. ResponsesAPIResponse from /v1/responses
streaming). Regression test for custom pricing on streaming responses."""
import litellm
from litellm.litellm_core_utils.litellm_logging import \
Logging as LiteLLMLoggingObj
from litellm.types.llms.openai import ResponsesAPIResponse
custom_model_id = "gpt-5-custom-pricing"
custom_input_cost = 125.0
custom_output_cost = 10.0
litellm.register_model(
model_cost={
custom_model_id: {
"input_cost_per_token": custom_input_cost,
"output_cost_per_token": custom_output_cost,
"max_tokens": 128000,
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"litellm_provider": "openai",
}
}
)
try:
logging_obj = LiteLLMLoggingObj(
model="gpt-5",
messages=[{"role": "user", "content": "Hi"}],
stream=True,
call_type="aresponses",
start_time=time.time(),
litellm_call_id="test-123",
function_id="test-fn",
)
logging_obj.update_environment_variables(
model="gpt-5",
user="",
optional_params={},
litellm_params={
"api_base": "",
"litellm_metadata": {
"model_info": {
"id": custom_model_id,
"input_cost_per_token": custom_input_cost,
"output_cost_per_token": custom_output_cost,
},
},
},
)
response_obj = ResponsesAPIResponse(
id="resp_abc",
created_at=1234567890,
model="gpt-5",
output=[],
usage={
"input_tokens": 10,
"output_tokens": 5,
"total_tokens": 15,
},
)
cost = logging_obj._response_cost_calculator(result=response_obj)
assert cost is not None, "Cost should not be None"
expected_cost = (10 * custom_input_cost) + (5 * custom_output_cost)
assert cost == pytest.approx(
expected_cost
), f"Expected {expected_cost}, got {cost}"
finally:
litellm.model_cost.pop(custom_model_id, None)
class TestGetRouterModelId:
"""Tests for the get_router_model_id helper method."""
def test_returns_id_from_litellm_metadata(self, logging_obj):
"""Should extract model_info.id from litellm_metadata."""
logging_obj.litellm_params = {
"litellm_metadata": {
"model_info": {"id": "custom-deploy-1"},
},
}
assert logging_obj.get_router_model_id() == "custom-deploy-1"
def test_returns_id_from_metadata(self, logging_obj):
"""Should fall back to metadata when litellm_metadata has no model_info."""
logging_obj.litellm_params = {
"metadata": {
"model_info": {"id": "custom-deploy-2"},
},
}
assert logging_obj.get_router_model_id() == "custom-deploy-2"
def test_prefers_litellm_metadata_over_metadata(self, logging_obj):
"""litellm_metadata should take priority over metadata."""
logging_obj.litellm_params = {
"litellm_metadata": {
"model_info": {"id": "from-litellm-meta"},
},
"metadata": {
"model_info": {"id": "from-meta"},
},
}
assert logging_obj.get_router_model_id() == "from-litellm-meta"
def test_returns_none_when_no_model_info(self, logging_obj):
"""Should return None when no model_info is present."""
logging_obj.litellm_params = {"api_base": ""}
assert logging_obj.get_router_model_id() is None
def test_returns_none_when_no_litellm_params(self):
"""Should return None when litellm_params is not set."""
from litellm.litellm_core_utils.litellm_logging import \
Logging as LiteLLMLoggingObj
obj = LiteLLMLoggingObj(
model="test",
messages=[],
stream=False,
call_type="completion",
start_time=time.time(),
litellm_call_id="x",
function_id="x",
)
# litellm_params exists but is empty by default
assert obj.get_router_model_id() is None
class TestAnthropicPassthroughCustomPricing:
"""Verify the Anthropic pass-through handler forwards custom pricing."""
def test_completion_cost_receives_custom_pricing_args(self):
"""_create_anthropic_response_logging_payload should pass
custom_pricing and router_model_id to litellm.completion_cost
when the logging object carries custom pricing in model_info."""
from unittest.mock import patch
from litellm.litellm_core_utils.litellm_logging import \
Logging as LiteLLMLoggingObj
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import \
AnthropicPassthroughLoggingHandler
logging_obj = LiteLLMLoggingObj(
model="claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "Hi"}],
stream=False,
call_type="anthropic_messages",
start_time=time.time(),
litellm_call_id="test-456",
function_id="test-fn",
)
logging_obj.update_environment_variables(
model="claude-sonnet-4-20250514",
user="",
optional_params={},
litellm_params={
"api_base": "",
"litellm_metadata": {
"model_info": {
"id": "claude-custom-pricing",
"input_cost_per_token": 0.5,
"output_cost_per_token": 1.5,
},
},
},
)
logging_obj.model_call_details["custom_llm_provider"] = "anthropic"
mock_response = ModelResponse()
mock_response.usage = {"prompt_tokens": 10, "completion_tokens": 5} # type: ignore
with patch("litellm.completion_cost", return_value=42.0) as mock_cost:
AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
litellm_model_response=mock_response,
model="claude-sonnet-4-20250514",
kwargs={},
start_time=time.time(),
end_time=time.time(),
logging_obj=logging_obj,
)
mock_cost.assert_called_once()
call_kwargs = mock_cost.call_args
assert call_kwargs.kwargs.get("custom_pricing") is True
assert call_kwargs.kwargs.get("router_model_id") == "claude-custom-pricing"
class TestUpdateFromKwargs:
"""Tests for the update_from_kwargs convenience wrapper."""
@ -245,9 +441,8 @@ class TestUpdateFromKwargs:
def test_custom_pricing_detected_via_litellm_metadata(self, logging_obj):
"""Custom pricing in litellm_metadata.model_info should set custom_pricing flag."""
from litellm.litellm_core_utils.litellm_logging import (
use_custom_pricing_for_model,
)
from litellm.litellm_core_utils.litellm_logging import \
use_custom_pricing_for_model
lm_meta = {
"model_info": {
@ -306,7 +501,8 @@ async def test_datadog_logger_not_shadowed_by_llm_obs(monkeypatch):
monkeypatch.setenv("DD_SITE", "us5.datadoghq.com")
from litellm.integrations.datadog.datadog import DataDogLogger
from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
from litellm.integrations.datadog.datadog_llm_obs import \
DataDogLLMObsLogger
from litellm.litellm_core_utils import litellm_logging as logging_module
logging_module._in_memory_loggers.clear()
@ -347,7 +543,8 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch):
) # no trailing slash on purpose
# Import after env vars are set (important if module-level caching exists)
from litellm.integrations.opentelemetry import OpenTelemetry # logger class
from litellm.integrations.opentelemetry import \
OpenTelemetry # logger class
from litellm.litellm_core_utils import litellm_logging as logging_module
logging_module._in_memory_loggers.clear()
@ -676,7 +873,8 @@ def test_success_handler_runs_guardrail_logging_hook_when_enabled(logging_obj):
def test_get_user_agent_tags():
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
tags = StandardLoggingPayloadSetup._get_user_agent_tags(
proxy_server_request={
@ -691,7 +889,8 @@ def test_get_user_agent_tags():
def test_get_request_tags():
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
tags = StandardLoggingPayloadSetup._get_request_tags(
litellm_params={"metadata": {"tags": ["test-tag"]}},
@ -718,7 +917,8 @@ def test_get_request_tags_from_metadata_and_litellm_metadata():
4. No tags in either
5. None values for metadata/litellm_metadata
"""
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
# Test case 1: Tags in metadata only
tags = StandardLoggingPayloadSetup._get_request_tags(
@ -799,7 +999,8 @@ def test_get_request_tags_does_not_mutate_original_tags():
would cause User-Agent tags to be duplicated because the function was mutating
the original tags list instead of creating a copy.
"""
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
# Create metadata with original tags
original_tags = ["custom-tag-1", "custom-tag-2"]
@ -859,7 +1060,8 @@ def test_get_request_tags_does_not_mutate_original_tags():
def test_get_extra_header_tags():
"""Test the _get_extra_header_tags method with various scenarios."""
import litellm
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
# Store original value to restore later
original_extra_headers = getattr(litellm, "extra_spend_tag_headers", None)
@ -1080,7 +1282,8 @@ async def test_e2e_generate_cold_storage_object_key_successful():
from datetime import datetime, timezone
from unittest.mock import patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
# Create test data
start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
@ -1122,7 +1325,8 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
# Create test data
start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
@ -1173,7 +1377,8 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
# Create test data
start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
@ -1220,7 +1425,8 @@ async def test_e2e_generate_cold_storage_object_key_not_configured():
from unittest.mock import patch
import litellm
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
# Create test data
start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
@ -1244,7 +1450,8 @@ def test_get_final_response_obj_with_empty_response_obj_and_list_init():
When response_obj is empty (falsy), the method should return init_response_obj if it's a list.
"""
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
# Create test objects
class TestObject1:
@ -1280,7 +1487,8 @@ def test_get_usage_as_dict():
"""
Test get_usage_as_dict returns usage as plain dict from response_obj or combined_usage_object.
"""
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
from litellm.types.utils import Usage
# Test case 1: None response_obj returns empty usage dict
@ -1318,7 +1526,8 @@ def test_append_system_prompt_messages():
"""
Test append_system_prompt_messages prepends system message from kwargs to messages list.
"""
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
# Test case 1: system in kwargs with existing messages
kwargs = {"system": "You are a helpful assistant"}
@ -1389,7 +1598,8 @@ async def test_async_success_handler_sets_standard_logging_object_for_pass_throu
from datetime import datetime
from unittest.mock import patch
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import \
Logging as LiteLLMLoggingObj
from litellm.types.utils import StandardPassThroughResponseObject
# Create a logging object for a pass-through endpoint
@ -1470,7 +1680,8 @@ async def test_async_success_handler_prevents_reprocessing_for_pass_through_endp
from datetime import datetime
from unittest.mock import patch
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import \
Logging as LiteLLMLoggingObj
from litellm.types.utils import StandardPassThroughResponseObject
# Create a logging object for a pass-through endpoint
@ -1546,7 +1757,8 @@ async def test_async_success_handler_sets_standard_logging_object_for_streaming_
from datetime import datetime
from unittest.mock import patch
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import \
Logging as LiteLLMLoggingObj
from litellm.types.utils import StandardPassThroughResponseObject
# Create a logging object for a streaming pass-through endpoint
@ -1602,7 +1814,8 @@ def test_get_error_information_error_code_priority():
Test get_error_information prioritizes 'code' attribute over 'status_code' attribute
and handles edge cases like empty strings and "None" string values.
"""
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.litellm_logging import \
StandardLoggingPayloadSetup
# Test case 1: Exception with 'code' attribute (ProxyException style)
class ProxyException(Exception):
@ -1795,7 +2008,8 @@ async def test_async_success_handler_preserves_response_cost_for_pass_through_en
by pass-through handlers (Gemini/Vertex)."""
from datetime import datetime
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import \
Logging as LiteLLMLoggingObj
from litellm.types.utils import ModelResponse, Usage
logging_obj = LiteLLMLoggingObj(

View file

@ -1340,6 +1340,121 @@ def test_is_chunk_non_empty_with_valid_tool_calls(
)
def _make_chunk(content: Optional[str]) -> ModelResponseStream:
return ModelResponseStream(
id="test",
created=1741037890,
model="test-model",
choices=[StreamingChoices(index=0, delta=Delta(content=content))],
)
def _build_chunks(pattern: list[str], N: int) -> list[ModelResponseStream]:
"""
Build a list of chunks based on a pattern specification.
"""
chunks = []
for i, p in enumerate(pattern):
if p == "same":
chunks.append(_make_chunk("same_chunk"))
elif p == "diff":
chunks.append(_make_chunk(f"chunk_{i}"))
else:
chunks.append(_make_chunk(p))
return chunks
_REPETITION_TEST_CASES = [
# Basic cases
pytest.param(
["same"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
True,
id="all_identical_raises",
),
pytest.param(
["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 1),
False,
id="below_threshold_no_raise",
),
pytest.param(
[None] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="none_content_no_raise",
),
pytest.param(
[""] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="empty_content_no_raise",
),
# Short content (len <= 2) should not raise
pytest.param(
["##"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="short_content_2chars_no_raise",
),
pytest.param(
["{"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="short_content_1char_no_raise",
),
pytest.param(
["ab"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="short_content_2chars_ab_no_raise",
),
# All different chunks
pytest.param(
["diff"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT,
False,
id="all_different_no_raise",
),
# One chunk different at various positions
pytest.param(
["different_first"] + ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 1),
False,
id="first_chunk_different_no_raise",
),
pytest.param(
["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 1) + ["different_last"],
False,
id="last_chunk_different_no_raise",
),
pytest.param(
["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 + 1) + ["different_mid"] + ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 + 1),
False,
id="middle_chunk_different_no_raise",
),
pytest.param(
["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - 2) + ["diff", "diff"],
False,
id="last_two_different_no_raise",
),
pytest.param(
["diff"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT + ["same"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT + ["diff"],
True,
id="in_between_same_and_diff_raise",
),
]
@pytest.mark.parametrize("chunks_pattern,should_raise", _REPETITION_TEST_CASES)
def test_raise_on_model_repetition(
initialized_custom_stream_wrapper: CustomStreamWrapper,
chunks_pattern: list,
should_raise: bool,
):
wrapper = initialized_custom_stream_wrapper
chunks = _build_chunks(chunks_pattern, len(chunks_pattern))
if should_raise:
with pytest.raises(litellm.InternalServerError) as exc_info:
for chunk in chunks:
wrapper.chunks.append(chunk)
wrapper.raise_on_model_repetition()
assert "repeating the same chunk" in str(exc_info.value)
else:
for chunk in chunks:
wrapper.chunks.append(chunk)
wrapper.raise_on_model_repetition()
def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj):
"""
Test that provider-reported usage from a post-finish_reason chunk

View file

@ -6,6 +6,7 @@ from litellm.types.llms.openai import (
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
)
from litellm.types.responses.main import OutputCodeInterpreterCall
def test_redacted_thinking_content_block_delta():
@ -479,14 +480,22 @@ def test_partial_json_chunk_accumulation():
# First partial chunk should return None (still accumulating)
result1 = iterator._parse_sse_data(f"data:{partial_chunk_1}")
assert result1 is None, "First partial chunk should return None while accumulating"
assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode"
assert iterator.accumulated_json == partial_chunk_1, "Should have accumulated first part"
assert (
iterator.chunk_type == "accumulated_json"
), "Should switch to accumulated_json mode"
assert (
iterator.accumulated_json == partial_chunk_1
), "Should have accumulated first part"
# Second partial chunk should complete the JSON and return a parsed result
result2 = iterator._parse_sse_data(f"data:{partial_chunk_2}")
assert result2 is not None, "Second chunk should return parsed result"
assert iterator.accumulated_json == "", "Buffer should be cleared after successful parse"
assert result2.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result2.choices[0].delta.content}'"
assert (
iterator.accumulated_json == ""
), "Buffer should be cleared after successful parse"
assert (
result2.choices[0].delta.content == "Hello"
), f"Expected 'Hello', got '{result2.choices[0].delta.content}'"
def test_complete_json_chunk_no_accumulation():
@ -503,7 +512,9 @@ def test_complete_json_chunk_no_accumulation():
assert result is not None, "Complete chunk should return parsed result immediately"
assert iterator.chunk_type == "valid_json", "Should remain in valid_json mode"
assert iterator.accumulated_json == "", "Buffer should remain empty"
assert result.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result.choices[0].delta.content}'"
assert (
result.choices[0].delta.content == "Hello"
), f"Expected 'Hello', got '{result.choices[0].delta.content}'"
def test_multiple_partial_chunks_accumulation():
@ -620,7 +631,9 @@ def test_web_search_tool_result_no_extra_tool_calls():
# Should have exactly 2 tool calls:
# 1. From content_block_start (server_tool_use) with id and name
# 2. From content_block_delta with the actual query
assert len(tool_calls_emitted) == 2, f"Expected 2 tool calls, got {len(tool_calls_emitted)}"
assert (
len(tool_calls_emitted) == 2
), f"Expected 2 tool calls, got {len(tool_calls_emitted)}"
# First tool call should have the id and name
assert tool_calls_emitted[0]["id"] == "srvtoolu_01ABC123"
@ -722,7 +735,10 @@ def test_web_search_tool_result_captured_in_provider_specific_fields():
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '{"query": "otter facts"}'},
"delta": {
"type": "input_json_delta",
"partial_json": '{"query": "otter facts"}',
},
},
# 4. content_block_stop for server_tool_use
{"type": "content_block_stop", "index": 0},
@ -822,7 +838,10 @@ def test_web_fetch_tool_result_captured_in_provider_specific_fields():
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '{"url": "https://example.com"}'},
"delta": {
"type": "input_json_delta",
"partial_json": '{"url": "https://example.com"}',
},
},
# 4. content_block_stop for server_tool_use
{"type": "content_block_stop", "index": 0},
@ -946,7 +965,7 @@ def test_web_fetch_tool_result_no_extra_tool_calls():
def test_container_in_provider_specific_fields_streaming():
"""
Test that container is captured in provider_specific_fields for streaming responses.
When container with skills is used, the container field should be present in
the provider_specific_fields of the message_delta chunk.
"""
@ -1025,7 +1044,9 @@ def test_container_in_provider_specific_fields_streaming():
]
# Verify container was captured
assert container_field is not None, "container should be captured in provider_specific_fields"
assert (
container_field is not None
), "container should be captured in provider_specific_fields"
assert (
container_field["id"] == "container_011CW9hA9zpZ8xD3bjjShy4p"
), "container id should match"
@ -1033,18 +1054,14 @@ def test_container_in_provider_specific_fields_streaming():
container_field["expires_at"] == "2025-12-16T04:57:16.913181Z"
), "expires_at should match"
assert len(container_field["skills"]) == 1, "Should have 1 skill"
assert (
container_field["skills"][0]["skill_id"] == "pptx"
), "skill_id should be pptx"
assert (
container_field["skills"][0]["version"] == "20251013"
), "version should match"
assert container_field["skills"][0]["skill_id"] == "pptx", "skill_id should be pptx"
assert container_field["skills"][0]["version"] == "20251013", "version should match"
def test_container_in_provider_specific_fields_non_streaming():
"""
Test that container is captured in provider_specific_fields for non-streaming responses.
When container with skills is used in non-streaming, the container field should be
present in the provider_specific_fields of the response.
"""
@ -1106,7 +1123,7 @@ def test_container_in_provider_specific_fields_non_streaming():
def test_container_absent_when_not_provided():
"""
Test that container is not added to provider_specific_fields when not provided.
This ensures we don't add empty or None container fields.
"""
iterator = ModelResponseIterator(
@ -1133,3 +1150,434 @@ def test_container_absent_when_not_provided():
assert (
"container" not in model_response.choices[0].delta.provider_specific_fields
), "container should not be present when not provided in delta"
def test_streaming_code_execution_produces_code_interpreter_results():
"""
Test that bash_code_execution_tool_result content blocks in streaming
produce code_interpreter_results in provider_specific_fields, so the
Responses API layer can use them without Anthropic-specific knowledge.
"""
chunks = [
{
"type": "message_start",
"message": {
"id": "msg_01XYZ",
"type": "message",
"role": "assistant",
"content": [],
"usage": {"input_tokens": 100, "output_tokens": 1},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "text",
"text": "",
},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": "Running code..."},
},
{"type": "content_block_stop", "index": 0},
{
"type": "content_block_start",
"index": 1,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_01ABC",
"name": "bash_code_execution",
"input": {"command": "echo hello"},
},
},
{"type": "content_block_stop", "index": 1},
{
"type": "content_block_start",
"index": 2,
"content_block": {
"type": "bash_code_execution_tool_result",
"tool_use_id": "srvtoolu_01ABC",
"content": {
"type": "bash_code_execution_result",
"stdout": "hello\n",
"stderr": "",
"return_code": 0,
},
},
},
{"type": "content_block_stop", "index": 2},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 50},
},
]
iterator = ModelResponseIterator(None, sync_stream=True)
found_code_interpreter_results = False
for chunk in chunks:
parsed = iterator.chunk_parser(chunk)
psf = None
if parsed.choices and parsed.choices[0].delta:
psf = getattr(parsed.choices[0].delta, "provider_specific_fields", None)
if psf and "code_interpreter_results" in psf:
found_code_interpreter_results = True
results = psf["code_interpreter_results"]
assert len(results) == 1
assert isinstance(results[0], OutputCodeInterpreterCall)
assert results[0].type == "code_interpreter_call"
assert results[0].id == "srvtoolu_01ABC"
assert results[0].code == "echo hello"
assert results[0].outputs is not None
assert len(results[0].outputs) == 1
assert results[0].outputs[0].logs == "hello\n"
assert found_code_interpreter_results, (
"code_interpreter_results should appear in provider_specific_fields "
"when bash_code_execution_tool_result is streamed"
)
def test_streaming_multiple_code_executions_no_duplicates():
"""
Test that multiple code executions in a single streaming response emit
cumulative code_interpreter_results on each chunk (matching stream_chunk_builder's
"last value wins" contract). The final emission must contain ALL results.
"""
chunks = [
{
"type": "message_start",
"message": {
"id": "msg_01XYZ",
"type": "message",
"role": "assistant",
"content": [],
"usage": {"input_tokens": 100, "output_tokens": 1},
},
},
# First code execution
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_01AAA",
"name": "bash_code_execution",
"input": {"command": "echo first"},
},
},
{"type": "content_block_stop", "index": 0},
{
"type": "content_block_start",
"index": 1,
"content_block": {
"type": "bash_code_execution_tool_result",
"tool_use_id": "srvtoolu_01AAA",
"content": {
"type": "bash_code_execution_result",
"stdout": "first\n",
"stderr": "",
"return_code": 0,
},
},
},
{"type": "content_block_stop", "index": 1},
# Second code execution
{
"type": "content_block_start",
"index": 2,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_01BBB",
"name": "bash_code_execution",
"input": {"command": "echo second"},
},
},
{"type": "content_block_stop", "index": 2},
{
"type": "content_block_start",
"index": 3,
"content_block": {
"type": "bash_code_execution_tool_result",
"tool_use_id": "srvtoolu_01BBB",
"content": {
"type": "bash_code_execution_result",
"stdout": "second\n",
"stderr": "",
"return_code": 0,
},
},
},
{"type": "content_block_stop", "index": 3},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 50},
},
]
iterator = ModelResponseIterator(None, sync_stream=True)
# Collect each emission of code_interpreter_results
emissions = []
for chunk in chunks:
parsed = iterator.chunk_parser(chunk)
psf = None
if parsed.choices and parsed.choices[0].delta:
psf = getattr(parsed.choices[0].delta, "provider_specific_fields", None)
if psf and "code_interpreter_results" in psf:
emissions.append(psf["code_interpreter_results"])
# Should have 2 emissions (one per tool_result block)
assert len(emissions) == 2, f"Expected 2 emissions, got {len(emissions)}"
# First emission: cumulative list with 1 result
assert len(emissions[0]) == 1
assert emissions[0][0].id == "srvtoolu_01AAA"
assert emissions[0][0].code == "echo first"
assert emissions[0][0].outputs[0].logs == "first\n"
# Second (final) emission: cumulative list with BOTH results
# This is what stream_chunk_builder will pick as "last value wins"
assert len(emissions[1]) == 2, (
f"Expected final emission to have 2 results, got {len(emissions[1])}. "
f"IDs: {[r.id for r in emissions[1]]}"
)
assert emissions[1][0].id == "srvtoolu_01AAA"
assert emissions[1][0].code == "echo first"
assert emissions[1][0].outputs[0].logs == "first\n"
assert emissions[1][1].id == "srvtoolu_01BBB"
assert emissions[1][1].code == "echo second"
assert emissions[1][1].outputs[0].logs == "second\n"
def test_streaming_code_execution_input_assembled_from_deltas():
"""
In real Anthropic streaming, content_block_start for server_tool_use has
input: {}. The actual input arrives via input_json_delta deltas and must
be assembled at content_block_stop so the code field is populated.
This test uses realistic chunk shapes (empty input in start, partial JSON
in deltas) to exercise the input assembly path.
"""
chunks = [
{
"type": "message_start",
"message": {
"id": "msg_01XYZ",
"type": "message",
"role": "assistant",
"content": [],
"usage": {"input_tokens": 100, "output_tokens": 1},
},
},
# server_tool_use with empty input (real streaming behaviour)
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_01AAA",
"name": "code_execution",
"input": {},
},
},
# Input arrives via deltas, split across two chunks
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "input_json_delta",
"partial_json": '{"comma',
},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "input_json_delta",
"partial_json": 'nd": "echo hello"}',
},
},
{"type": "content_block_stop", "index": 0},
# Tool result
{
"type": "content_block_start",
"index": 1,
"content_block": {
"type": "bash_code_execution_tool_result",
"tool_use_id": "srvtoolu_01AAA",
"content": {
"type": "bash_code_execution_result",
"stdout": "hello\n",
"stderr": "",
"return_code": 0,
},
},
},
{"type": "content_block_stop", "index": 1},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 50},
},
]
iterator = ModelResponseIterator(None, sync_stream=True)
code_results = None
for chunk in chunks:
parsed = iterator.chunk_parser(chunk)
psf = None
if parsed.choices and parsed.choices[0].delta:
psf = getattr(parsed.choices[0].delta, "provider_specific_fields", None)
if psf and "code_interpreter_results" in psf:
code_results = psf["code_interpreter_results"]
# The code field must contain the assembled input, not be empty
assert code_results is not None, "No code_interpreter_results emitted"
assert len(code_results) == 1
assert code_results[0].id == "srvtoolu_01AAA"
assert code_results[0].code == "echo hello"
assert code_results[0].outputs[0].logs == "hello\n"
def test_empty_output_produces_null_outputs():
"""
When both stdout and stderr are empty, outputs should be None
(matching OpenAI's native behavior) rather than [{logs: ""}].
"""
chunks = [
{
"type": "message_start",
"message": {
"id": "msg_01XYZ",
"type": "message",
"role": "assistant",
"content": [],
"usage": {"input_tokens": 100, "output_tokens": 1},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_01AAA",
"name": "bash_code_execution",
"input": {"command": "true"},
},
},
{"type": "content_block_stop", "index": 0},
{
"type": "content_block_start",
"index": 1,
"content_block": {
"type": "bash_code_execution_tool_result",
"tool_use_id": "srvtoolu_01AAA",
"content": {
"type": "bash_code_execution_result",
"stdout": "",
"stderr": "",
"return_code": 0,
},
},
},
{"type": "content_block_stop", "index": 1},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 50},
},
]
iterator = ModelResponseIterator(None, sync_stream=True)
code_results = None
for chunk in chunks:
parsed = iterator.chunk_parser(chunk)
psf = None
if parsed.choices and parsed.choices[0].delta:
psf = getattr(parsed.choices[0].delta, "provider_specific_fields", None)
if psf and "code_interpreter_results" in psf:
code_results = psf["code_interpreter_results"]
assert code_results is not None, "No code_interpreter_results emitted"
assert len(code_results) == 1
assert code_results[0].id == "srvtoolu_01AAA"
assert (
code_results[0].outputs is None
), f"Expected outputs=None for empty execution, got {code_results[0].outputs}"
def test_non_bash_tool_result_skipped():
"""
Tool result types other than bash_code_execution_tool_result (e.g.
text_editor_code_execution_tool_result) should be skipped and NOT
produce code_interpreter_call items.
"""
chunks = [
{
"type": "message_start",
"message": {
"id": "msg_01XYZ",
"type": "message",
"role": "assistant",
"content": [],
"usage": {"input_tokens": 100, "output_tokens": 1},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_01AAA",
"name": "text_editor",
"input": {"command": "view", "path": "/tmp/test.py"},
},
},
{"type": "content_block_stop", "index": 0},
# text_editor result — should NOT become a code_interpreter_call
{
"type": "content_block_start",
"index": 1,
"content_block": {
"type": "text_editor_code_execution_tool_result",
"tool_use_id": "srvtoolu_01AAA",
"content": [
{"type": "text", "text": "file contents here"},
],
},
},
{"type": "content_block_stop", "index": 1},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 50},
},
]
iterator = ModelResponseIterator(None, sync_stream=True)
code_results = None
for chunk in chunks:
parsed = iterator.chunk_parser(chunk)
psf = None
if parsed.choices and parsed.choices[0].delta:
psf = getattr(parsed.choices[0].delta, "provider_specific_fields", None)
if psf and "code_interpreter_results" in psf:
code_results = psf["code_interpreter_results"]
# code_interpreter_results should be emitted but empty (no bash results)
assert (
code_results is not None
), "Expected code_interpreter_results key to be emitted"
assert (
len(code_results) == 0
), f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}"

View file

@ -0,0 +1,268 @@
"""
Tests for the Responses API _extract_tool_result_output_items path,
the non-streaming _hidden_params propagation of code_interpreter_results,
and mock end-to-end streaming integration.
"""
from unittest.mock import MagicMock
from litellm.llms.anthropic.chat.handler import ModelResponseIterator
from litellm.main import stream_chunk_builder
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.types.responses.main import (
OutputCodeInterpreterCall,
OutputCodeInterpreterCallLog,
)
from litellm.types.utils import Choices, Message, ModelResponse
def _make_model_response(code_interpreter_results=None, provider_specific_fields=None):
"""Helper to build a ModelResponse with provider_specific_fields on the message."""
psf = provider_specific_fields or {}
if code_interpreter_results is not None:
psf["code_interpreter_results"] = code_interpreter_results
msg = Message(content="test", provider_specific_fields=psf if psf else None)
choice = Choices(index=0, message=msg, finish_reason="stop")
resp = ModelResponse()
resp.choices = [choice]
return resp
def test_extract_tool_result_output_items_from_pydantic_objects():
"""Non-streaming path: code_interpreter_results are Pydantic OutputCodeInterpreterCall objects."""
items = [
OutputCodeInterpreterCall(
type="code_interpreter_call",
id="srvtoolu_01AAA",
code="echo hello",
container_id=None,
status="completed",
outputs=[OutputCodeInterpreterCallLog(type="logs", logs="hello\n")],
),
OutputCodeInterpreterCall(
type="code_interpreter_call",
id="srvtoolu_01BBB",
code="echo world",
container_id=None,
status="completed",
outputs=[OutputCodeInterpreterCallLog(type="logs", logs="world\n")],
),
]
resp = _make_model_response(code_interpreter_results=items)
result = LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(resp)
assert len(result) == 2
assert result[0].id == "srvtoolu_01AAA"
assert result[1].id == "srvtoolu_01BBB"
def test_extract_tool_result_output_items_from_dicts():
"""Streaming path: after model_dump(), code_interpreter_results are plain dicts.
_extract_tool_result_output_items reconstructs them as Pydantic objects."""
items = [
{
"type": "code_interpreter_call",
"id": "srvtoolu_01AAA",
"code": "echo hello",
"container_id": None,
"status": "completed",
"outputs": [{"type": "logs", "logs": "hello\n"}],
},
]
resp = _make_model_response(code_interpreter_results=items)
result = LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(resp)
assert len(result) == 1
assert isinstance(result[0], OutputCodeInterpreterCall)
assert result[0].id == "srvtoolu_01AAA"
def test_extract_tool_result_output_items_empty():
"""No code_interpreter_results → empty list."""
resp = _make_model_response()
result = LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(resp)
assert result == []
def test_extract_tool_result_output_items_no_provider_specific_fields():
"""Message with no provider_specific_fields → empty list."""
msg = Message(content="test")
choice = Choices(index=0, message=msg, finish_reason="stop")
resp = ModelResponse()
resp.choices = [choice]
result = LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(resp)
assert result == []
def test_in_place_substitution_preserves_ordering():
"""
function_call items matching code_interpreter_results should be replaced
in-place, preserving the original output ordering.
Simulates: [message, function_call(exec1), function_call(regular), function_call(exec2)]
Expected: [message, code_interpreter_call(exec1), function_call(regular), code_interpreter_call(exec2)]
"""
code_results = [
OutputCodeInterpreterCall(
type="code_interpreter_call",
id="srvtoolu_01AAA",
code="echo first",
container_id=None,
status="completed",
outputs=[OutputCodeInterpreterCallLog(type="logs", logs="first\n")],
),
OutputCodeInterpreterCall(
type="code_interpreter_call",
id="srvtoolu_01CCC",
code="echo third",
container_id=None,
status="completed",
outputs=[OutputCodeInterpreterCallLog(type="logs", logs="third\n")],
),
]
resp = _make_model_response(code_interpreter_results=code_results)
# Build a mock responses_output list with interleaved items
class MockItem:
def __init__(self, type, call_id=None):
self.type = type
self.call_id = call_id
msg_item = MockItem(type="message")
fc_exec1 = MockItem(type="function_call", call_id="srvtoolu_01AAA")
fc_regular = MockItem(type="function_call", call_id="srvtoolu_01BBB")
fc_exec2 = MockItem(type="function_call", call_id="srvtoolu_01CCC")
responses_output = [msg_item, fc_exec1, fc_regular, fc_exec2]
# Apply the same logic as _transform_chat_completion_choices_to_responses_output
tool_result_items = (
LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(resp)
)
if tool_result_items:
result_by_id = {
(item.get("id") if isinstance(item, dict) else item.id): item
for item in tool_result_items
}
replaced_ids = set(result_by_id.keys())
responses_output = [
(
result_by_id[getattr(item, "call_id", None)]
if (
getattr(item, "type", None) == "function_call"
and getattr(item, "call_id", None) in replaced_ids
)
else item
)
for item in responses_output
]
# Verify ordering: message, code_interpreter(AAA), function_call(BBB), code_interpreter(CCC)
assert len(responses_output) == 4
assert responses_output[0].type == "message"
assert responses_output[1].type == "code_interpreter_call"
assert responses_output[1].id == "srvtoolu_01AAA"
assert responses_output[2].type == "function_call"
assert responses_output[2].call_id == "srvtoolu_01BBB"
assert responses_output[3].type == "code_interpreter_call"
assert responses_output[3].id == "srvtoolu_01CCC"
def test_end_to_end_streaming_chunks_to_code_interpreter_output():
"""
Mock end-to-end test: Anthropic SSE chunks ModelResponseIterator
stream_chunk_builder _extract_tool_result_output_items final output
with code_interpreter_call items replacing function_call items.
This exercises the full streaming data flow without a live server.
"""
# Realistic Anthropic streaming chunks for a single code execution
raw_chunks = [
{
"type": "message_start",
"message": {
"id": "msg_01XYZ",
"type": "message",
"role": "assistant",
"content": [],
"usage": {"input_tokens": 100, "output_tokens": 1},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "server_tool_use",
"id": "srvtoolu_01AAA",
"name": "bash_code_execution",
"input": {},
},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "input_json_delta",
"partial_json": '{"command": "echo e2e_test"}',
},
},
{"type": "content_block_stop", "index": 0},
{
"type": "content_block_start",
"index": 1,
"content_block": {
"type": "bash_code_execution_tool_result",
"tool_use_id": "srvtoolu_01AAA",
"content": {
"type": "bash_code_execution_result",
"stdout": "e2e_test\n",
"stderr": "",
"return_code": 0,
},
},
},
{"type": "content_block_stop", "index": 1},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 50},
},
]
# Step 1: Parse chunks through ModelResponseIterator (Anthropic handler)
iterator = ModelResponseIterator(None, sync_stream=True)
parsed_chunks = []
for chunk in raw_chunks:
parsed = iterator.chunk_parser(chunk)
d = parsed.model_dump()
# In production, CustomStreamWrapper sets the model on each chunk;
# stream_chunk_builder requires it.
d["model"] = "claude-sonnet-4-20250514"
parsed_chunks.append(d)
# Step 2: Assemble via stream_chunk_builder (simulates end-of-stream)
assembled = stream_chunk_builder(chunks=parsed_chunks)
assert assembled is not None
# Verify stream_chunk_builder picked up code_interpreter_results via last-value-wins
psf = assembled.choices[0].message.provider_specific_fields
assert psf is not None
assert "code_interpreter_results" in psf
code_results = psf["code_interpreter_results"]
assert len(code_results) == 1
# After model_dump + stream_chunk_builder, results are plain dicts
assert code_results[0]["id"] == "srvtoolu_01AAA"
assert code_results[0]["code"] == "echo e2e_test"
# Step 3: Extract via _extract_tool_result_output_items (Responses API layer)
tool_result_items = (
LiteLLMCompletionResponsesConfig._extract_tool_result_output_items(assembled)
)
assert len(tool_result_items) == 1
item = tool_result_items[0]
# Items are reconstructed as Pydantic OutputCodeInterpreterCall objects
assert isinstance(item, OutputCodeInterpreterCall)
assert item.type == "code_interpreter_call"
assert item.id == "srvtoolu_01AAA"
assert item.code == "echo e2e_test"
assert item.outputs[0].logs == "e2e_test\n"

View file

@ -1,20 +1,27 @@
"""
Tests for Anthropic OAuth token handling in common_utils.
Tests for Anthropic authentication and environment variable handling in common_utils.
Verifies that OAuth tokens (sk-ant-oat*) are sent via Authorization: Bearer
instead of x-api-key, per Anthropic's OAuth specification.
Verifies that:
- OAuth tokens (sk-ant-oat*) produce Authorization: Bearer headers with OAuth beta flags.
- Regular API keys produce x-api-key headers.
- ANTHROPIC_AUTH_TOKEN produces Authorization: Bearer headers,
matching the official Anthropic SDK behavior.
- ANTHROPIC_BASE_URL is used as a fallback for base URL resolution.
- ANTHROPIC_API_KEY / ANTHROPIC_API_BASE take precedence over their aliases.
"""
import os
import sys
from unittest.mock import patch
sys.path.insert(
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))
)
# Fake OAuth token for testing (not a real secret)
# Fake tokens for testing (not real secrets)
FAKE_OAUTH_TOKEN = "sk-ant-oat01-fake-token-for-testing-123456789abcdef"
FAKE_REGULAR_KEY = "sk-ant-api03-regular-key-for-testing-123456789"
FAKE_AUTH_TOKEN = "sk-ant-aut01-fake-auth-token-for-testing-123456789"
class TestOptionallyHandleAnthropicOAuth:
@ -697,3 +704,430 @@ class TestProxyOAuthHeaderForwarding:
assert cleaned["authorization"] == oauth_token
# Proxy key must be stripped
assert "x-litellm-api-key" not in cleaned
class TestGetAnthropicHeadersWithAuthToken:
"""Tests for get_anthropic_headers with auth_token parameter."""
def test_auth_token_uses_bearer_header(self):
"""auth_token should produce Authorization: Bearer header."""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
config = AnthropicModelInfo()
headers = config.get_anthropic_headers(
api_key=None,
auth_token=FAKE_AUTH_TOKEN,
computer_tool_used=False,
prompt_caching_set=False,
pdf_used=False,
is_vertex_request=False,
)
assert headers["authorization"] == f"Bearer {FAKE_AUTH_TOKEN}"
assert "x-api-key" not in headers
# auth_token should NOT set OAuth-specific flags
assert "anthropic-dangerous-direct-browser-access" not in headers
def test_auth_token_includes_standard_headers(self):
"""auth_token path should include standard Anthropic headers."""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
config = AnthropicModelInfo()
headers = config.get_anthropic_headers(
api_key=None,
auth_token=FAKE_AUTH_TOKEN,
computer_tool_used=False,
prompt_caching_set=False,
pdf_used=False,
is_vertex_request=False,
)
assert headers["anthropic-version"] == "2023-06-01"
assert headers["accept"] == "application/json"
assert headers["content-type"] == "application/json"
def test_api_key_takes_precedence_over_auth_token(self):
"""When both api_key and auth_token are provided, api_key wins."""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
config = AnthropicModelInfo()
headers = config.get_anthropic_headers(
api_key=FAKE_REGULAR_KEY,
auth_token=FAKE_AUTH_TOKEN,
computer_tool_used=False,
prompt_caching_set=False,
pdf_used=False,
is_vertex_request=False,
)
assert headers["x-api-key"] == FAKE_REGULAR_KEY
assert "authorization" not in headers
class TestValidateEnvironmentAuthToken:
"""Tests for validate_environment with auth_token resolution."""
def test_auth_token_env_var_produces_bearer_header(self):
"""validate_environment should use Bearer auth when only ANTHROPIC_AUTH_TOKEN is set."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
config = AnthropicModelInfo()
with mock_patch.dict(
"os.environ",
{"ANTHROPIC_AUTH_TOKEN": FAKE_AUTH_TOKEN},
clear=True,
):
headers = config.validate_environment(
headers={},
model="claude-sonnet-4-5-20250929",
messages=[{"role": "user", "content": "Hello"}],
optional_params={},
litellm_params={},
api_key=None,
api_base=None,
)
assert headers["authorization"] == f"Bearer {FAKE_AUTH_TOKEN}"
assert "x-api-key" not in headers
assert "anthropic-dangerous-direct-browser-access" not in headers
def test_api_key_param_takes_precedence_over_auth_token_env_var(self):
"""validate_environment should prefer explicit api_key over ANTHROPIC_AUTH_TOKEN."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
config = AnthropicModelInfo()
with mock_patch.dict(
"os.environ",
{"ANTHROPIC_AUTH_TOKEN": FAKE_AUTH_TOKEN},
clear=True,
):
headers = config.validate_environment(
headers={},
model="claude-sonnet-4-5-20250929",
messages=[{"role": "user", "content": "Hello"}],
optional_params={},
litellm_params={},
api_key=FAKE_REGULAR_KEY,
api_base=None,
)
assert headers["x-api-key"] == FAKE_REGULAR_KEY
assert "authorization" not in headers
def test_raises_when_no_credentials(self):
"""validate_environment should raise when neither API key nor auth token is available."""
from unittest.mock import patch as mock_patch
import pytest
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
config = AnthropicModelInfo()
with mock_patch.dict("os.environ", {}, clear=True):
with pytest.raises(
Exception, match="ANTHROPIC_API_KEY.*ANTHROPIC_AUTH_TOKEN"
):
config.validate_environment(
headers={},
model="claude-sonnet-4-5-20250929",
messages=[{"role": "user", "content": "Hello"}],
optional_params={},
litellm_params={},
api_key=None,
api_base=None,
)
def test_resolves_api_key_from_env_when_param_is_none(self):
"""validate_environment should resolve ANTHROPIC_API_KEY from env when api_key param is None."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
config = AnthropicModelInfo()
with mock_patch.dict(
"os.environ",
{"ANTHROPIC_API_KEY": FAKE_REGULAR_KEY},
clear=True,
):
headers = config.validate_environment(
headers={},
model="claude-sonnet-4-5-20250929",
messages=[{"role": "user", "content": "Hello"}],
optional_params={},
litellm_params={},
api_key=None,
api_base=None,
)
assert headers["x-api-key"] == FAKE_REGULAR_KEY
assert "authorization" not in headers
class TestGetAuthToken:
"""Tests for AnthropicModelInfo.get_auth_token() static method."""
def test_returns_env_var_value(self):
"""get_auth_token returns the ANTHROPIC_AUTH_TOKEN env var value."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
with mock_patch.dict(
"os.environ", {"ANTHROPIC_AUTH_TOKEN": FAKE_AUTH_TOKEN}, clear=True
):
assert AnthropicModelInfo.get_auth_token() == FAKE_AUTH_TOKEN
def test_returns_none_when_not_set(self):
"""get_auth_token returns None when ANTHROPIC_AUTH_TOKEN is not set."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
with mock_patch.dict("os.environ", {}, clear=True):
assert AnthropicModelInfo.get_auth_token() is None
def test_explicit_param_takes_precedence(self):
"""Explicit auth_token param takes precedence over env var."""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
explicit_token = "sk-ant-aut01-explicit-token-override-123456789"
assert AnthropicModelInfo.get_auth_token(explicit_token) == explicit_token
class TestGetAuthHeader:
"""Tests for AnthropicModelInfo.get_auth_header() centralized helper."""
def test_returns_x_api_key_when_api_key_provided(self):
"""Explicit api_key param should return x-api-key header."""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
result = AnthropicModelInfo.get_auth_header(api_key=FAKE_REGULAR_KEY)
assert result == {"x-api-key": FAKE_REGULAR_KEY}
def test_returns_x_api_key_from_env(self):
"""ANTHROPIC_API_KEY env var should return x-api-key header."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
with mock_patch.dict(
"os.environ",
{"ANTHROPIC_API_KEY": FAKE_REGULAR_KEY},
clear=True,
):
result = AnthropicModelInfo.get_auth_header()
assert result == {"x-api-key": FAKE_REGULAR_KEY}
def test_returns_bearer_from_auth_token_env(self):
"""ANTHROPIC_AUTH_TOKEN env var should return Authorization: Bearer header."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
with mock_patch.dict(
"os.environ",
{"ANTHROPIC_AUTH_TOKEN": FAKE_AUTH_TOKEN},
clear=True,
):
result = AnthropicModelInfo.get_auth_header()
assert result == {"authorization": f"Bearer {FAKE_AUTH_TOKEN}"}
def test_api_key_takes_precedence_over_auth_token(self):
"""ANTHROPIC_API_KEY should take precedence over ANTHROPIC_AUTH_TOKEN."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
with mock_patch.dict(
"os.environ",
{
"ANTHROPIC_API_KEY": FAKE_REGULAR_KEY,
"ANTHROPIC_AUTH_TOKEN": FAKE_AUTH_TOKEN,
},
clear=True,
):
result = AnthropicModelInfo.get_auth_header()
assert result == {"x-api-key": FAKE_REGULAR_KEY}
def test_explicit_api_key_overrides_env_auth_token(self):
"""Explicit api_key param should override ANTHROPIC_AUTH_TOKEN env var."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
with mock_patch.dict(
"os.environ",
{"ANTHROPIC_AUTH_TOKEN": FAKE_AUTH_TOKEN},
clear=True,
):
result = AnthropicModelInfo.get_auth_header(api_key=FAKE_REGULAR_KEY)
assert result == {"x-api-key": FAKE_REGULAR_KEY}
def test_returns_none_when_no_credentials(self):
"""Should return None when neither api_key nor auth_token is available."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
with mock_patch.dict("os.environ", {}, clear=True):
result = AnthropicModelInfo.get_auth_header()
assert result is None
def test_oauth_token_uses_bearer_not_x_api_key(self):
"""OAuth token (sk-ant-oat*) should return Authorization: Bearer, not x-api-key."""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
result = AnthropicModelInfo.get_auth_header(api_key=FAKE_OAUTH_TOKEN)
assert result == {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"}
def test_oauth_token_from_env_uses_bearer(self):
"""OAuth token in ANTHROPIC_API_KEY env var should return Authorization: Bearer."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
with mock_patch.dict(
"os.environ",
{"ANTHROPIC_API_KEY": FAKE_OAUTH_TOKEN},
clear=True,
):
result = AnthropicModelInfo.get_auth_header()
assert result == {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"}
class TestGetApiBaseFallbackChain:
"""Tests for AnthropicModelInfo.get_api_base() fallback to ANTHROPIC_BASE_URL."""
def test_explicit_param_takes_precedence(self):
"""Explicit api_base param takes precedence over all env vars."""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
assert (
AnthropicModelInfo.get_api_base("https://explicit.example.com")
== "https://explicit.example.com"
)
def test_defaults_to_anthropic_api(self):
"""get_api_base returns the default Anthropic API base when no env vars are set."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
with mock_patch.dict("os.environ", {}, clear=True):
assert AnthropicModelInfo.get_api_base() == "https://api.anthropic.com"
def test_api_base_env_preferred_over_base_url_env(self):
"""ANTHROPIC_API_BASE takes precedence over ANTHROPIC_BASE_URL."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
with mock_patch.dict(
"os.environ",
{
"ANTHROPIC_API_BASE": "https://api-base.example.com",
"ANTHROPIC_BASE_URL": "https://base-url.example.com",
},
clear=True,
):
assert AnthropicModelInfo.get_api_base() == "https://api-base.example.com"
def test_falls_back_to_base_url_env(self):
"""get_api_base falls back to ANTHROPIC_BASE_URL when ANTHROPIC_API_BASE is not set."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
with mock_patch.dict(
"os.environ",
{"ANTHROPIC_BASE_URL": "https://base-url.example.com"},
clear=True,
):
assert AnthropicModelInfo.get_api_base() == "https://base-url.example.com"
class TestPassthroughAuthToken:
"""Tests for passthrough messages endpoint with ANTHROPIC_AUTH_TOKEN."""
def test_passthrough_auth_token_uses_bearer_header(self):
"""Passthrough endpoint should use Bearer auth when only ANTHROPIC_AUTH_TOKEN is set."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
AnthropicMessagesConfig,
)
config = AnthropicMessagesConfig()
with mock_patch.dict(
"os.environ", {"ANTHROPIC_AUTH_TOKEN": FAKE_AUTH_TOKEN}, clear=True
):
updated_headers, _ = config.validate_anthropic_messages_environment(
headers={},
model="claude-sonnet-4-5-20250929",
messages=[{"role": "user", "content": "Hello"}],
optional_params={},
litellm_params={},
api_key=None,
api_base=None,
)
assert updated_headers["authorization"] == f"Bearer {FAKE_AUTH_TOKEN}"
assert "x-api-key" not in updated_headers
assert "anthropic-dangerous-direct-browser-access" not in updated_headers
def test_passthrough_api_key_takes_precedence(self):
"""Passthrough endpoint should prefer ANTHROPIC_API_KEY over ANTHROPIC_AUTH_TOKEN."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
AnthropicMessagesConfig,
)
config = AnthropicMessagesConfig()
with mock_patch.dict(
"os.environ",
{"ANTHROPIC_API_KEY": FAKE_REGULAR_KEY, "ANTHROPIC_AUTH_TOKEN": FAKE_AUTH_TOKEN},
clear=True,
):
updated_headers, _ = config.validate_anthropic_messages_environment(
headers={},
model="claude-sonnet-4-5-20250929",
messages=[{"role": "user", "content": "Hello"}],
optional_params={},
litellm_params={},
api_key=None,
api_base=None,
)
assert updated_headers["x-api-key"] == FAKE_REGULAR_KEY
assert "authorization" not in updated_headers
def test_passthrough_get_complete_url_honours_base_url_env(self):
"""get_complete_url should use ANTHROPIC_BASE_URL when api_base is None."""
from unittest.mock import patch as mock_patch
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
AnthropicMessagesConfig,
)
config = AnthropicMessagesConfig()
with mock_patch.dict(
"os.environ",
{"ANTHROPIC_BASE_URL": "https://custom.example.com"},
clear=True,
):
url = config.get_complete_url(
api_base=None,
api_key=FAKE_REGULAR_KEY,
model="claude-sonnet-4-5-20250929",
optional_params={},
litellm_params={},
)
assert url == "https://custom.example.com/v1/messages"

View file

@ -336,7 +336,7 @@ class TestAnthropicFilesHandler:
"extra_body": None
}
with patch.object(handler.anthropic_model_info, "get_api_key", return_value=None):
with patch.object(handler.anthropic_model_info, "get_auth_header", return_value=None):
with pytest.raises(ValueError, match="Missing Anthropic API Key"):
await handler.afile_content(
file_content_request=file_content_request,

View file

@ -1666,3 +1666,141 @@ async def test_oauth_authorize_prefers_request_scope_over_server_config():
redirect_url = response.headers["location"]
assert "scope=custom_scope1+custom_scope2" in redirect_url or "scope=custom_scope1%20custom_scope2" in redirect_url
assert "default_scope" not in redirect_url
@pytest.mark.asyncio
async def test_token_endpoint_refresh_token_grant():
"""Test that token endpoint supports refresh_token grant type."""
try:
from fastapi import Request
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
token_endpoint,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy._types import MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
except ImportError:
pytest.skip("MCP discoverable endpoints not available")
# Clear registry
global_mcp_server_manager.registry.clear()
# Create mock OAuth2 server
oauth2_server = MCPServer(
server_id="google_mcp",
name="google_mcp",
server_name="google_mcp",
alias="google_mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="test_client_id",
client_secret="test_secret",
authorization_url="https://accounts.google.com/o/oauth2/v2/auth",
token_url="https://oauth2.googleapis.com/token",
scopes=["openid", "email"],
)
global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server
mock_request = MagicMock(spec=Request)
mock_request.base_url = "https://proxy.litellm.example/"
mock_request.headers = {}
# Mock httpx client response with new tokens
mock_response = MagicMock()
mock_response.json.return_value = {
"access_token": "new_access_token",
"token_type": "Bearer",
"expires_in": 3599,
"refresh_token": "new_refresh_token",
}
mock_response.raise_for_status = MagicMock()
mock_async_client = MagicMock()
mock_async_client.post = AsyncMock(return_value=mock_response)
with patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client"
) as mock_get_client:
mock_get_client.return_value = mock_async_client
response = await token_endpoint(
request=mock_request,
grant_type="refresh_token",
code=None,
redirect_uri=None,
client_id="test_client_id",
mcp_server_name="google_mcp",
client_secret="test_secret",
refresh_token="rt-test",
scope="openid email",
)
# Verify the POST was called with refresh_token grant data
mock_async_client.post.assert_called_once()
call_args = mock_async_client.post.call_args
assert call_args[1]["data"]["grant_type"] == "refresh_token"
assert call_args[1]["data"]["refresh_token"] == "rt-test"
assert call_args[1]["data"]["client_id"] == "test_client_id"
assert call_args[1]["data"]["client_secret"] == "test_secret"
assert call_args[1]["data"]["scope"] == "openid email"
# Verify response contains the new tokens
import json
token_data = json.loads(response.body)
assert token_data["access_token"] == "new_access_token"
assert token_data["refresh_token"] == "new_refresh_token"
@pytest.mark.asyncio
async def test_token_endpoint_authorization_code_missing_code():
"""Test that authorization_code grant rejects missing code param."""
try:
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
exchange_token_with_server,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy._types import MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
except ImportError:
pytest.skip("MCP discoverable endpoints not available")
global_mcp_server_manager.registry.clear()
server = MCPServer(
server_id="test_server",
name="test_server",
server_name="test_server",
alias="test_server",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
client_id="cid",
token_url="https://example.com/token",
)
global_mcp_server_manager.registry[server.server_id] = server
mock_request = MagicMock()
mock_request.base_url = "https://proxy.example/"
mock_request.headers = {}
with pytest.raises(HTTPException) as exc_info:
await exchange_token_with_server(
request=mock_request,
mcp_server=server,
grant_type="authorization_code",
code=None,
redirect_uri="https://example.com/cb",
client_id="cid",
client_secret=None,
code_verifier=None,
)
assert exc_info.value.status_code == 400
assert "code is required" in str(exc_info.value.detail)

View file

@ -1,6 +1,6 @@
import os
import sys
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import pytest
from fastapi import FastAPI
@ -11,6 +11,7 @@ sys.path.insert(
)
from litellm.proxy.discovery_endpoints.ui_discovery_endpoints import router
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
def test_ui_discovery_endpoints_with_defaults():
@ -245,9 +246,9 @@ def test_ui_discovery_endpoints_with_admin_ui_enabled():
patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \
patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \
patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False):
response = client.get("/.well-known/litellm-ui-config")
assert response.status_code == 200
data = response.json()
assert data["server_root_path"] == "/"
@ -256,3 +257,53 @@ def test_ui_discovery_endpoints_with_admin_ui_enabled():
assert data["admin_ui_disabled"] is False
assert data["sso_configured"] is False
def test_ui_discovery_endpoints_is_control_plane_true_when_workers_configured():
app = FastAPI()
app.include_router(router)
client = TestClient(app)
mock_config = MagicMock()
mock_config.worker_registry = [
WorkerRegistryEntry(
worker_id="team-a", name="Team A", url="https://worker-1:4001"
),
]
with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \
patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \
patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \
patch("litellm.proxy.proxy_server.proxy_config", mock_config), \
patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False):
response = client.get("/.well-known/litellm-ui-config")
assert response.status_code == 200
data = response.json()
assert data["is_control_plane"] is True
assert len(data["workers"]) == 1
assert data["workers"][0]["worker_id"] == "team-a"
assert data["workers"][0]["name"] == "Team A"
assert data["workers"][0]["url"] == "https://worker-1:4001"
def test_ui_discovery_endpoints_is_control_plane_false_when_no_workers():
app = FastAPI()
app.include_router(router)
client = TestClient(app)
mock_config = MagicMock()
mock_config.worker_registry = []
with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \
patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \
patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \
patch("litellm.proxy.proxy_server.proxy_config", mock_config), \
patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False):
response = client.get("/.well-known/litellm-ui-config")
assert response.status_code == 200
data = response.json()
assert data["is_control_plane"] is False
assert data["workers"] == []

View file

@ -1458,10 +1458,6 @@ async def test_unblock_key_supports_both_sk_and_hashed_tokens(monkeypatch):
return_value=mock_key_record
)
# Mock get_key_object and _cache_key_object functions
mock_key_object = MagicMock()
mock_key_object.blocked = True # Initially blocked
# Mock hash_token function
def mock_hash_token(token):
if token == "sk-test123456789":
@ -1482,19 +1478,12 @@ async def test_unblock_key_supports_both_sk_and_hashed_tokens(monkeypatch):
) # Disable audit logs for simpler test
# Mock get_key_object and _cache_key_object
async def mock_get_key_object(**kwargs):
return mock_key_object
async def mock_cache_key_object(**kwargs):
async def mock_delete_cache_key_object(**kwargs):
pass
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_key_object",
mock_get_key_object,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._cache_key_object",
mock_cache_key_object,
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
mock_delete_cache_key_object,
)
# Create mock request and user auth
@ -1519,11 +1508,9 @@ async def test_unblock_key_supports_both_sk_and_hashed_tokens(monkeypatch):
)
assert result == mock_key_record
assert mock_key_object.blocked == False # Should be updated to unblocked
# Reset mocks for second test
mock_prisma_client.db.litellm_verificationtoken.update.reset_mock()
mock_key_object.blocked = True # Reset to blocked state
# Test Case 2: Using already hashed token
hashed_token_request = BlockKeyRequest(key=test_hashed_token)
@ -1541,7 +1528,6 @@ async def test_unblock_key_supports_both_sk_and_hashed_tokens(monkeypatch):
)
assert result == mock_key_record
assert mock_key_object.blocked == False # Should be updated to unblocked
@pytest.mark.asyncio
@ -1579,6 +1565,249 @@ async def test_unblock_key_invalid_key_format(monkeypatch):
assert "Invalid key format" in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_block_key_nonexistent_key_returns_404(monkeypatch):
"""
Test that block_key returns 404 (not misleading 401) when the key
doesn't exist in the database, even when the caller is authenticated
as a proxy admin.
Previously, block_key would call get_key_object() for cache refresh,
which raised a 401 ProxyException with 'Authentication Error' making
it look like an auth failure when it was really a missing-key error.
"""
from litellm.proxy._types import BlockKeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import block_key
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
# find_unique returns None → key does not exist
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=None
)
def mock_hash_token(token):
return "abcd1234" * 8 # 64-char hex
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
)
monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
monkeypatch.setattr("litellm.store_audit_logs", False)
mock_request = MagicMock()
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
)
data = BlockKeyRequest(key="sk-does-not-exist-key")
with pytest.raises(ProxyException) as exc_info:
await block_key(
data=data,
http_request=mock_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
assert exc_info.value.code == "404"
assert "not found" in str(exc_info.value.message).lower()
# Must NOT contain "Authentication Error"
assert "Authentication Error" not in str(exc_info.value.message)
# update should never be called since the key doesn't exist
mock_prisma_client.db.litellm_verificationtoken.update.assert_not_called()
@pytest.mark.asyncio
async def test_unblock_key_nonexistent_key_returns_404(monkeypatch):
"""
Test that unblock_key returns 404 (not misleading 401) when the key
doesn't exist in the database.
"""
from litellm.proxy._types import BlockKeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import (
unblock_key,
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
# find_unique returns None → key does not exist
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=None
)
def mock_hash_token(token):
return "abcd1234" * 8
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
)
monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
monkeypatch.setattr("litellm.store_audit_logs", False)
mock_request = MagicMock()
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
)
data = BlockKeyRequest(key="sk-does-not-exist-key")
with pytest.raises(ProxyException) as exc_info:
await unblock_key(
data=data,
http_request=mock_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
assert exc_info.value.code == "404"
assert "not found" in str(exc_info.value.message).lower()
assert "Authentication Error" not in str(exc_info.value.message)
mock_prisma_client.db.litellm_verificationtoken.update.assert_not_called()
@pytest.mark.asyncio
async def test_update_key_nonexistent_key_returns_404(monkeypatch):
"""
Test that update_key_fn returns 404 (not misleading 401) when the body
key doesn't exist in the database, even when the caller is authenticated
as a proxy admin via the Authorization header.
"""
from litellm.proxy.management_endpoints.key_management_endpoints import (
update_key_fn,
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
# find_unique returns None → key does not exist
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=None
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
mock_request = MagicMock()
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
)
data = UpdateKeyRequest(key="sk-does-not-exist-key")
with pytest.raises(ProxyException) as exc_info:
await update_key_fn(
request=mock_request,
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
assert exc_info.value.code == "404"
assert "not found" in str(exc_info.value.message).lower()
assert "Authentication Error" not in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_block_key_existing_key_succeeds(monkeypatch):
"""
Test that block_key successfully blocks an existing key and
invalidates the cache entry.
"""
from litellm.proxy._types import BlockKeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import block_key
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
mock_proxy_logging_obj = MagicMock()
test_hashed_token = "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
mock_key_record = MagicMock()
mock_key_record.token = test_hashed_token
mock_key_record.blocked = False
mock_key_record.model_dump_json.return_value = (
f'{{"token": "{test_hashed_token}", "blocked": false}}'
)
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=mock_key_record
)
mock_updated_record = MagicMock()
mock_updated_record.token = test_hashed_token
mock_updated_record.blocked = True
mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
return_value=mock_updated_record
)
def mock_hash_token(token):
if token.startswith("sk-"):
return test_hashed_token
return token
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
)
monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
monkeypatch.setattr("litellm.store_audit_logs", False)
# Mock _delete_cache_key_object
async def mock_delete_cache_key_object(**kwargs):
pass
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
mock_delete_cache_key_object,
)
mock_request = MagicMock()
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
)
data = BlockKeyRequest(key="sk-test123456789")
result = await block_key(
data=data,
http_request=mock_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
# Verify the key was found and updated
mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once_with(
where={"token": test_hashed_token}
)
mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once_with(
where={"token": test_hashed_token}, data={"blocked": True}
)
assert result == mock_updated_record
@pytest.mark.asyncio
async def test_validate_key_team_change_with_member_permissions():
"""
@ -4871,14 +5100,16 @@ async def test_validate_max_budget():
async def test_get_and_validate_existing_key():
"""
Test _get_and_validate_existing_key helper function.
Tests:
1. Successfully retrieve existing key
2. Key not found raises HTTPException
2. Key not found raises ProxyException
3. Database not connected raises HTTPException
"""
from fastapi import HTTPException
from litellm.proxy._types import ProxyException
# Test Case 1: Successfully retrieve existing key
mock_prisma_client = AsyncMock()
mock_key = LiteLLM_VerificationToken(
@ -4887,39 +5118,49 @@ async def test_get_and_validate_existing_key():
models=["gpt-4"],
team_id=None,
)
mock_prisma_client.get_data = AsyncMock(return_value=mock_key)
result = await _get_and_validate_existing_key(
token="test-key-123",
prisma_client=mock_prisma_client,
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=mock_key
)
assert result == mock_key
mock_prisma_client.get_data.assert_called_once_with(
token="test-key-123",
table_name="key",
query_type="find_unique",
)
# Test Case 2: Key not found raises HTTPException
mock_prisma_client.get_data = AsyncMock(return_value=None)
with pytest.raises(HTTPException) as exc_info:
await _get_and_validate_existing_key(
token="non-existent-key",
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
return_value="hashed-test-key-123",
):
result = await _get_and_validate_existing_key(
token="test-key-123",
prisma_client=mock_prisma_client,
)
assert exc_info.value.status_code == 404
assert "Key not found" in str(exc_info.value.detail)
assert result == mock_key
mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once_with(
where={"token": "hashed-test-key-123"}
)
# Test Case 2: Key not found raises ProxyException
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=None
)
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
return_value="hashed-non-existent-key",
):
with pytest.raises(ProxyException) as exc_info:
await _get_and_validate_existing_key(
token="non-existent-key",
prisma_client=mock_prisma_client,
)
assert str(exc_info.value.code) == "404"
assert "Key not found" in exc_info.value.message
# Test Case 3: Database not connected raises HTTPException
with pytest.raises(HTTPException) as exc_info:
await _get_and_validate_existing_key(
token="test-key-123",
prisma_client=None,
)
assert exc_info.value.status_code == 500
assert "Database not connected" in str(exc_info.value.detail)
@ -4960,75 +5201,82 @@ async def test_process_single_key_update():
"tags": ["production"],
}
mock_prisma_client.get_data = AsyncMock(return_value=existing_key)
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=existing_key
)
mock_updated_key_obj = MagicMock()
mock_updated_key_obj.model_dump.return_value = updated_key_data
mock_prisma_client.update_data = AsyncMock(
return_value={"data": mock_updated_key_obj}
)
# Mock prepare_key_update_data
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data"
) as mock_prepare:
mock_prepare.return_value = {"max_budget": 100.0, "tags": ["production"]}
# Mock TeamMemberPermissionChecks
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint"
) as mock_permission_check:
mock_permission_check.return_value = None
# Mock _delete_cache_key_object
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
) as mock_delete_cache:
mock_delete_cache.return_value = None
# Mock hash_token (imported from litellm.proxy._types)
with patch(
"litellm.proxy._types.hash_token"
) as mock_hash:
mock_hash.return_value = "hashed-test-key-123"
# Mock KeyManagementEventHooks
# Mock _hash_token_if_needed
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
return_value="hashed-test-key-123",
):
# Create update request
key_update_item = BulkUpdateKeyRequestItem(
key="test-key-123",
max_budget=100.0,
tags=["production"],
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
# Call the function
result = await _process_single_key_update(
key_update_item=key_update_item,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
prisma_client=mock_prisma_client,
user_api_key_cache=mock_user_api_key_cache,
proxy_logging_obj=mock_proxy_logging_obj,
llm_router=mock_llm_router,
)
# Verify results
assert result is not None
assert "token" not in result # Token should be removed
assert result.get("max_budget") == 100.0
assert result.get("tags") == ["production"]
# Verify mocks were called
mock_prisma_client.get_data.assert_called_once()
mock_prisma_client.update_data.assert_called_once()
mock_delete_cache.assert_called_once()
# Mock KeyManagementEventHooks
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
):
# Create update request
key_update_item = BulkUpdateKeyRequestItem(
key="test-key-123",
max_budget=100.0,
tags=["production"],
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
# Call the function
result = await _process_single_key_update(
key_update_item=key_update_item,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
prisma_client=mock_prisma_client,
user_api_key_cache=mock_user_api_key_cache,
proxy_logging_obj=mock_proxy_logging_obj,
llm_router=mock_llm_router,
)
# Verify results
assert result is not None
assert "token" not in result # Token should be removed
assert result.get("max_budget") == 100.0
assert result.get("tags") == ["production"]
# Verify mocks were called
mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once()
mock_prisma_client.update_data.assert_called_once()
mock_delete_cache.assert_called_once()
@pytest.mark.asyncio
@ -5090,7 +5338,7 @@ async def test_bulk_update_keys_success(monkeypatch):
"tags": ["staging"],
}
mock_prisma_client.get_data = AsyncMock(
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
side_effect=[existing_key_1, existing_key_2]
)
mock_updated_key_1_obj = MagicMock()
@ -5103,7 +5351,7 @@ async def test_bulk_update_keys_success(monkeypatch):
{"data": mock_updated_key_2_obj},
]
)
# Patch dependencies
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
@ -5115,7 +5363,7 @@ async def test_bulk_update_keys_success(monkeypatch):
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_llm_router)
# Mock helper functions
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data"
@ -5124,7 +5372,7 @@ async def test_bulk_update_keys_success(monkeypatch):
{"max_budget": 100.0, "tags": ["production"]},
{"max_budget": 200.0, "tags": ["staging"]},
]
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint"
):
@ -5135,45 +5383,49 @@ async def test_bulk_update_keys_success(monkeypatch):
"litellm.proxy._types.hash_token"
) as mock_hash:
mock_hash.side_effect = ["hashed-key-1", "hashed-key-2"]
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
side_effect=["hashed-key-1", "hashed-key-2"],
):
# Create request
request_data = BulkUpdateKeyRequest(
keys=[
BulkUpdateKeyRequestItem(
key="test-key-1",
max_budget=100.0,
tags=["production"],
),
BulkUpdateKeyRequestItem(
key="test-key-2",
max_budget=200.0,
tags=["staging"],
),
]
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
# Call endpoint
response = await bulk_update_keys(
data=request_data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
# Verify response
assert response.total_requested == 2
assert len(response.successful_updates) == 2
assert len(response.failed_updates) == 0
assert response.successful_updates[0].key == "test-key-1"
assert response.successful_updates[1].key == "test-key-2"
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
):
# Create request
request_data = BulkUpdateKeyRequest(
keys=[
BulkUpdateKeyRequestItem(
key="test-key-1",
max_budget=100.0,
tags=["production"],
),
BulkUpdateKeyRequestItem(
key="test-key-2",
max_budget=200.0,
tags=["staging"],
),
]
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
# Call endpoint
response = await bulk_update_keys(
data=request_data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
# Verify response
assert response.total_requested == 2
assert len(response.successful_updates) == 2
assert len(response.failed_updates) == 0
assert response.successful_updates[0].key == "test-key-1"
assert response.successful_updates[1].key == "test-key-2"
@pytest.mark.asyncio
@ -5218,7 +5470,7 @@ async def test_bulk_update_keys_partial_failures(monkeypatch):
}
# First key exists, second key doesn't exist
mock_prisma_client.get_data = AsyncMock(
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
side_effect=[existing_key_1, None] # Second key not found
)
mock_updated_key_1_obj = MagicMock()
@ -5226,7 +5478,9 @@ async def test_bulk_update_keys_partial_failures(monkeypatch):
mock_prisma_client.update_data = AsyncMock(
return_value={"data": mock_updated_key_1_obj}
)
# Mock get_data for the error handler path (used to fetch key_info on failure)
mock_prisma_client.get_data = AsyncMock(return_value=None)
# Patch dependencies
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
@ -5238,13 +5492,13 @@ async def test_bulk_update_keys_partial_failures(monkeypatch):
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_llm_router)
# Mock helper functions
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data"
) as mock_prepare:
mock_prepare.return_value = {"max_budget": 100.0, "tags": ["production"]}
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint"
):
@ -5255,46 +5509,50 @@ async def test_bulk_update_keys_partial_failures(monkeypatch):
"litellm.proxy._types.hash_token"
) as mock_hash:
mock_hash.return_value = "hashed-key-1"
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
side_effect=["hashed-key-1", "hashed-non-existent-key"],
):
# Create request with one valid and one invalid key
request_data = BulkUpdateKeyRequest(
keys=[
BulkUpdateKeyRequestItem(
key="test-key-1",
max_budget=100.0,
tags=["production"],
),
BulkUpdateKeyRequestItem(
key="non-existent-key",
max_budget=200.0,
tags=["staging"],
),
]
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
# Call endpoint
response = await bulk_update_keys(
data=request_data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
# Verify response
assert response.total_requested == 2
assert len(response.successful_updates) == 1
assert len(response.failed_updates) == 1
assert response.successful_updates[0].key == "test-key-1"
assert response.failed_updates[0].key == "non-existent-key"
assert "Key not found" in response.failed_updates[0].failed_reason
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
):
# Create request with one valid and one invalid key
request_data = BulkUpdateKeyRequest(
keys=[
BulkUpdateKeyRequestItem(
key="test-key-1",
max_budget=100.0,
tags=["production"],
),
BulkUpdateKeyRequestItem(
key="non-existent-key",
max_budget=200.0,
tags=["staging"],
),
]
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
# Call endpoint
response = await bulk_update_keys(
data=request_data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
# Verify response
assert response.total_requested == 2
assert len(response.successful_updates) == 1
assert len(response.failed_updates) == 1
assert response.successful_updates[0].key == "test-key-1"
assert response.failed_updates[0].key == "non-existent-key"
assert "Key not found" in response.failed_updates[0].failed_reason
@pytest.mark.parametrize(
@ -7379,19 +7637,12 @@ def _setup_block_unblock_mocks(monkeypatch, mock_key_team_id=None):
monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
monkeypatch.setattr("litellm.store_audit_logs", False)
async def mock_get_key_object(**kwargs):
return mock_key_object
async def mock_cache_key_object(**kwargs):
async def mock_delete_cache_key_object(**kwargs):
pass
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.get_key_object",
mock_get_key_object,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._cache_key_object",
mock_cache_key_object,
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
mock_delete_cache_key_object,
)
return mock_prisma_client, test_hashed_token
@ -7638,16 +7889,9 @@ async def test_update_key_non_budget_fields_allowed_for_internal_user(monkeypatc
monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
async def mock_cache_key_object(**kwargs):
pass
async def mock_delete_cache_key_object(**kwargs):
pass
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._cache_key_object",
mock_cache_key_object,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
mock_delete_cache_key_object,

View file

@ -1519,6 +1519,8 @@ class TestTemporaryMCPSessionEndpoints:
client_id="client",
client_secret="secret",
code_verifier="verifier",
refresh_token=None,
scope=None,
)
assert result is exchange_response
@ -1532,6 +1534,56 @@ class TestTemporaryMCPSessionEndpoints:
client_id="client",
client_secret="secret",
code_verifier="verifier",
refresh_token=None,
scope=None,
)
@pytest.mark.asyncio
async def test_mcp_token_proxies_refresh_token_grant(self):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
mcp_token,
)
request = MagicMock()
server = generate_mock_mcp_server_config_record(server_id="server-1")
exchange_response = {"access_token": "new-token", "refresh_token": "new-rt"}
with (
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404",
return_value=server,
) as get_server,
patch(
"litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server",
AsyncMock(return_value=exchange_response),
) as exchange_mock,
):
result = await mcp_token(
request=request,
server_id="server-1",
grant_type="refresh_token",
code=None,
redirect_uri=None,
client_id="client",
client_secret="secret",
code_verifier=None,
refresh_token="rt-123",
scope=None,
)
assert result is exchange_response
get_server.assert_called_once_with("server-1")
exchange_mock.assert_awaited_once_with(
request=request,
mcp_server=server,
grant_type="refresh_token",
code=None,
redirect_uri=None,
client_id="client",
client_secret="secret",
code_verifier=None,
refresh_token="rt-123",
scope=None,
)
@pytest.mark.asyncio

View file

@ -6441,3 +6441,53 @@ async def test_list_team_v1_batches_key_queries():
assert result[0].keys == [key1, key2]
assert result[1].team_id == "team-2"
assert result[1].keys == [key3]
def test_new_team_request_accepts_team_member_budget_duration():
"""Test that NewTeamRequest does not silently drop team_member_budget_duration."""
from litellm.proxy._types import NewTeamRequest
request = NewTeamRequest(
team_member_budget=20.0,
team_member_budget_duration="30d",
)
assert request.team_member_budget == 20.0
assert request.team_member_budget_duration == "30d"
@pytest.mark.asyncio
async def test_create_team_member_budget_table_with_duration():
"""Verify that create_team_member_budget_table passes budget_duration
through to the new_budget call when team_member_budget_duration is provided."""
from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, LitellmUserRoles
from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
mock_budget_response = MagicMock(budget_id="budget-abc")
mock_admin = UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
)
data = NewTeamRequest(
team_alias="test-team",
team_member_budget=20.0,
team_member_budget_duration="30d",
)
with patch(
"litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
new_callable=AsyncMock,
return_value=mock_budget_response,
) as mock_new_budget:
result = await TeamMemberBudgetHandler.create_team_member_budget_table(
data=data,
new_team_data_json={"metadata": None},
user_api_key_dict=mock_admin,
team_member_budget=20.0,
team_member_budget_duration="30d",
)
mock_new_budget.assert_awaited_once()
budget_request = mock_new_budget.call_args.kwargs["budget_obj"]
assert budget_request.budget_duration == "30d"
assert budget_request.max_budget == 20.0
assert result["metadata"]["team_member_budget_id"] == "budget-abc"

View file

@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from fastapi import Request
from fastapi import HTTPException, Request
from litellm._uuid import uuid
@ -5160,3 +5160,99 @@ def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch):
assert result.extra_fields["missing_field"] is None
assert result.extra_fields["another_missing"] is None
class TestValidateReturnTo:
"""Tests for SSOAuthenticationHandler._validate_return_to"""
def test_rejects_when_no_control_plane_url_configured(self, monkeypatch):
"""return_to should be rejected if control_plane_url is not in general_settings."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings", {}
)
with pytest.raises(HTTPException) as exc_info:
SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
assert exc_info.value.status_code == 400
assert "not configured" in exc_info.value.detail
def test_allows_matching_origin(self, monkeypatch):
"""return_to matching the configured control_plane_url origin should pass."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com"},
)
# Should not raise
SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui?page=models")
def test_allows_matching_origin_with_trailing_slash(self, monkeypatch):
"""Trailing slash on control_plane_url should not affect origin comparison."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com/"},
)
SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
def test_rejects_prefix_attack(self, monkeypatch):
"""return_to like cp.example.com.evil.com must be rejected (not just prefix match)."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com"},
)
with pytest.raises(HTTPException) as exc_info:
SSOAuthenticationHandler._validate_return_to("https://cp.example.com.evil.com/steal")
assert exc_info.value.status_code == 400
def test_rejects_different_origin(self, monkeypatch):
"""return_to pointing to a completely different domain should be rejected."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com"},
)
with pytest.raises(HTTPException) as exc_info:
SSOAuthenticationHandler._validate_return_to("https://evil.com/phish")
assert exc_info.value.status_code == 400
def test_case_insensitive_hostname(self, monkeypatch):
"""Hostname comparison should be case-insensitive per RFC 3986."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://CP.Example.COM"},
)
# Should not raise
SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
def test_rejects_scheme_mismatch(self, monkeypatch):
"""http:// must be rejected when control_plane_url uses https://."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com"},
)
with pytest.raises(HTTPException) as exc_info:
SSOAuthenticationHandler._validate_return_to("http://cp.example.com/ui")
assert exc_info.value.status_code == 400
def test_rejects_port_mismatch(self, monkeypatch):
"""Non-default port must be rejected."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com"},
)
with pytest.raises(HTTPException) as exc_info:
SSOAuthenticationHandler._validate_return_to("https://cp.example.com:8443/ui")
assert exc_info.value.status_code == 400
def test_allows_explicit_default_port(self, monkeypatch):
"""https://host:443 should match https://host (default port normalisation)."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com"},
)
SSOAuthenticationHandler._validate_return_to("https://cp.example.com:443/ui")
def test_allows_matching_custom_port(self, monkeypatch):
"""Both sides on the same custom port should match."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com:3000"},
)
SSOAuthenticationHandler._validate_return_to("https://cp.example.com:3000/ui")

View file

@ -280,6 +280,47 @@ class TestProxyInitializationHelpers:
assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}"
mock_uvicorn_run.assert_called_once()
@patch("uvicorn.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False)
def test_proxy_default_api_version_uses_azure_default(
self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run
):
"""Proxy default api_version should match litellm.AZURE_DEFAULT_API_VERSION for consistency."""
from click.testing import CliRunner
import litellm
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with patch.dict(os.environ, clean_env, clear=True), patch.dict(
"sys.modules",
{
"proxy_server": mock_proxy_module,
"litellm.proxy.proxy_server": mock_proxy_module,
},
), patch(
"litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
) as mock_get_args:
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(run_server, ["--local", "--skip_server_startup"])
assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}"
mock_proxy_module.save_worker_config.assert_called_once()
call_kwargs = mock_proxy_module.save_worker_config.call_args[1]
assert call_kwargs["api_version"] == litellm.AZURE_DEFAULT_API_VERSION
@patch("uvicorn.run")
@patch("builtins.print")
def test_keepalive_timeout_flag(self, mock_print, mock_uvicorn_run):

View file

@ -236,6 +236,217 @@ def test_login_v2_returns_json_on_invalid_json_body(monkeypatch):
assert isinstance(data["error"], dict)
def test_login_v3_rejected_without_control_plane_url(monkeypatch):
"""v3/login returns 404 when control_plane_url is not configured."""
mock_prisma_client = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
client = TestClient(app)
response = client.post(
"/v3/login",
json={"username": "alice", "password": "secret"},
)
assert response.status_code == 404
assert "control_plane_url" in response.json()["error"]["message"]
def test_login_v3_returns_code(monkeypatch):
"""v3/login returns an opaque code, not the JWT directly."""
mock_prisma_client = MagicMock()
monkeypatch.setattr(
"litellm.proxy.auth.login_utils.authenticate_user",
AsyncMock(return_value={"user_id": "test-user"}),
)
monkeypatch.setattr(
"litellm.proxy.auth.login_utils.create_ui_token_object",
MagicMock(return_value={"user_id": "test-user"}),
)
monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com"},
)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_config = MagicMock()
mock_config.worker_registry = []
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
client = TestClient(app)
response = client.post(
"/v3/login",
json={"username": "alice", "password": "secret"},
)
assert response.status_code == 200
data = response.json()
assert "code" in data
assert data["expires_in"] == 60
assert "token" not in data
def test_login_v3_exchange_happy_path(monkeypatch):
"""Full flow: v3/login returns code, v3/login/exchange redeems it for JWT."""
mock_prisma_client = MagicMock()
monkeypatch.setattr(
"litellm.proxy.auth.login_utils.authenticate_user",
AsyncMock(return_value={"user_id": "test-user"}),
)
monkeypatch.setattr(
"litellm.proxy.auth.login_utils.create_ui_token_object",
MagicMock(return_value={"user_id": "test-user"}),
)
monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com"},
)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_config = MagicMock()
mock_config.worker_registry = []
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
client = TestClient(app)
# Step 1: login — get code
login_response = client.post(
"/v3/login",
json={"username": "alice", "password": "secret"},
)
assert login_response.status_code == 200
code = login_response.json()["code"]
# Step 2: exchange — get JWT
exchange_response = client.post(
"/v3/login/exchange",
json={"code": code},
)
assert exchange_response.status_code == 200
exchange_data = exchange_response.json()
assert exchange_data["token"] == "signed-token"
assert "redirect_url" in exchange_data
assert exchange_response.cookies.get("token") == "signed-token"
def test_login_v3_exchange_single_use(monkeypatch):
"""Code can only be redeemed once."""
mock_prisma_client = MagicMock()
monkeypatch.setattr(
"litellm.proxy.auth.login_utils.authenticate_user",
AsyncMock(return_value={"user_id": "test-user"}),
)
monkeypatch.setattr(
"litellm.proxy.auth.login_utils.create_ui_token_object",
MagicMock(return_value={"user_id": "test-user"}),
)
monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com"},
)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mock_config = MagicMock()
mock_config.worker_registry = []
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
client = TestClient(app)
login_response = client.post(
"/v3/login",
json={"username": "alice", "password": "secret"},
)
code = login_response.json()["code"]
# First exchange succeeds
first = client.post("/v3/login/exchange", json={"code": code})
assert first.status_code == 200
# Second exchange fails
second = client.post("/v3/login/exchange", json={"code": code})
assert second.status_code == 401
def test_login_v3_exchange_invalid_code(monkeypatch):
"""Random code returns 401."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com"},
)
client = TestClient(app)
response = client.post(
"/v3/login/exchange",
json={"code": "nonexistent-code"},
)
assert response.status_code == 401
def test_login_v3_exchange_rejected_without_control_plane_url(monkeypatch):
"""v3/login/exchange returns 404 when control_plane_url is not configured."""
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
client = TestClient(app)
response = client.post(
"/v3/login/exchange",
json={"code": "some-code"},
)
assert response.status_code == 404
assert "control_plane_url" in response.json()["error"]["message"]
def test_login_v3_returns_json_on_proxy_exception(monkeypatch):
"""Test that /v3/login returns JSON error when ProxyException is raised"""
from litellm.proxy._types import ProxyErrorTypes, ProxyException
mock_prisma_client = MagicMock()
mock_authenticate_user = AsyncMock(
side_effect=ProxyException(
message="Invalid credentials",
type=ProxyErrorTypes.auth_error,
param="password",
code=401,
)
)
monkeypatch.setattr(
"litellm.proxy.auth.login_utils.authenticate_user",
mock_authenticate_user,
)
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"control_plane_url": "https://cp.example.com"},
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
client = TestClient(app)
response = client.post(
"/v3/login",
json={"username": "alice", "password": "wrong"},
)
assert response.status_code == 401
assert response.headers["content-type"] == "application/json"
data = response.json()
assert "error" in data
assert data["error"]["message"] == "Invalid credentials"
assert data["error"]["type"] == "auth_error"
def test_fallback_login_has_no_deprecation_banner(client_no_auth):
response = client_no_auth.get("/fallback/login")

View file

@ -1,16 +1,10 @@
"use client";
import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView";
import { useState } from "react";
interface ProxySettings {
PROXY_BASE_URL: string;
PROXY_LOGOUT_URL: string;
LITELLM_UI_API_DOC_BASE_URL?: string | null;
}
import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings";
const APIReferencePage = () => {
const [proxySettings, setProxySettings] = useState<ProxySettings>({ PROXY_BASE_URL: "", PROXY_LOGOUT_URL: "" });
const proxySettings = useProxySettings();
return <APIReferenceView proxySettings={proxySettings} />;
};

View file

@ -195,7 +195,7 @@ const menuItems: MenuItemCfg[] = [
icon: <UserOutlined style={{ fontSize: 18 }} />,
roles: all_admin_roles,
},
{ key: "14", page: "api_ref", label: "API Reference", icon: <ApiOutlined style={{ fontSize: 18 }} /> },
{ key: "14", page: "api-reference", label: "API Reference", icon: <ApiOutlined style={{ fontSize: 18 }} /> },
{
key: "16",
page: "model-hub-table",

View file

@ -0,0 +1,34 @@
import { describe, it, expect } from "vitest";
import { createQueryKeys } from "./queryKeysFactory";
describe("createQueryKeys", () => {
const keys = createQueryKeys("books");
it("should return the resource name as the base key", () => {
expect(keys.all).toEqual(["books"]);
});
it("should generate a lists key", () => {
expect(keys.lists()).toEqual(["books", "list"]);
});
it("should generate a list key with params", () => {
expect(keys.list({ page: 1, limit: 10 })).toEqual([
"books",
"list",
{ params: { page: 1, limit: 10 } },
]);
});
it("should generate a list key with undefined params when none provided", () => {
expect(keys.list()).toEqual(["books", "list", { params: undefined }]);
});
it("should generate a details key", () => {
expect(keys.details()).toEqual(["books", "detail"]);
});
it("should generate a detail key for a specific ID", () => {
expect(keys.detail("123")).toEqual(["books", "detail", "123"]);
});
});

View file

@ -3,8 +3,8 @@ import { loginCall, LoginRequest } from "@/components/networking";
export const useLogin = () => {
return useMutation({
mutationFn: async ({ username, password }: LoginRequest) => {
const result = await loginCall(username, password);
mutationFn: async ({ username, password, useV3 }: LoginRequest) => {
const result = await loginCall(username, password, useV3);
return result;
},
});

View file

@ -0,0 +1,21 @@
import { useState, useEffect } from "react";
import { fetchProxySettings } from "@/utils/proxyUtils";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
export default function useProxySettings() {
const { accessToken } = useAuthorized();
const [proxySettings, setProxySettings] = useState({
PROXY_BASE_URL: "",
PROXY_LOGOUT_URL: "",
LITELLM_UI_API_DOC_BASE_URL: null as string | null,
});
useEffect(() => {
if (!accessToken) return;
fetchProxySettings(accessToken).then((settings) => {
if (settings) setProxySettings(settings);
});
}, [accessToken]);
return proxySettings;
}

View file

@ -28,6 +28,8 @@ const mockUIConfig: LiteLLMWellKnownUiConfig = {
proxy_base_url: "https://proxy.example.com",
auto_redirect_to_sso: true,
admin_ui_disabled: false,
is_control_plane: false,
workers: [],
};
describe("useUIConfig", () => {
@ -102,6 +104,8 @@ describe("useUIConfig", () => {
auto_redirect_to_sso: false,
sso_configured: false,
admin_ui_disabled: true,
is_control_plane: false,
workers: [],
};
// Mock successful API call with different data

View file

@ -0,0 +1,54 @@
import { render, screen } from "@testing-library/react";
import React from "react";
import { describe, expect, it, vi } from "vitest";
import TeamsHeaderTabs from "./TeamsHeaderTabs";
vi.mock("@tremor/react", () => ({
TabGroup: ({ children, ...props }: any) => <div data-testid="tab-group" {...props}>{children}</div>,
TabList: ({ children, ...props }: any) => <div data-testid="tab-list" {...props}>{children}</div>,
Tab: ({ children, ...props }: any) => <button {...props}>{children}</button>,
TabPanels: ({ children, ...props }: any) => <div data-testid="tab-panels" {...props}>{children}</div>,
Text: ({ children, ...props }: any) => <span {...props}>{children}</span>,
Icon: ({ onClick, ...props }: any) => <button data-testid="refresh-icon" onClick={onClick} />,
}));
vi.mock("@heroicons/react/outline", () => ({
RefreshIcon: () => <svg data-testid="refresh-svg" />,
}));
const renderTabs = (props: Partial<Parameters<typeof TeamsHeaderTabs>[0]> = {}) => {
const defaults = {
lastRefreshed: "",
onRefresh: vi.fn(),
userRole: "Internal User",
children: <div data-testid="panel-content">Panel</div>,
};
return render(<TeamsHeaderTabs {...defaults} {...props} />);
};
describe("TeamsHeaderTabs", () => {
it("should render 'Your Teams' and 'Available Teams' tabs", () => {
renderTabs();
expect(screen.getByText("Your Teams")).toBeInTheDocument();
expect(screen.getByText("Available Teams")).toBeInTheDocument();
});
it("should render 'Default Team Settings' tab when user is Admin", () => {
renderTabs({ userRole: "Admin" });
expect(screen.getByText("Default Team Settings")).toBeInTheDocument();
});
it("should not render 'Default Team Settings' tab for non-admin users", () => {
renderTabs({ userRole: "Internal User" });
expect(screen.queryByText("Default Team Settings")).not.toBeInTheDocument();
});
it("should display last refreshed time when provided", () => {
renderTabs({ lastRefreshed: "2024-06-01 12:00:00" });
expect(screen.getByText("Last Refreshed: 2024-06-01 12:00:00")).toBeInTheDocument();
});
});

View file

@ -0,0 +1,129 @@
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import React from "react";
import { describe, expect, it, vi } from "vitest";
import { Team } from "@/components/key_team_helpers/key_list";
import TeamsTable from "./TeamsTable";
vi.mock("@tremor/react", () => ({
Button: React.forwardRef<HTMLButtonElement, any>(({ children, ...props }, ref) =>
React.createElement("button", { ...props, ref }, children),
),
Icon: ({ onClick, ...props }: any) => <button data-testid={props["data-testid"] || "icon-btn"} onClick={onClick} aria-label={props["aria-label"]} />,
Table: ({ children }: any) => <table>{children}</table>,
TableHead: ({ children }: any) => <thead>{children}</thead>,
TableBody: ({ children }: any) => <tbody>{children}</tbody>,
TableRow: ({ children }: any) => <tr>{children}</tr>,
TableHeaderCell: ({ children }: any) => <th>{children}</th>,
TableCell: ({ children, ...props }: any) => <td {...props}>{children}</td>,
Text: ({ children }: any) => <span>{children}</span>,
}));
vi.mock("antd", () => ({
Tooltip: ({ children }: any) => <>{children}</>,
}));
vi.mock("@heroicons/react/outline", () => ({
PencilAltIcon: () => <svg data-testid="pencil-icon" />,
TrashIcon: () => <svg data-testid="trash-icon" />,
}));
vi.mock("@/utils/dataUtils", () => ({
formatNumberWithCommas: (val: number, decimals: number) =>
val != null ? val.toFixed(decimals) : "N/A",
}));
vi.mock("@/app/(dashboard)/teams/components/TeamsTable/ModelsCell", () => ({
default: ({ team }: any) => <td data-testid="models-cell">{team.models.join(",")}</td>,
}));
vi.mock("@/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell", () => ({
default: ({ team }: any) => <td data-testid="role-cell">{team.team_id}</td>,
}));
const makeTeam = (overrides: Partial<Team> = {}): Team => ({
team_id: "team-abc1234",
team_alias: "Platform",
models: ["gpt-4"],
max_budget: 500,
budget_duration: null,
tpm_limit: null,
rpm_limit: null,
organization_id: "org-1",
created_at: "2024-06-01T00:00:00Z",
keys: [],
members_with_roles: [],
spend: 123.4567,
...overrides,
});
const defaultPerTeamInfo = {
"team-abc1234": {
keys: [{ token: "tok-1" } as any, { token: "tok-2" } as any],
team_info: {
members_with_roles: [{ user_id: "u1", role: "admin" } as any],
},
},
};
const renderTable = (overrides: Partial<Parameters<typeof TeamsTable>[0]> = {}) => {
const defaults = {
teams: [makeTeam()],
currentOrg: null,
perTeamInfo: defaultPerTeamInfo,
userRole: "Admin",
userId: "user-1",
setSelectedTeamId: vi.fn(),
setEditTeam: vi.fn(),
onDeleteTeam: vi.fn(),
};
return render(<TeamsTable {...defaults} {...overrides} />);
};
describe("TeamsTable", () => {
it("should render table headers", () => {
renderTable();
expect(screen.getByText("Team Name")).toBeInTheDocument();
expect(screen.getByText("Team ID")).toBeInTheDocument();
expect(screen.getByText("Created")).toBeInTheDocument();
expect(screen.getByText("Spend (USD)")).toBeInTheDocument();
expect(screen.getByText("Budget (USD)")).toBeInTheDocument();
expect(screen.getByText("Models")).toBeInTheDocument();
expect(screen.getByText("Organization")).toBeInTheDocument();
expect(screen.getByText("Your Role")).toBeInTheDocument();
expect(screen.getByText("Info")).toBeInTheDocument();
});
it("should render team rows with team data", () => {
renderTable();
expect(screen.getByText("Platform")).toBeInTheDocument();
expect(screen.getByText("team-ab...")).toBeInTheDocument();
expect(screen.getByText("org-1")).toBeInTheDocument();
});
it("should show edit and delete icons for Admin users", () => {
renderTable({ userRole: "Admin" });
expect(screen.getAllByTestId("icon-btn").length).toBeGreaterThanOrEqual(2);
});
it("should not show edit and delete icons for non-Admin users", () => {
renderTable({ userRole: "Internal User" });
// Only the team ID button should be present, no icon-btn for edit/delete
const iconBtns = screen.queryAllByTestId("icon-btn");
expect(iconBtns).toHaveLength(0);
});
it("should call setSelectedTeamId when team ID button is clicked", async () => {
const user = userEvent.setup();
const setSelectedTeamId = vi.fn();
renderTable({ setSelectedTeamId });
await user.click(screen.getByText("team-ab..."));
expect(setSelectedTeamId).toHaveBeenCalledWith("team-abc1234");
});
});

View file

@ -41,6 +41,17 @@ vi.mock("@/app/(dashboard)/hooks/login/useLogin", () => ({
})),
}));
vi.mock("@/hooks/useWorker", () => ({
useWorker: vi.fn(() => ({
isControlPlane: false,
workers: [],
selectedWorkerId: null,
selectedWorker: null,
selectWorker: vi.fn(),
disconnectFromWorker: vi.fn(),
})),
}));
import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
import { getCookie } from "@/utils/cookieUtils";
import { isJwtExpired } from "@/utils/jwtUtils";
@ -108,7 +119,7 @@ describe("LoginPage", () => {
);
await waitFor(() => {
expect(mockReplace).toHaveBeenCalledWith("http://localhost:4000/ui");
expect(mockReplace).toHaveBeenCalledWith("/ui");
});
});
@ -189,7 +200,7 @@ describe("LoginPage", () => {
);
await waitFor(() => {
expect(mockReplace).toHaveBeenCalledWith("http://localhost:4000/ui");
expect(mockReplace).toHaveBeenCalledWith("/ui");
});
expect(mockPush).not.toHaveBeenCalled();

View file

@ -3,14 +3,15 @@
import { useLogin } from "@/app/(dashboard)/hooks/login/useLogin";
import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
import LoadingScreen from "@/components/common_components/LoadingScreen";
import { getProxyBaseUrl } from "@/components/networking";
import { getCookie } from "@/utils/cookieUtils";
import { exchangeLoginCode, getProxyBaseUrl, switchToWorkerUrl } from "@/components/networking";
import { clearTokenCookies, getCookie } from "@/utils/cookieUtils";
import { isJwtExpired } from "@/utils/jwtUtils";
import { consumeReturnUrl, getReturnUrl, isValidReturnUrl } from "@/utils/returnUrlUtils";
import { InfoCircleOutlined } from "@ant-design/icons";
import { Alert, Button, Card, Form, Input, Popover, Space, Typography } from "antd";
import { InfoCircleOutlined, CloudServerOutlined } from "@ant-design/icons";
import { Alert, Button, Card, Form, Input, Popover, Select, Space, Typography } from "antd";
import { useRouter } from "next/navigation";
import { useEffect, useState } from "react";
import { useWorker } from "@/hooks/useWorker";
function LoginPageContent() {
const [username, setUsername] = useState("");
@ -19,6 +20,17 @@ function LoginPageContent() {
const { data: uiConfig, isLoading: isConfigLoading } = useUIConfig();
const loginMutation = useLogin();
const router = useRouter();
const { workers, selectWorker } = useWorker();
const [selectedWorkerId, setSelectedWorkerId] = useState<string | null>(null);
// Pre-select worker from URL param (e.g. /ui/login?worker=team-b)
useEffect(() => {
const params = new URLSearchParams(window.location.search);
const workerParam = params.get("worker");
if (workerParam) {
setSelectedWorkerId(workerParam);
}
}, []);
useEffect(() => {
if (isConfigLoading) {
@ -31,6 +43,44 @@ function LoginPageContent() {
return;
}
// Cross-origin SSO: worker redirected back with a single-use code.
// Exchange it for the JWT via the worker's /v3/login/exchange endpoint.
const params = new URLSearchParams(window.location.search);
const ssoCode = params.get("code");
if (ssoCode) {
const workerUrl = localStorage.getItem("litellm_worker_url");
exchangeLoginCode(ssoCode, workerUrl).then(() => {
params.delete("code");
const cleanSearch = params.toString();
window.history.replaceState(null, "", window.location.pathname + (cleanSearch ? `?${cleanSearch}` : ""));
router.replace("/ui/?login=success");
});
return;
}
// Backwards compat: handle direct token in URL (legacy flow)
const urlToken = params.get("token");
if (urlToken && !isJwtExpired(urlToken)) {
document.cookie = `token=${urlToken}; path=/; SameSite=Lax`;
params.delete("token");
const cleanSearch = params.toString();
window.history.replaceState(
null,
"",
window.location.pathname + (cleanSearch ? `?${cleanSearch}` : ""),
);
router.replace("/ui/?login=success");
return;
}
// If switching workers on a control plane, clear the old token and show login
const switchingWorker = params.has("worker");
if (switchingWorker && uiConfig?.is_control_plane) {
clearTokenCookies();
setIsLoading(false);
return;
}
const rawToken = getCookie("token");
if (rawToken && !isJwtExpired(rawToken)) {
// User already logged in - redirect to return URL or default
@ -38,7 +88,7 @@ function LoginPageContent() {
if (returnUrl) {
router.replace(returnUrl);
} else {
router.replace(`${getProxyBaseUrl()}/ui`);
router.replace("/ui");
}
return;
}
@ -58,16 +108,35 @@ function LoginPageContent() {
}, [isConfigLoading, router, uiConfig]);
const handleSubmit = () => {
// If a worker is selected, point proxyBaseUrl at it before login
const selectedWorker = workers.find((w) => w.worker_id === selectedWorkerId);
if (selectedWorker) {
switchToWorkerUrl(selectedWorker.url);
}
loginMutation.mutate(
{ username, password },
{ username, password, useV3: !!selectedWorker },
{
onSuccess: (data) => {
// Check if we have a return URL to use instead of the default redirect
const returnUrl = consumeReturnUrl();
if (returnUrl) {
router.push(returnUrl);
// Update the worker context with the selected worker
if (selectedWorker) {
selectWorker(selectedWorker.worker_id);
// Stay on the CP's UI — proxyBaseUrl already points at the worker
router.push("/ui/?login=success");
} else {
router.push(data.redirect_url);
// Normal (non-control-plane) login — follow the server's redirect
const returnUrl = consumeReturnUrl();
if (returnUrl) {
router.push(returnUrl);
} else {
router.push(data.redirect_url);
}
}
},
onError: () => {
// Reset proxyBaseUrl on login failure
if (selectedWorker) {
switchToWorkerUrl(null);
}
},
},
@ -154,6 +223,22 @@ function LoginPageContent() {
{error && <Alert message={error} type="error" showIcon />}
<Form onFinish={handleSubmit} layout="vertical" requiredMark={true}>
{uiConfig?.is_control_plane && workers.length > 0 && (
<Form.Item label="Worker" style={{ marginBottom: 16 }}>
<Select
value={selectedWorkerId || undefined}
onChange={(value) => setSelectedWorkerId(value)}
placeholder="Choose a worker to connect to"
size="large"
suffixIcon={<CloudServerOutlined />}
options={workers.map((w) => ({
label: w.name,
value: w.worker_id,
}))}
/>
</Form.Item>
)}
<Form.Item
label="Username"
name="username"
@ -209,10 +294,20 @@ function LoginPageContent() {
</Popover>
) : (
<Button
disabled={isLoginLoading}
onClick={() =>
router.push(`${getProxyBaseUrl()}/sso/key/generate`)
}
disabled={isLoginLoading || (!!selectedWorkerId && workers.length === 0)}
onClick={() => {
const selectedWorker = workers.find((w) => w.worker_id === selectedWorkerId);
if (selectedWorker) {
// Store worker selection so useWorker hook restores it after redirect
localStorage.setItem("litellm_selected_worker_id", selectedWorkerId!);
switchToWorkerUrl(selectedWorker.url);
}
// SSO on the worker (or this instance if no worker), always
// include return_to so the callback redirects back here
const ssoBase = selectedWorker?.url ?? getProxyBaseUrl();
const returnTo = encodeURIComponent(window.location.origin + "/ui/login");
router.push(`${ssoBase}/sso/key/generate?return_to=${returnTo}`);
}}
block
size="large"
>

View file

@ -0,0 +1,95 @@
import { render, screen } from "@testing-library/react";
import React from "react";
import { describe, expect, it, vi } from "vitest";
import { OnboardingForm } from "./OnboardingForm";
const mockUseOnboardingCredentials = vi.fn();
const mockClaimToken = vi.fn();
vi.mock("next/navigation", () => ({
useSearchParams: () => new URLSearchParams("invitation_id=inv-123"),
}));
vi.mock("jwt-decode", () => ({
jwtDecode: vi.fn(() => ({
user_email: "alice@example.com",
user_id: "user-1",
key: "access-tok",
})),
}));
vi.mock("@/app/(dashboard)/hooks/onboarding/useOnboarding", () => ({
useOnboardingCredentials: (...args: unknown[]) => mockUseOnboardingCredentials(...args),
useClaimOnboardingToken: () => ({ mutate: mockClaimToken, isPending: false }),
}));
vi.mock("@/components/networking", () => ({
getProxyBaseUrl: vi.fn(() => ""),
}));
vi.mock("./OnboardingLoadingView", () => ({
OnboardingLoadingView: () => <div data-testid="loading-view">Loading</div>,
}));
vi.mock("./OnboardingErrorView", () => ({
OnboardingErrorView: () => <div data-testid="error-view">Error</div>,
}));
vi.mock("./OnboardingFormBody", () => ({
OnboardingFormBody: ({ variant, userEmail }: { variant: string; userEmail: string }) => (
<div data-testid="form-body" data-variant={variant} data-email={userEmail}>
Form Body
</div>
),
}));
describe("OnboardingForm", () => {
it("should render loading view when credentials are loading", () => {
mockUseOnboardingCredentials.mockReturnValue({
data: undefined,
isLoading: true,
isError: false,
});
render(<OnboardingForm variant="signup" />);
expect(screen.getByTestId("loading-view")).toBeInTheDocument();
});
it("should render error view when credentials fail to load", () => {
mockUseOnboardingCredentials.mockReturnValue({
data: undefined,
isLoading: false,
isError: true,
});
render(<OnboardingForm variant="signup" />);
expect(screen.getByTestId("error-view")).toBeInTheDocument();
});
it("should render form body with decoded email when credentials are loaded", () => {
mockUseOnboardingCredentials.mockReturnValue({
data: { token: "fake-jwt-token" },
isLoading: false,
isError: false,
});
render(<OnboardingForm variant="signup" />);
expect(screen.getByTestId("form-body")).toBeInTheDocument();
expect(screen.getByTestId("form-body")).toHaveAttribute("data-email", "alice@example.com");
});
it("should pass variant prop to OnboardingFormBody", () => {
mockUseOnboardingCredentials.mockReturnValue({
data: { token: "fake-jwt-token" },
isLoading: false,
isError: false,
});
render(<OnboardingForm variant="reset_password" />);
expect(screen.getByTestId("form-body")).toHaveAttribute("data-variant", "reset_password");
});
});

View file

@ -1,6 +1,5 @@
"use client";
import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView";
import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider";
import OldModelDashboard from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView";
import PlaygroundPage from "@/app/(dashboard)/playground/page";
@ -9,7 +8,7 @@ import AgentsPanel from "@/components/agents";
import BudgetPanel from "@/components/budgets/budget_panel";
import CacheDashboard from "@/components/cache_dashboard";
import ClaudeCodePluginsPanel from "@/components/claude_code_plugins";
import { fetchTeams } from "@/components/common_components/fetch_teams";
import { teamListCall as v2TeamListCall } from "@/app/(dashboard)/hooks/teams/useTeams";
import LoadingScreen from "@/components/common_components/LoadingScreen";
import { CostTrackingSettings } from "@/components/CostTrackingSettings";
import GeneralSettings from "@/components/general_settings";
@ -48,7 +47,7 @@ import { buildLoginUrlWithReturn, consumeReturnUrl, normalizeUrlForCompare, stor
import { formatUserRole, isAdminRole } from "@/utils/roles";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { jwtDecode } from "jwt-decode";
import { useSearchParams } from "next/navigation";
import { useRouter, useSearchParams } from "next/navigation";
import { Suspense, useEffect, useMemo, useRef, useState } from "react";
import { ConfigProvider, theme } from "antd";
@ -75,6 +74,16 @@ interface ProxySettings {
LITELLM_UI_API_DOC_BASE_URL?: string | null;
}
/**
* Map of legacy query-param page keys new path-based route segments.
* When a user visits ?page=<key>, they are redirected to /ui/<value>.
* Add entries here as pages are migrated from the if/else chain to path-based routes.
*/
const LEGACY_REDIRECTS: Record<string, string> = {
api_ref: "api-reference",
"api-reference": "api-reference",
};
function CreateKeyPageContent() {
const [userRole, setUserRole] = useState("");
const [premiumUser, setPremiumUser] = useState(false);
@ -90,6 +99,7 @@ function CreateKeyPageContent() {
});
const [showSSOBanner, setShowSSOBanner] = useState<boolean>(true);
const router = useRouter();
const searchParams = useSearchParams()!;
const [modelData, setModelData] = useState<any>({ data: [] });
const [token, setToken] = useState<string | null>(null);
@ -243,6 +253,15 @@ function CreateKeyPageContent() {
}
}, [redirectToLogin]);
// Redirect legacy query-param pages to their new path-based routes
const isLegacyRedirect = page in LEGACY_REDIRECTS;
useEffect(() => {
if (!authLoading && isLegacyRedirect) {
const base = (proxyBaseUrl || "") + "/ui";
router.replace(`${base}/${LEGACY_REDIRECTS[page]}`);
}
}, [authLoading, isLegacyRedirect, page, router]);
// Check for a stored return URL after successful authentication
// This handles the case where user comes back from SSO and we need to redirect to the original URL
useEffect(() => {
@ -339,7 +358,9 @@ function CreateKeyPageContent() {
fetchUserModels(userID, userRole, accessToken, setUserModels);
}
if (accessToken && userID && userRole) {
fetchTeams(accessToken, userID, userRole, null, setTeams);
v2TeamListCall(accessToken, 1, 100, {
userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null,
}).then((response) => setTeams(response.teams ?? [])).catch(console.error);
}
if (accessToken) {
fetchOrganizations(accessToken, setOrganizations);
@ -427,7 +448,7 @@ function CreateKeyPageContent() {
setShowClaudeCodePrompt(true);
};
if (authLoading || redirectToLogin) {
if (authLoading || redirectToLogin || isLegacyRedirect) {
return <LoadingScreen />;
}
@ -536,8 +557,6 @@ function CreateKeyPageContent() {
<AdminPanel
proxySettings={proxySettings}
/>
) : page == "api_ref" ? (
<APIReferenceView proxySettings={proxySettings} />
) : page == "logging-and-alerts" ? (
<Settings userID={userID} userRole={userRole} accessToken={accessToken} premiumUser={premiumUser} />
) : page == "budgets" ? (

View file

@ -0,0 +1,146 @@
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { vi } from "vitest";
import { flexRender, getCoreRowModel, useReactTable } from "@tanstack/react-table";
import { getAgentHubTableColumns, AgentHubData } from "./AgentHubTableColumns";
const mockAgent: AgentHubData = {
agent_id: "agent-1",
protocolVersion: "1.0",
name: "Test Agent",
description: "A test agent for unit testing",
url: "https://agent.example.com",
version: "2.0",
capabilities: { streaming: true, caching: false },
defaultInputModes: ["text"],
defaultOutputModes: ["text", "image"],
skills: [
{ id: "s1", name: "Skill One", description: "First skill" },
{ id: "s2", name: "Skill Two", description: "Second skill" },
{ id: "s3", name: "Skill Three", description: "Third skill" },
],
is_public: true,
};
function TestTable({
data,
publicPage = false,
showModal = vi.fn(),
copyToClipboard = vi.fn(),
}: {
data: AgentHubData[];
publicPage?: boolean;
showModal?: ReturnType<typeof vi.fn>;
copyToClipboard?: ReturnType<typeof vi.fn>;
}) {
const columns = getAgentHubTableColumns(showModal, copyToClipboard, publicPage);
const table = useReactTable({ data, columns, getCoreRowModel: getCoreRowModel() });
return (
<table>
<thead>
{table.getHeaderGroups().map((hg) => (
<tr key={hg.id}>
{hg.headers.map((h) => (
<th key={h.id}>{flexRender(h.column.columnDef.header, h.getContext())}</th>
))}
</tr>
))}
</thead>
<tbody>
{table.getRowModel().rows.map((row) => (
<tr key={row.id}>
{row.getVisibleCells().map((cell) => (
<td key={cell.id}>{flexRender(cell.column.columnDef.cell, cell.getContext())}</td>
))}
</tr>
))}
</tbody>
</table>
);
}
describe("AgentHubTableColumns", () => {
it("should render", () => {
render(<TestTable data={[mockAgent]} />);
expect(screen.getByText("Test Agent")).toBeInTheDocument();
});
it("should display the agent description", () => {
render(<TestTable data={[mockAgent]} />);
// Description appears in both the description column and the mobile view within agent name column
expect(screen.getAllByText("A test agent for unit testing").length).toBeGreaterThanOrEqual(1);
});
it("should display the version with a 'v' prefix", () => {
render(<TestTable data={[mockAgent]} />);
expect(screen.getByText("v2.0")).toBeInTheDocument();
});
it("should display the protocol version", () => {
render(<TestTable data={[mockAgent]} />);
expect(screen.getByText("1.0")).toBeInTheDocument();
});
it("should show skill count with correct pluralization", () => {
render(<TestTable data={[mockAgent]} />);
expect(screen.getByText("3 skills")).toBeInTheDocument();
});
it("should show first two skills and '+1' for overflow", () => {
render(<TestTable data={[mockAgent]} />);
expect(screen.getByText("Skill One")).toBeInTheDocument();
expect(screen.getByText("Skill Two")).toBeInTheDocument();
expect(screen.getByText("+1")).toBeInTheDocument();
});
it("should show only true capabilities as badges", () => {
render(<TestTable data={[mockAgent]} />);
expect(screen.getByText("streaming")).toBeInTheDocument();
expect(screen.queryByText("caching")).not.toBeInTheDocument();
});
it("should display I/O modes", () => {
render(<TestTable data={[mockAgent]} />);
// "In:" and "Out:" are in <span> children; getByText with exact:false
// matches against the element's full textContent across child nodes
expect(screen.getByText((_, el) =>
el?.tagName === "P" && el.textContent === "In: text"
)).toBeInTheDocument();
expect(screen.getByText((_, el) =>
el?.tagName === "P" && el.textContent === "Out: text, image"
)).toBeInTheDocument();
});
it("should display 'Yes' badge for public agents", () => {
render(<TestTable data={[mockAgent]} />);
expect(screen.getByText("Yes")).toBeInTheDocument();
});
it("should display 'No' badge for non-public agents", () => {
const privateAgent = { ...mockAgent, is_public: false };
render(<TestTable data={[privateAgent]} />);
expect(screen.getByText("No")).toBeInTheDocument();
});
it("should display a Details button", () => {
render(<TestTable data={[mockAgent]} />);
expect(screen.getByRole("button", { name: /details|info/i })).toBeInTheDocument();
});
it("should show '-' when agent has no capabilities", () => {
const noCapAgent = { ...mockAgent, capabilities: {} };
render(<TestTable data={[noCapAgent]} />);
// The dash is rendered in the capabilities column
expect(screen.getByText("-")).toBeInTheDocument();
});
it("should show singular 'skill' for one skill", () => {
const oneSkillAgent = {
...mockAgent,
skills: [{ id: "s1", name: "Only Skill", description: "One" }],
};
render(<TestTable data={[oneSkillAgent]} />);
expect(screen.getByText("1 skill")).toBeInTheDocument();
});
});

View file

@ -194,7 +194,6 @@ export const getAgentHubTableColumns = (
return publicA - publicB;
},
cell: ({ row }) => {
console.log(`CHECKPOINT 1: ${JSON.stringify(row.original)}`);
const agent = row.original;
return agent.is_public === true ? (

View file

@ -1,5 +1,6 @@
import React from "react";
import { Modal, Form, message } from "antd";
import { Modal, Form } from "antd";
import MessageManager from "@/components/molecules/message_manager";
import {
AccessGroupBaseForm,
AccessGroupFormValues,
@ -37,7 +38,7 @@ export function AccessGroupCreateModal({
createMutation.mutate(params, {
onSuccess: () => {
message.success("Access group created successfully");
MessageManager.success("Access group created successfully");
form.resetFields();
onSuccess?.();
onCancel();

View file

@ -1,5 +1,6 @@
import React, { useEffect } from "react";
import { Modal, Form, message } from "antd";
import { Modal, Form } from "antd";
import MessageManager from "@/components/molecules/message_manager";
import {
AccessGroupBaseForm,
AccessGroupFormValues,
@ -55,7 +56,7 @@ export function AccessGroupEditModal({
{ accessGroupId: accessGroup.access_group_id, params },
{
onSuccess: () => {
message.success("Access group updated successfully");
MessageManager.success("Access group updated successfully");
onSuccess?.();
onCancel();
},

View file

@ -3,7 +3,6 @@ import {
Modal,
Typography,
Divider,
message,
Table,
Select,
InputNumber,
@ -14,6 +13,7 @@ import {
import { userBulkUpdateUserCall, teamBulkMemberAddCall, Member } from "./networking";
import { UserEditView } from "./user_edit_view";
import NotificationsManager from "./molecules/notifications_manager";
import MessageManager from "@/components/molecules/message_manager";
const { Text, Title } = Typography;
@ -188,7 +188,7 @@ const BulkEditUserModal: React.FC<BulkEditUserModalProps> = ({
}
if (failedTeams.length > 0) {
message.warning(`Failed to add users to ${failedTeams.length} team(s)`);
MessageManager.warning(`Failed to add users to ${failedTeams.length} team(s)`);
}
}

View file

@ -1,4 +1,5 @@
import { Form, Modal, Input, message } from "antd";
import { Form, Modal, Input } from "antd";
import MessageManager from "@/components/molecules/message_manager";
import { useEffect } from "react";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import { useCloudZeroCreate } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate";
@ -31,7 +32,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
},
{
onSuccess: () => {
message.success("CloudZero integration created successfully");
MessageManager.success("CloudZero integration created successfully");
form.resetFields();
onOk();
},
@ -39,7 +40,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
if (error?.errorFields) {
return;
}
message.error(error?.message || "Failed to create CloudZero integration");
MessageManager.error(error?.message || "Failed to create CloudZero integration");
},
},
);
@ -47,7 +48,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
if (error?.errorFields) {
return;
}
message.error(error?.message || "Failed to create CloudZero integration");
MessageManager.error(error?.message || "Failed to create CloudZero integration");
}
};

View file

@ -3,7 +3,8 @@ import { useCloudZeroExport } from "@/app/(dashboard)/hooks/cloudzero/useCloudZe
import { useCloudZeroDeleteSettings } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import DeleteResourceModal from "@/components/common_components/DeleteResourceModal";
import { Alert, Button, Card, Descriptions, Divider, message, Popconfirm, Tag } from "antd";
import { Alert, Button, Card, Descriptions, Divider, Popconfirm, Tag } from "antd";
import MessageManager from "@/components/molecules/message_manager";
import { CheckCircle, Edit, Play, Trash2, Upload } from "lucide-react";
import { useState } from "react";
import CloudZeroUpdateModal from "./CloudZeroUpdateModal";
@ -30,10 +31,10 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
{ limit: 10 },
{
onSuccess: (data) => {
message.success("Dry run completed successfully");
MessageManager.success("Dry run completed successfully");
},
onError: (error) => {
message.error(error?.message || "Failed to perform dry run");
MessageManager.error(error?.message || "Failed to perform dry run");
},
},
);
@ -48,10 +49,10 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
{ operation: "replace_hourly" },
{
onSuccess: () => {
message.success("Data successfully exported to CloudZero");
MessageManager.success("Data successfully exported to CloudZero");
},
onError: (error) => {
message.error(error?.message || "Failed to export data");
MessageManager.error(error?.message || "Failed to export data");
},
},
);
@ -79,12 +80,12 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
deleteMutation.mutate(undefined, {
onSuccess: () => {
message.success("CloudZero integration deleted successfully");
MessageManager.success("CloudZero integration deleted successfully");
setIsDeleteModalOpen(false);
onSettingsUpdated();
},
onError: (error) => {
message.error(error?.message || "Failed to delete CloudZero integration");
MessageManager.error(error?.message || "Failed to delete CloudZero integration");
},
});
};

View file

@ -1,6 +1,7 @@
import { useCloudZeroUpdateSettings } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import { Form, Input, message, Modal } from "antd";
import { Form, Input, Modal } from "antd";
import MessageManager from "@/components/molecules/message_manager";
import { useEffect } from "react";
import { CloudZeroSettings } from "./types";
@ -39,7 +40,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
},
{
onSuccess: () => {
message.success("CloudZero integration updated successfully");
MessageManager.success("CloudZero integration updated successfully");
form.resetFields();
onOk();
},
@ -47,7 +48,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
if (error?.errorFields) {
return;
}
message.error(error?.message || "Failed to update CloudZero integration");
MessageManager.error(error?.message || "Failed to update CloudZero integration");
},
},
);
@ -55,7 +56,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
if (error?.errorFields) {
return;
}
message.error(error?.message || "Failed to update CloudZero integration");
MessageManager.error(error?.message || "Failed to update CloudZero integration");
}
};

View file

@ -0,0 +1,73 @@
import { renderWithProviders, screen } from "../../../tests/test-utils";
import userEvent from "@testing-library/user-event";
import { vi } from "vitest";
import UsageExportHeader from "./UsageExportHeader";
import type { EntitySpendData } from "./types";
vi.mock("./EntityUsageExportModal", () => ({
default: ({ isOpen, onClose }: { isOpen: boolean; onClose: () => void }) =>
isOpen ? (
<div data-testid="export-modal">
<button onClick={onClose}>Close</button>
</div>
) : null,
}));
const defaultProps = {
dateValue: { from: new Date("2025-01-01"), to: new Date("2025-01-31") },
entityType: "team" as const,
spendData: {
results: [],
metadata: {
total_spend: 0,
total_api_requests: 0,
total_successful_requests: 0,
total_failed_requests: 0,
total_tokens: 0,
},
} satisfies EntitySpendData,
};
describe("UsageExportHeader", () => {
it("should render", () => {
renderWithProviders(<UsageExportHeader {...defaultProps} />);
expect(screen.getByRole("button", { name: /export data/i })).toBeInTheDocument();
});
it("should open the export modal when the export button is clicked", async () => {
const user = userEvent.setup();
renderWithProviders(<UsageExportHeader {...defaultProps} />);
await user.click(screen.getByRole("button", { name: /export data/i }));
expect(screen.getByTestId("export-modal")).toBeInTheDocument();
});
it("should close the export modal when onClose is called", async () => {
const user = userEvent.setup();
renderWithProviders(<UsageExportHeader {...defaultProps} />);
await user.click(screen.getByRole("button", { name: /export data/i }));
await user.click(screen.getByRole("button", { name: /close/i }));
expect(screen.queryByTestId("export-modal")).not.toBeInTheDocument();
});
it("should not show filter dropdown when showFilters is false", () => {
renderWithProviders(<UsageExportHeader {...defaultProps} showFilters={false} />);
expect(screen.queryByText(/filter/i)).not.toBeInTheDocument();
});
it("should show filter dropdown when showFilters is true and options provided", () => {
renderWithProviders(
<UsageExportHeader
{...defaultProps}
showFilters
filterLabel="Team"
filterPlaceholder="Select teams"
filterOptions={[
{ label: "Team A", value: "team-a" },
{ label: "Team B", value: "team-b" },
]}
onFiltersChange={vi.fn()}
/>,
);
expect(screen.getByText("Team")).toBeInTheDocument();
});
});

View file

@ -0,0 +1,98 @@
import { render, screen, act } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { vi } from "vitest";
import { GuardrailConfig } from "./GuardrailConfig";
describe("GuardrailConfig", () => {
const defaultProps = {
guardrailName: "Content Safety",
guardrailType: "Content Safety",
provider: "bedrock",
};
afterEach(() => {
vi.useRealTimers();
});
it("should render", () => {
render(<GuardrailConfig {...defaultProps} />);
expect(screen.getByText("Parameters")).toBeInTheDocument();
});
it("should display the guardrail name in the parameters description", () => {
render(<GuardrailConfig {...defaultProps} />);
expect(screen.getByText(/Configure Content Safety behavior/)).toBeInTheDocument();
});
// Note: Version history entries are hardcoded placeholders in the component.
// These assertions will need updating when wired to real API data.
it("should show version history when 'View history' is clicked", async () => {
const user = userEvent.setup();
render(<GuardrailConfig {...defaultProps} />);
await user.click(screen.getByRole("button", { name: /view history/i }));
expect(screen.getByText("Initial configuration")).toBeInTheDocument();
expect(screen.getByText("Added custom categories list")).toBeInTheDocument();
});
it("should toggle version history text between View/Hide", async () => {
const user = userEvent.setup();
render(<GuardrailConfig {...defaultProps} />);
const button = screen.getByRole("button", { name: /view history/i });
await user.click(button);
expect(screen.getByRole("button", { name: /hide history/i })).toBeInTheDocument();
});
it("should show custom code textarea when custom code override is toggled on", async () => {
const user = userEvent.setup();
render(<GuardrailConfig {...defaultProps} />);
// Walk up from "Custom Code Override" heading to find the enclosing section,
// then locate the switch within it
const heading = screen.getByText("Custom Code Override");
let container = heading.parentElement;
let customCodeSwitch: Element | null = null;
while (container && !customCodeSwitch) {
customCodeSwitch = container.querySelector('[role="switch"]');
container = container.parentElement;
}
if (!customCodeSwitch) {
throw new Error("Could not find the Custom Code Override switch via DOM traversal");
}
await user.click(customCodeSwitch);
expect(screen.getByPlaceholderText(/async def evaluate/)).toBeInTheDocument();
});
it("should hide custom code textarea when custom code override is off", () => {
render(<GuardrailConfig {...defaultProps} />);
// There's an input for categories, but no textarea
expect(screen.queryByPlaceholderText(/async def evaluate/)).not.toBeInTheDocument();
});
it("should show the re-run button in idle state", () => {
render(<GuardrailConfig {...defaultProps} />);
expect(screen.getByRole("button", { name: /re-run on failing logs/i })).toBeInTheDocument();
});
it("should show loading state when re-run is clicked", async () => {
vi.useFakeTimers({ shouldAdvanceTime: true });
const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime });
render(<GuardrailConfig {...defaultProps} />);
await user.click(screen.getByRole("button", { name: /re-run on failing logs/i }));
expect(screen.getByText(/Running on 10 samples/)).toBeInTheDocument();
});
it("should show success message after re-run completes", async () => {
vi.useFakeTimers({ shouldAdvanceTime: true });
const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime });
render(<GuardrailConfig {...defaultProps} />);
await user.click(screen.getByRole("button", { name: /re-run on failing logs/i }));
await act(async () => { vi.advanceTimersByTime(2500); });
expect(screen.getByText(/7\/10 would now pass/)).toBeInTheDocument();
});
it("should display the Revert and Save buttons", () => {
render(<GuardrailConfig {...defaultProps} />);
expect(screen.getByRole("button", { name: /revert/i })).toBeInTheDocument();
// The component's hardcoded default version is "v3", so Save shows "v4"
expect(screen.getByRole("button", { name: /save as v\d+/i })).toBeInTheDocument();
});
});

View file

@ -138,4 +138,18 @@ describe("DocsMenu", () => {
await user.click(button);
expect(button).toHaveAttribute("aria-expanded", "true");
});
it("should close menu when clicking outside", async () => {
const user = userEvent.setup();
renderWithProviders(
<div>
<DocsMenu items={items} />
<button>Outside</button>
</div>,
);
await user.click(screen.getByRole("button", { name: /docs/i }));
expect(screen.getByText("Custom pricing")).toBeInTheDocument();
await user.click(screen.getByRole("button", { name: /outside/i }));
expect(screen.queryByText("Custom pricing")).not.toBeInTheDocument();
});
});

View file

@ -0,0 +1,38 @@
"use client";
import React from "react";
import { Select } from "antd";
import { CloudServerOutlined } from "@ant-design/icons";
import { useWorker } from "@/hooks/useWorker";
interface WorkerDropdownProps {
onWorkerSwitch: (workerId: string) => void;
}
const WorkerDropdown: React.FC<WorkerDropdownProps> = ({ onWorkerSwitch }) => {
const { isControlPlane, selectedWorker, workers } = useWorker();
if (!isControlPlane || !selectedWorker) return null;
return (
<Select
showSearch
filterOption={(input, option) =>
(option?.label as string ?? "").toLowerCase().includes(input.toLowerCase())
}
value={selectedWorker.worker_id}
style={{ minWidth: 180 }}
suffixIcon={<CloudServerOutlined />}
options={workers.map((w) => ({
label: w.name,
value: w.worker_id,
disabled: w.worker_id === selectedWorker.worker_id,
}))}
onChange={(newWorkerId) => {
onWorkerSwitch(newWorkerId);
}}
/>
);
};
export default WorkerDropdown;

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